Update app_secure_client.py and lib_secure_server.py WIP

app_secure_client.py

@@ -17,6 +17,159 @@ import hashlib
 from Crypto import Random
 from Crypto.PublicKey import RSA
 
+HDRLEN = 2
+
+class Message:
+    def __init__(self, sock, addr):
+        self._recv_buffer = bytearray(b"")
+        self.sock = sock
+        self.key = RSA.generate(2048, Random.new().read)
+        self.sessionkey = None
+        self.serverkey = None
+        self.addr = addr
+
+    def _json_encode(self,obj, encoding):
+        return json.dumps(obj, ensure_ascii=False).encode(encoding)
+
+    def _json_decode(self,json_bytes, encoding):
+        tiow = io.TextIOWrapper(
+            io.BytesIO(json_bytes), encoding=encoding, newline=""
+        )
+        obj = json.load(tiow)
+        tiow.close()
+        return obj
+
+    def _create_message(self, *, content_bytes, content_type, content_encoding):
+        jsonheader = {
+            "byteorder": sys.byteorder,
+            "content-type": content_type,
+            "content-encoding": content_encoding,
+            "content-length": len(content_bytes),
+        }
+        jsonheader_bytes = self._json_encode(jsonheader, "utf-8")
+        message_hdr = struct.pack(">H", len(jsonheader_bytes))
+        message = message_hdr + jsonheader_bytes + content_bytes
+        return message
+
+    def send_request(self, request):
+        content = request["content"]
+        content_type = request["type"]
+        content_encoding = request["encoding"]
+        if content_type == "text/json":
+            req = {
+                "content_bytes": self._json_encode(content, content_encoding),
+                "content_type": content_type,
+                "content_encoding": content_encoding,
+            }
+        else:
+            req = {
+                "content_bytes": content,
+                "content_type": content_type,
+                "content_encoding": content_encoding,
+            }
+        msg = self._create_message(**req)
+        print(f"sending request {repr(msg)} to {self.addr}")
+        self.sock.sendall(msg)
+
+    def read_to_buffer(self, length):
+        while len(self._recv_buffer) < length:
+            self._recv_buffer += self.sock.recv(4096)
+
+    def process_response(self):
+        # read and process protoheader (fixed length of 2 byte in network order)
+        self.read_to_buffer(HDRLEN)
+        _jsonheader_len = struct.unpack(">H", self._recv_buffer[:HDRLEN])[0]
+        self._recv_buffer = self._recv_buffer[HDRLEN:]
+
+        # read and process jsonheader
+        self.read_to_buffer(_jsonheader_len)
+        jsonheader = self._json_decode(self._recv_buffer[:_jsonheader_len], "utf-8")
+        self._recv_buffer = self._recv_buffer[_jsonheader_len:]
+        for reqhdr in ("byteorder", "content-length", "content-type",
+                       "content-encoding",):
+            if reqhdr not in jsonheader:
+                raise ValueError(f'Missing required header "{reqhdr}".')
+
+        content_len = jsonheader["content-length"]
+        self.read_to_buffer(content_len)
+        data = self._recv_buffer[:content_len]
+        self._recv_buffer = self._recv_buffer[content_len:]
+        if jsonheader["content-type"] == "text/json":
+            encoding = jsonheader["content-encoding"]
+            response = self._json_decode(data, encoding)
+            print("received response", repr(response), "from", self.addr)
+            content = response
+            result = content.get("result")
+            print(f"got result: {result}")
+        else:
+            # Binary or unknown content-type
+            response = data
+            print(
+                f'received {jsonheader["content-type"]} response from',
+                self.addr,
+            )
+            content = response
+            print(f"got response: {repr(content)}")
+
+    def _send_handshake(self, data):
+            response = {
+                "content_bytes": data,
+                "content_type": "binary/custom-server-binary-type",
+                "content_encoding": "binary",
+            }
+            message = self._create_message(**response)
+            self.sock.sendall(message)
+
+    def close(self):
+        self.sock.shutdown(socket.SHUT_RDWR)
+        self.sock.close()
+        self.sock = None
+
+    def _recv_handshake(self):
+        # read and process protoheader (fixed length of 2 byte in network order)
+        self.read_to_buffer(HDRLEN)
+        _jsonheader_len = struct.unpack(">H", self._recv_buffer[:HDRLEN])[0]
+        self._recv_buffer = self._recv_buffer[HDRLEN:]
+
+        # read and process jsonheader
+        self.read_to_buffer(_jsonheader_len)
+        jsonheader = self._json_decode(self._recv_buffer[:_jsonheader_len], "utf-8")
+        self._recv_buffer = self._recv_buffer[_jsonheader_len:]
+        for reqhdr in ("byteorder", "content-length", "content-type",
+                       "content-encoding",):
+            if reqhdr not in jsonheader:
+                raise ValueError(f'Missing required header "{reqhdr}".')
+
+        content_len = jsonheader["content-length"]
+        self.read_to_buffer(content_len)
+        data = self._recv_buffer[:content_len]
+        self._recv_buffer = self._recv_buffer[content_len:]
+        sesskeyhash = data[-128:]
+        data = data[:-128]
+        encryptedsess = data[-256:]
+        halfsesskey = self.key.decrypt(bytes(encryptedsess))
+        data = data[:-256]
+        pubkeyhash = data[-128:]
+        pubkey = data[:-128]
+        if hashlib.sha3_512(pubkey).hexdigest().encode() != pubkeyhash:
+            self.close()
+            raise Exception('Server public key does not match hash')
+        if hashlib.sha3_512(halfsesskey).hexdigest().encode() != sesskeyhash:
+            self.close()
+            raise Exception('Server public key does not match hash')
+        self.sessionkey = halfsesskey
+        self.serverkey = RSA.importKey(pubkey)
+        print(f"Received server secrets correctly.")
+
+    def send_passphrase(self, passphrase):
+        passphrase = passphrase.encode()
+        passhash = hashlib.sha3_512(passphrase).hexdigest().encode()
+        sesshash = hashlib.sha3_512(self.sessionkey).hexdigest().encode()
+        data = self.serverkey.encrypt(passhash + sesshash, None)[0]
+        print(f'Sending passphrase + hash + sessionkey hash...')
+        self._send_handshake(data)
+
+
 def create_request(action, value):
     if action == "search":
         return dict(
@@ -32,134 +185,7 @@ def create_request(action, value):
         )
 
 
