Copy some functions over from the signal docs

2026-10-22 20:46:42 -04:00
parent 4d733408eb
commit 2573c12295
2 changed files with 76 additions and 0 deletions
@@ -49,6 +49,80 @@ class MessengerClient:
        self.dh = GENERATE_DH()
        self.certificate = None

+    def RatchetInitAlice(state, SK, bob_dh_public_key):
+        state.DHs = 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 = {}
+
+    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 = {}
+
+    def RatchetSendKey(state):
+        state.CKs, mk = KDF_CK(state.CKs)
+        Ns = state.Ns
+        state.Ns += 1
+        return Ns, mk
+
+    def RatchetEncrypt(state, plaintext, AD):
+        Ns, mk = RatchetSendKey(state)
+        header = HEADER(state.DHs, state.PN, Ns)
+        return header, ENCRYPT(mk, plaintext, CONCAT(AD, header))
+
+    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)
+        state.CKr, mk = KDF_CK(state.CKr)
+        state.Nr += 1
+        return mk
+
+    def RatchetDecrypt(state, header, ciphertext, AD):
+        mk = RatchetReceiveKey(state, header)
+        return DECRYPT(mk, ciphertext, CONCAT(AD, header))
+
+    def TrySkippedMessageKeys(state, header):
+        if (header.dh, header.n) in state.MKSKIPPED:
+            mk = state.MKSKIPPED[header.dh, header.n]
+            del state.MKSKIPPED[header.dh, header.n]
+            return mk
+        else:
+            return None
+
+    def SkipMessageKeys(state, until):
+        if state.Nr + MAX_SKIP < until:
+            raise Error()
+        if state.CKr != None:
+            while state.Nr < until:
+                state.CKr, mk = KDF_CK(state.CKr)
+                state.MKSKIPPED[state.DHr, state.Nr] = mk
+                state.Nr += 1
+
+    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)))