New networkbot request: AUTH to authenticate bots

bot.py

@@ -19,9 +19,11 @@
 from datetime import datetime
 from datetime import timedelta
 from queue import Queue
+import os
+import random
+import re
 import threading
 import time
-import re
 
 import consumer
 import event
@@ -34,11 +36,23 @@ import response
 ID_letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
 
 class Bot:
-    def __init__(self, ip, realname, mp=list()):
+    def __init__(self, ip, realname, key=None, mp=list()):
         # Bot general informations
         self.version     = 3.3
         self.version_txt = "3.3-dev"
 
+        if os.path.exists("datas/dhk"):
+            try:
+                self.key = int.from_bytes(open("datas/dhk", 'rb').read(), byteorder='big', signed=False)
+            except:
+                self.key = None
+        else:
+            self.key = None
+        if self.key is None:
+            print ("No key found for association, please wait while I generate it for you...")
+            self.key = random.getrandbits(1024)
+            open("datas/dhk", 'wb').write(self.key.to_bytes((self.key.bit_length() // 8) + 1, byteorder='big', signed=False))
+
         # Save various informations
         self.ip = ip
         self.realname = realname

networkbot.py

@@ -32,6 +32,7 @@ class NetworkBot:
         self.context = context
         self.srv = srv
         self.dest = dest
+        self.key = None
 
         self.dcc = dcc # DCC connection to the other bot
         if self.dcc is not None:
@@ -220,6 +221,23 @@ class NetworkBot:
                 self.my_tag = random.randint(0,255)
             self.send_ack(tag)
 
+        elif cmd == b'AUTH': # Authenticate bot
+            if len(args) == 3: # Request pairing
+                g = int.from_bytes(args[0].decode('hex'), byteorder='big')
+                p = int.from_bytes(args[1].decode('hex'), byteorder='big')
+                A = int.from_bytes(args[2].decode('hex'), byteorder='big')
+
+                self.key = pow(A, b, p)
+
+                b = random.getrandbits(1024)
+                B = pow(g, b, p)
+
+                self.send_response_final(tag, pack("!x", pow(r, self.key, q)))
+            elif len(args) == 2: # Auth
+                r = int.from_bytes(args[0].decode('hex'), byteorder='big')
+                q = int.from_bytes(args[1].decode('hex'), byteorder='big')
+                self.send_response_final(tag, pack("!x", pow(r, self.key, q)))
+
         elif cmd == b'FETCH': # Get known commands
             for name in ["cmd_hook", "ask_hook", "msg_hook"]:
                 elts = self.context.create_cache(name)