-def _json_encode(obj, encoding):
-    return json.dumps(obj, ensure_ascii=False).encode(encoding)
-
-
-def _json_decode(json_bytes, encoding):
-    tiow = io.TextIOWrapper(
-        io.BytesIO(json_bytes), encoding=encoding, newline=""
-    )
-    obj = json.load(tiow)
-    tiow.close()
-    return obj
-
-
-def _create_message(*, content_bytes, content_type, content_encoding):
-    jsonheader = {
-        "byteorder": sys.byteorder,
-        "content-type": content_type,
-        "content-encoding": content_encoding,
-        "content-length": len(content_bytes),
-    }
-    jsonheader_bytes = _json_encode(jsonheader, "utf-8")
-    message_hdr = struct.pack(">H", len(jsonheader_bytes))
-    message = message_hdr + jsonheader_bytes + content_bytes
-    return message
-
-
-def create_message(request):
-    content = request["content"]
-    content_type = request["type"]
-    content_encoding = request["encoding"]
-    if content_type == "text/json":
-        req = {
-            "content_bytes": _json_encode(content, content_encoding),
-            "content_type": content_type,
-            "content_encoding": content_encoding,
-        }
-    else:
-        req = {
-            "content_bytes": content,
-            "content_type": content_type,
-            "content_encoding": content_encoding,
-        }
-    message = _create_message(**req)
-    return message
-
-
-def read_to_buffer(sock, length, buf : bytearray):
-    while len(buf) < length:
-        data = sock.recv(4096)
-        buf += data
-
-
-def process_response(sock, addr, jsonheader, _recv_buffer):
-    content_len = jsonheader["content-length"]
-    read_to_buffer(sock, content_len, _recv_buffer)
-    data = _recv_buffer[:content_len]
-    _recv_buffer = _recv_buffer[content_len:]
-    if jsonheader["content-type"] == "text/json":
-        encoding = jsonheader["content-encoding"]
-        response = _json_decode(data, encoding)
-        print("received response", repr(response), "from", addr)
-        content = response
-        result = content.get("result")
-        print(f"got result: {result}")
-    else:
-        # Binary or unknown content-type
-        response = data
-        print(
-            f'received {jsonheader["content-type"]} response from',
-            addr,
-        )
-        content = response
-        print(f"got response: {repr(content)}")
-
-
-def send_client_key(sock, data):
-        response = {
-            "content_bytes": data,
-            "content_type": "binary/custom-server-binary-type",
-            "content_encoding": "binary",
-        }
-        message = _create_message(**response)
-        print(f'Sending client public key and hash...')
-        sock.sendall(message)
-
-def recv_server_secrets(sock, _recv_buffer, key):
-    HDRLEN = 2
-    # read and process protoheader (fixed length of 2 byte in network order)
-    read_to_buffer(sock, HDRLEN, _recv_buffer)
-    _jsonheader_len = struct.unpack(">H", _recv_buffer[:HDRLEN])[0]
-    _recv_buffer = _recv_buffer[HDRLEN:]
-
-    # read and process jsonheader
-    read_to_buffer(sock, _jsonheader_len, _recv_buffer)
-    jsonheader = _json_decode(_recv_buffer[:_jsonheader_len], "utf-8")
-    _recv_buffer = _recv_buffer[_jsonheader_len:]
-    for reqhdr in ("byteorder", "content-length", "content-type",
-                   "content-encoding",):
-        if reqhdr not in jsonheader:
-            raise ValueError(f'Missing required header "{reqhdr}".')
-
-    content_len = jsonheader["content-length"]
-    read_to_buffer(sock, content_len, _recv_buffer)
-    data = _recv_buffer[:content_len]
-    _recv_buffer = _recv_buffer[content_len:]
-    sesskeyhash = data[-128:]
-    data = data[:-128]
-    encryptedsess = data[-256:]
-    halfsesskey = key.decrypt(bytes(encryptedsess))
-    data = data[:-256]
-    pubkeyhash = data[-128:]
-    pubkey = data[:-128]
-    if hashlib.sha3_512(pubkey).hexdigest().encode() != pubkeyhash:
-        sock.shutdown(socket.SHUT_RDWR)
-        sock.close()
-        raise Exception('Server public key does not match hash')
-    if hashlib.sha3_512(halfsesskey).hexdigest().encode() != sesskeyhash:
-        sock.shutdown(socket.SHUT_RDWR)
-        sock.close()
-        raise Exception('Server public key does not match hash')
-    print(f"Received server secrets correctly.")
-    return RSA.importKey(pubkey), halfsesskey
-
-
 def main():
