Throw in safety checks that return none

messenger.py

@@ -1,7 +1,8 @@
 import os
 import pickle
-import string
 import hmac
+import copy
+import json
 from dataclasses import dataclass
 from types import SimpleNamespace
 
@@ -45,6 +46,8 @@ def ENCRYPT(mk, plaintext, associated_data):
     return nonce + AESGCM(mk).encrypt(nonce, plaintext, associated_data)
 
 def DECRYPT(mk, ciphertext, associated_data):
+    if mk is None or not isinstance(ciphertext, bytes) or len(ciphertext) < 12:
+        return None
     return AESGCM(mk).decrypt(ciphertext[:12], ciphertext[12:], associated_data)
 
 def HEADER(dh_pair, pn, n):
@@ -93,8 +96,8 @@ class MessengerClient:
         self.certificate = None
 
     @staticmethod
-    def RatchetInitAlice(state, SK, bob_dh_public_key):
+    def RatchetInitAlice(state, SK, bob_dh_public_key, alice_dh_key_pair=None):
-        state.DHs = GENERATE_DH()
+        state.DHs = alice_dh_key_pair or GENERATE_DH()
         state.DHr = bob_dh_public_key
         state.RK, state.CKs = KDF_RK(SK, DH(state.DHs, state.DHr))
         state.CKr = None
@@ -130,17 +133,14 @@ class MessengerClient:
 
     @staticmethod
     def RatchetReceiveKey(state, header):
-        # mk = TrySkippedMessageKeys(state, header)
+        if not isinstance(header, MessageHeader) or header.dh is None:
-        # if mk != None:
+            return None
-        #     return mk
+        if header.pn < 0 or header.n < 0:
-        # if header.dh != state.DHr:
+            return None
-            # SkipMessageKeys(state, header.pn)
-            # DHRatchet(state, header)
-        # SkipMessageKeys(state, header.n)
         if state.DHr is None or header.dh.public_bytes(serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint) != state.DHr.public_bytes(serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint):
             MessengerClient.DHRatchet(state, header)
         if header.n != state.Nr:
-            return None # out of order
+            return None
         state.CKr, mk = KDF_CK(state.CKr)
         state.Nr += 1
         return mk
@@ -172,31 +172,47 @@ class MessengerClient:
         self.certs[name] = public_key
 
     def sendMessage(self, name, message):
+        if name not in self.certs:
+            return None
         state = self.conns.get(name)
         if state is None:
             state = SimpleNamespace()
-            self.conns[name] = state
             initial_secret = DH(self.dh, self.certs[name])
-            self.RatchetInitAlice(state, initial_secret, self.certs[name])
+            self.RatchetInitAlice(
+                state, initial_secret, self.certs[name], self.dh
+            )
+            self.conns[name] = state
         plaintext = message.encode("utf-8")
         return self.RatchetEncrypt(state, plaintext, b"")
 
     def receiveMessage(self, name, header, ciphertext):
+        if name not in self.certs:
+            return None
         state = self.conns.get(name)
         if state is None:
             state = SimpleNamespace()
             initial_secret = DH(self.dh, self.certs[name])
-            self.conns[name] = state
             self.RatchetInitBob(state, initial_secret, self.dh)
+        temp_state = copy.copy(state) # fixes the issue with bad ciphertexts corrupting the state
         try:
-            plaintext = self.RatchetDecrypt(state, header, ciphertext, b"")
+            plaintext = self.RatchetDecrypt(temp_state, header, ciphertext, b"")
-        except InvalidTag:
+            message = plaintext.decode("utf-8")
+        except (
+            InvalidTag,
+            ValueError,
+            TypeError,
+            AttributeError,
+            IndexError,
+            UnicodeDecodeError,
+            pickle.UnpicklingError,
+            EOFError,
+            KeyError,
+        ):
             return None
-        return plaintext.decode("utf-8")
+        self.conns[name] = temp_state
+        return message
 
 
-    # this is technically hybrid encryption instead of just el gamal but I blame the instructions instead of myself
-    # maybe i'll change this if I want to use a late day
     def report(self, name, message):
         report = (name, message)
         plaintext = pickle.dumps(report)