Update app_secure_client.py and lib_secure_server.py - fully functional

app_secure_client.py

@@ -16,8 +16,10 @@ import struct
 import hashlib
 from Crypto import Random
 from Crypto.PublicKey import RSA
+from Crypto.Cipher import AES
 
 HDRLEN = 2
+PASSWORD = 'adminadmin'
 
 class Message:
     def __init__(self, sock, addr):
@@ -27,6 +29,7 @@ class Message:
         self.sessionkey = None
         self.serverkey = None
         self.addr = addr
+        self.handshake = False
 
     def _json_encode(self,obj, encoding):
         return json.dumps(obj, ensure_ascii=False).encode(encoding)
@@ -40,15 +43,24 @@ class Message:
         return obj
 
     def _create_message(self, *, content_bytes, content_type, content_encoding):
+        if self.handshake:
+            # encrypt content_bytes
+            key = self.sessionkey + self.sessionkey[::-1]
+            iv = Random.new().read(AES.block_size)
+            cipher = AES.new(key, AES.MODE_CFB, iv)
+            data = iv + cipher.encrypt(content_bytes)
+        else:
+            data = content_bytes
+
         jsonheader = {
             "byteorder": sys.byteorder,
             "content-type": content_type,
             "content-encoding": content_encoding,
-            "content-length": len(content_bytes),
+            "content-length": len(data),
         }
         jsonheader_bytes = self._json_encode(jsonheader, "utf-8")
         message_hdr = struct.pack(">H", len(jsonheader_bytes))
-        message = message_hdr + jsonheader_bytes + content_bytes
+        message = message_hdr + jsonheader_bytes + data
         return message
 
     def send_request(self, request):
@@ -93,6 +105,12 @@ class Message:
         content_len = jsonheader["content-length"]
         self.read_to_buffer(content_len)
         data = self._recv_buffer[:content_len]
+        # decrypt data
+        iv = bytes(data[:16])
+        encrypted = bytes(data[16:])
+        key = self.sessionkey + self.sessionkey[::-1]
+        cipher = AES.new(key, AES.MODE_CFB, iv)
+        data = cipher.decrypt(encrypted)
         self._recv_buffer = self._recv_buffer[content_len:]
         if jsonheader["content-type"] == "text/json":
             encoding = jsonheader["content-encoding"]
@@ -168,6 +186,7 @@ class Message:
         data = self.serverkey.encrypt(passhash + sesshash, None)[0]
         print(f'Sending passphrase + hash + sessionkey hash...')
         self._send_handshake(data)
+        self.handshake = True
 
 
 def create_request(action, value):
@@ -205,7 +224,7 @@ def main():
     message._send_handshake(data)
     # receiving server secrets - public key and session key
     message._recv_handshake()
-    message.send_passphrase('adminadmin')
+    message.send_passphrase(PASSWORD)
     print('Handshake completed. Sever would close the sock if password wrong')
 
     action, value = sys.argv[3], sys.argv[4]

lib_secure_server.py

@@ -14,6 +14,7 @@ import struct
 import hashlib
 from Crypto import Random
 from Crypto.PublicKey import RSA
+from Crypto.Cipher import AES
 
 
 request_search = {
@@ -113,6 +114,13 @@ class Message:
         message = message_hdr + jsonheader_bytes + content_bytes
         return message
 
+    def encrypt_data(self, data):
+        key = self.sessionkey + self.sessionkey[::-1]
+        iv = Random.new().read(AES.block_size)
+        cipher = AES.new(key, AES.MODE_CFB, iv)
+        msg = iv + cipher.encrypt(data)
+        return msg
+
     def _create_response_json_content(self):
         action = self.request.get("action")
         if action == "search":
@@ -122,17 +130,20 @@ class Message:
         else:
             content = {"result": f'Error: invalid action "{action}".'}
         content_encoding = "utf-8"
+        # encrypt content_bytes
+        data = self.encrypt_data(self._json_encode(content, content_encoding))
+
         response = {
-            "content_bytes": self._json_encode(content, content_encoding),
+            "content_bytes": data,
             "content_type": "text/json",
             "content_encoding": content_encoding,
         }
         return response
 
     def _create_response_binary_content(self):
+        msg = self.encrypt_data(b"First 10 bytes of request: " + self.request[:10])
         response = {
-            "content_bytes": b"First 10 bytes of request: "
-            + self.request[:10],
+            "content_bytes": msg,
             "content_type": "binary/custom-server-binary-type",
             "content_encoding": "binary",
         }
@@ -270,7 +281,12 @@ class Message:
             self._reset_header()
             self._set_selector_events_mask("r")
         else:
-            print(f"Receiving request (unencrypted yet)")
+            print(f"Receiving request (encrypted)")
+            iv = data[:16]
+            encrypted = data[16:]
+            key = self.sessionkey + self.sessionkey[::-1]
+            cipher = AES.new(key, AES.MODE_CFB, iv)
+            data = cipher.decrypt(encrypted)
             if self.jsonheader["content-type"] == "text/json":
                 encoding = self.jsonheader["content-encoding"]
                 self.request = self._json_decode(data, encoding)