-    HDRLEN = 2
-    _recv_buffer = bytearray(b"")
-
     if len(sys.argv) != 5:
         print("usage:", sys.argv[0], "<host> <port> <action> <value>")
         sys.exit(1)
@@ -170,39 +196,25 @@ def main():
     sock = socket.create_connection(addr)
     sock.settimeout(None)
 
+    message = Message(sock, addr)
     # security handshake
-    key = RSA.generate(2048, Random.new().read)
+    publickey = message.key.publickey().exportKey('PEM')
-    publickey = key.publickey().exportKey('PEM')
     keyhash = hashlib.sha3_512(publickey).hexdigest().encode()
     data = publickey + keyhash
-    send_client_key(sock, data)
+    print(f'Sending client public key and hash...')
+    message._send_handshake(data)
     # receiving server secrets - public key and session key
-    serverkey, sessionkey = recv_server_secrets(sock, _recv_buffer, key)
+    message._recv_handshake()
-    # acknowledge session key and passphrase
+    message.send_passphrase('adminadmin')
-    # TODO send encrypted key and passphrase back
+    print('Handshake completed. Sever would close the sock if password wrong')
 
     action, value = sys.argv[3], sys.argv[4]
     request = create_request(action, value)
-    message = create_message(request)
-    print(f"sending request {repr(message)} to {addr}")
-    sock.sendall(message)
 
