diff --git a/messenger.py b/messenger.py index 5090e8b..88c2333 100644 --- a/messenger.py +++ b/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): - state.DHs = GENERATE_DH() + def RatchetInitAlice(state, SK, bob_dh_public_key, alice_dh_key_pair=None): + 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 mk != None: - # return mk - # if header.dh != state.DHr: - # SkipMessageKeys(state, header.pn) - # DHRatchet(state, header) - # SkipMessageKeys(state, header.n) + if not isinstance(header, MessageHeader) or header.dh is None: + return None + if header.pn < 0 or header.n < 0: + return None 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"") - except InvalidTag: + plaintext = self.RatchetDecrypt(temp_state, header, ciphertext, b"") + 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)