add byteorder

main.py

@@ -54,6 +54,8 @@ print("Testing Reporting")
 content = "inappropriate message contents"
 reportPT, reportCT = alice.report("Bob", content)
 decryptedReport = server.decryptReport(reportCT)
+# print(reportPT)
+# print(decryptedReport)
 if decryptedReport != reportPT:
     error("report did not decrypt properly")
     print(reportPT)

messenger.py

@@ -1,6 +1,12 @@
 import os
 import pickle
-import string
+import hmac
+import copy
+import json
+from dataclasses import dataclass
+from types import SimpleNamespace
+
+from cryptography.exceptions import InvalidTag
 from cryptography.hazmat.primitives import hashes, serialization
 from cryptography.hazmat.primitives.asymmetric import ec
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
@@ -8,17 +14,72 @@ from cryptography.hazmat.primitives.kdf.hkdf import HKDF
 
 CURVE = ec.SECP256R1()
 
+@dataclass
+class MessageHeader:
+    dh: object
+    pn: int
+    n: int
+
 def GENERATE_DH():
     return ec.generate_private_key(CURVE)
 
+def DH(dh_pair, dh_pub):
+    return dh_pair.exchange(ec.ECDH(), dh_pub)
+
+def KDF_RK(rk, dh_out):
+    output = HKDF(
+        algorithm=hashes.SHA256(),
+        length=64,
+        salt=rk,
+        info=b"",
+    ).derive(dh_out)
+    return output[:32], output[32:]
+
+def KDF_CK(ck):
+    return (
+        hmac.new(ck, b"\x01", "sha256").digest(),
+        hmac.new(ck, b"\x02", "sha256").digest(),
+    )
+
+def ENCRYPT(mk, plaintext, associated_data):
+    nonce = os.urandom(12)
+    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):
+    public_key = dh_pair.public_key()
+    return MessageHeader(public_key, int(pn), int(n))
+
+def CONCAT(ad, header):
+    header_tup = pickle.dumps(
+        (header.dh.public_bytes(serialization.Encoding.X962,serialization.PublicFormat.UncompressedPoint), header.pn, header.n),
+    )
+    return len(ad).to_bytes(8, "big") + ad + header_tup
+
 class MessengerServer:
     def __init__(self, server_signing_key, server_decryption_key):
         self.server_signing_key = server_signing_key
         self.server_decryption_key = server_decryption_key
 
     def decryptReport(self, ct):
-        raise Exception("not implemented!")
-        return
+        data = pickle.loads(ct)
+        key = ec.EllipticCurvePublicKey.from_encoded_point(
+            CURVE, data["bytes"]
+        )
+        shared = DH(self.server_decryption_key, key)
+        derived_key = HKDF(
+            algorithm=hashes.SHA256(),
+            length=32,
+            salt=None,
+            info=b""
+        ).derive(shared)
+        aesgcm = AESGCM(derived_key)
+        plaintext = aesgcm.decrypt(data["nonce"], data["ciphertext"], data["bytes"])
+        return plaintext.decode("latin1")
 
     def signCert(self, cert):
         return self.server_signing_key.sign(cert, ec.ECDSA(hashes.SHA256()))
@@ -34,6 +95,71 @@ class MessengerClient:
         self.dh = GENERATE_DH()
         self.certificate = None
 
+    @staticmethod
+    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
+        state.Ns = 0
+        state.Nr = 0
+        state.PN = 0
+        state.MKSKIPPED = {}
+
+    @staticmethod
+    def RatchetInitBob(state, SK, bob_dh_key_pair):
+        state.DHs = bob_dh_key_pair
+        state.DHr = None
+        state.RK = SK
+        state.CKs = None
+        state.CKr = None
+        state.Ns = 0
+        state.Nr = 0
+        state.PN = 0
+        state.MKSKIPPED = {}
+
+    @staticmethod
+    def RatchetSendKey(state):
+        state.CKs, mk = KDF_CK(state.CKs)
+        Ns = state.Ns
+        state.Ns += 1
+        return Ns, mk
+
+    @staticmethod
+    def RatchetEncrypt(state, plaintext, AD):
+        Ns, mk = MessengerClient.RatchetSendKey(state)
+        header = HEADER(state.DHs, state.PN, Ns)
+        return header, ENCRYPT(mk, plaintext, CONCAT(AD, header))
+
+    @staticmethod
+    def RatchetReceiveKey(state, header):
+        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
+        state.CKr, mk = KDF_CK(state.CKr)
+        state.Nr += 1
+        return mk
+
+    @staticmethod
+    def RatchetDecrypt(state, header, ciphertext, AD):
+        mk = MessengerClient.RatchetReceiveKey(state, header)
+        return DECRYPT(mk, ciphertext, CONCAT(AD, header))
+
+    @staticmethod
+    def DHRatchet(state, header):
+        state.PN = state.Ns
+        state.Ns = 0
+        state.Nr = 0
+        state.DHr = header.dh
+        state.RK, state.CKr = KDF_RK(state.RK, DH(state.DHs, state.DHr))
+        state.DHs = GENERATE_DH()
+        state.RK, state.CKs = KDF_RK(state.RK, DH(state.DHs, state.DHr))
+
     def generateCertificate(self):
         self.certificate = pickle.dumps((self.name, self.dh.public_key().public_bytes(serialization.Encoding.X962,
         serialization.PublicFormat.UncompressedPoint)))
@@ -46,13 +172,69 @@ class MessengerClient:
         self.certs[name] = public_key
 
     def sendMessage(self, name, message):
-        raise Exception("not implemented!")
-        return
+        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.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):
-        raise Exception("not implemented!")
-        return
+        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.RatchetInitBob(state, initial_secret, self.dh)
+        temp_state = copy.copy(state) # fixes the issue with bad ciphertexts corrupting the state
+        try:
+            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
+        self.conns[name] = temp_state
+        return message
+
 
     def report(self, name, message):
-        raise Exception("not implemented!")
-        return
+        report = (name, message)
+        plaintext = pickle.dumps(report)
+        temp_key = GENERATE_DH()
+        temp_bytes = temp_key.public_key().public_bytes(
+            serialization.Encoding.X962,
+            serialization.PublicFormat.UncompressedPoint,
+        )
+        shared_secret = DH(temp_key, self.server_encryption_pk)
+        hkdf = HKDF(
+            algorithm=hashes.SHA256(),
+            length=32,
+            salt=None,
+            info=b""
+        )
+        key = hkdf.derive(shared_secret)
+        nonce = os.urandom(12)
+        aesgcm = AESGCM(key)
+        ciphertext = aesgcm.encrypt(nonce, plaintext, temp_bytes) # not sure if passing in temp_bytes is correct, but it seems like it should be
+        data = {
+            "bytes": temp_bytes,
+            "nonce": nonce,
+            "ciphertext": ciphertext
+        }
+        return plaintext.decode("latin1"), pickle.dumps(data)