-    # read and process protoheader (fixed length of 2 byte in network order)
+    message.send_request(request)
-    read_to_buffer(sock, HDRLEN, _recv_buffer)
-    _jsonheader_len = struct.unpack(">H", _recv_buffer[:HDRLEN])[0]
-    _recv_buffer = _recv_buffer[HDRLEN:]
-
-    # read and process jsonheader
-    read_to_buffer(sock, _jsonheader_len, _recv_buffer)
-    jsonheader = _json_decode(_recv_buffer[:_jsonheader_len], "utf-8")
-    _recv_buffer = _recv_buffer[_jsonheader_len:]
-    for reqhdr in ("byteorder", "content-length", "content-type",
-                   "content-encoding",):
-        if reqhdr not in jsonheader:
-            raise ValueError(f'Missing required header "{reqhdr}".')
 
     # read and process response
-    process_response(sock, addr, jsonheader, _recv_buffer)
+    message.process_response()
 
     sock.shutdown(socket.SHUT_RDWR)
     sock.close()

lib_secure_server.py

@@ -37,7 +37,7 @@ class Message:
         # added for security
         self.key = RSA.generate(2048, Random.new().read)
         self.sessionkey = os.urandom(8)
-        self.passhash = passhash # pass phrase hash
+        self.password = passhash # pass phrase hash
         self.clientkey = None # client public key
         self.handshake = False
         print(f"initiating message associated with {addr}")
@@ -81,8 +81,10 @@ class Message:
                 pass
             else:
                 self._send_buffer = self._send_buffer[sent:]
+                if not self.handshake:
+                    self._set_selector_events_mask("r")
                 # Close when the buffer is drained. The response has been sent.
-                if sent and not self._send_buffer:
+                elif sent and not self._send_buffer:
                     self.close()
 
     def _json_encode(self, obj, encoding):
@@ -163,6 +165,10 @@ class Message:
             if self.request is None:
                 self.process_request()
 
+        if not self.handshake:
+            self._jsonheader_len = None
+            self.jsonheader = None
+
     def write(self):
         if self.request:
             if not self.response_created:
@@ -233,7 +239,7 @@ class Message:
             return
         data = self._recv_buffer[:content_len]
         self._recv_buffer = self._recv_buffer[content_len:]
-        if not self.handshake: # data is client public key plus hash
+        if self.clientkey is None: # data is client public key plus hash
             key = data[:-128]
             keyhash = data[-128:]
             if hashlib.sha3_512(key).hexdigest().encode() == keyhash:
@@ -242,6 +248,20 @@ class Message:
                 self.sock.close()
                 raise Exception('Client plublic key does not match key hash')
             self.send_my_secrets()
+            self._set_selector_events_mask("w")
+        elif not self.handshake: # data is passhass + sesshash
+            passphrase = self.key.decrypt(bytes(data))
+            sesshash = passphrase[-128:]
+            passhash = passphrase[:-128]
+            if passhash != self.password:
+                self.sock.close()
+                raise Exception('Password mismatch')
+            if hashlib.sha3_512(self.sessionkey).hexdigest().encode() != sesshash:
+                self.sock.close()
+                raise Exception('Session key hash mismatch')
+            self.handshake = True
+            print('Handshake Done.')
+            self._set_selector_events_mask("r")
         else:
             if self.jsonheader["content-type"] == "text/json":
                 encoding = self.jsonheader["content-encoding"]
@@ -254,8 +274,8 @@ class Message:
                     f'received {self.jsonheader["content-type"]} request from',
                     self.addr,
                 )
-        # Set selector to listen for write events, we're done reading.
+            # Set selector to listen for write events, we're done reading.
-        self._set_selector_events_mask("w")
+            self._set_selector_events_mask("w")
 
     def create_response(self):
         if self.jsonheader["content-type"] == "text/json":