initial commit

ios_client/EncryptedChat/Crypto/SenderKeyState.swift

@@ -0,0 +1,175 @@
+import Foundation
+import CryptoKit
+
+/// Sender key chain for group messaging.
+/// Each sender in a group has their own chain. Others receive the initial key via pairwise ratchet.
+/// Matches Python: SenderKeyState class in crypto_utils.py
+class SenderKeyState {
+
+    let senderKey: Data
+    let chainId: Data
+    private(set) var chainKey: Data
+    private(set) var n: Int
+    private var knownKeys: [Int: Data]
+
+    /// Initialize with optional sender key (generates random 32B if nil).
+    /// Matches Python: SenderKeyState.__init__(sender_key=None)
+    init(senderKey: Data? = nil) {
+        let key = senderKey ?? Data.randomBytes(32)
+        self.senderKey = key
+        self.chainId = Data(SHA256.hash(data: key))
+        self.chainKey = CryptoUtils.hkdfDerive(
+            inputKey: key,
+            salt: Data(repeating: 0x00, count: 32),
+            info: Data(Constants.senderKeyChainInfo.utf8),
+            length: 32
+        )
+        self.n = 0
+        self.knownKeys = [:]
+    }
+
+    /// Private init for import
+    private init(senderKey: Data, chainId: Data, chainKey: Data, n: Int, knownKeys: [Int: Data]) {
+        self.senderKey = senderKey
+        self.chainId = chainId
+        self.chainKey = chainKey
+        self.n = n
+        self.knownKeys = knownKeys
+    }
+
+    // MARK: - Encrypt
+
+    /// Encrypt with current chain key.
+    /// Returns (chainId hex, n, ciphertext with tag, nonce).
+    /// Matches Python: SenderKeyState.encrypt(plaintext)
+    func encrypt(_ plaintext: Data) throws -> (chainIdHex: String, n: Int, ciphertext: Data, nonce: Data) {
+        let (newCK, messageKey) = CryptoUtils.kdfCK(chainKey: chainKey)
+        chainKey = newCK
+
+        let nonce = Data.randomBytes(12)
+        // AAD = chainId + bigEndian(UInt32(n))
+        let aad = chainId + UInt32(n).bigEndianData
+        let ctWithTag = try CryptoUtils.aesGcmEncrypt(plaintext, key: messageKey, nonce: nonce, aad: aad)
+
+        let result = (chainIdHex: chainId.hexString, n: n, ciphertext: ctWithTag, nonce: nonce)
+        n += 1
+        return result
+    }
+
+    // MARK: - Decrypt
+
+    /// Decrypt a group message. Fast-forwards the chain if needed.
+    /// State is snapshotted before modification and restored on failure.
+    /// Matches Python: SenderKeyState.decrypt(chain_id_hex, n, ciphertext, nonce)
+    func decrypt(chainIdHex: String, n: Int, ciphertext: Data, nonce: Data) throws -> Data {
+        guard let expectedChainId = Data(hexString: chainIdHex) else {
+            throw CryptoError.senderKeyError("Invalid chain ID hex")
+        }
+        guard expectedChainId == chainId else {
+            throw CryptoError.senderKeyError("Chain ID mismatch")
+        }
+
+        if n - self.n > Constants.maxSenderKeySkip {
+            throw CryptoError.senderKeyError("Sender key skip too large (\(n - self.n) > \(Constants.maxSenderKeySkip))")
+        }
+
+        // Snapshot before fast-forward
+        let snapChainKey = chainKey
+        let snapN = self.n
+        let snapKnown = knownKeys
+
+        do {
+            // Fast-forward the chain to reach message n
+            while self.n <= n {
+                let (newCK, mk) = CryptoUtils.kdfCK(chainKey: chainKey)
+                chainKey = newCK
+                knownKeys[self.n] = mk
+                self.n += 1
+            }
+
+            guard let mk = knownKeys.removeValue(forKey: n) else {
+                throw CryptoError.senderKeyError("Message key for n=\(n) not available")
+            }
+
+            let aad = chainId + UInt32(n).bigEndianData
+            return try CryptoUtils.aesGcmDecrypt(ciphertext, key: mk, nonce: nonce, aad: aad)
+        } catch {
+            // Restore state on failure
+            chainKey = snapChainKey
+            self.n = snapN
+            knownKeys = snapKnown
+            throw error
+        }
+    }
+
+    // MARK: - Key Export/Import
+
+    /// Export sender key for distribution to group members.
+    /// Matches Python: SenderKeyState.export_key()
+    func exportKey() -> Data {
+        let dict: [String: Any] = ["sender_key": senderKey.hexString]
+        return try! JSONSerialization.data(withJSONObject: dict)
+    }
+
+    /// Initialize a receiving SenderKeyState from an exported key.
+    /// Matches Python: SenderKeyState.from_key(exported_key)
+    static func fromKey(_ exportedKey: Data) throws -> SenderKeyState {
+        guard let dict = try JSONSerialization.jsonObject(with: exportedKey) as? [String: Any],
+              let senderKeyHex = dict["sender_key"] as? String,
+              let senderKey = Data(hexString: senderKeyHex) else {
+            throw CryptoError.stateImportFailed("Invalid sender key export")
+        }
+        return SenderKeyState(senderKey: senderKey)
+    }
+
+    // MARK: - Full State Export/Import
+
+    /// Serialize full state for persistent storage.
+    /// Matches Python: SenderKeyState.export_state()
+    func exportState() -> Data {
+        var knownKeysDict: [String: String] = [:]
+        for (k, v) in knownKeys {
+            knownKeysDict[String(k)] = v.hexString
+        }
+        let state: [String: Any] = [
+            "sender_key": senderKey.hexString,
+            "chain_id": chainId.hexString,
+            "chain_key": chainKey.hexString,
+            "n": n,
+            "known_keys": knownKeysDict,
+        ]
+        return try! JSONSerialization.data(withJSONObject: state)
+    }
+
+    /// Deserialize full state.
+    /// Matches Python: SenderKeyState.import_state(data)
+    static func importState(_ data: Data) throws -> SenderKeyState {
+        guard let state = try JSONSerialization.jsonObject(with: data) as? [String: Any],
+              let senderKeyHex = state["sender_key"] as? String,
+              let senderKey = Data(hexString: senderKeyHex),
+              let chainIdHex = state["chain_id"] as? String,
+              let chainId = Data(hexString: chainIdHex),
+              let chainKeyHex = state["chain_key"] as? String,
+              let chainKey = Data(hexString: chainKeyHex),
+              let n = state["n"] as? Int else {
+            throw CryptoError.stateImportFailed("Invalid sender key state")
+        }
+
+        var knownKeys: [Int: Data] = [:]
+        if let knownKeysDict = state["known_keys"] as? [String: String] {
+            for (k, v) in knownKeysDict {
+                if let idx = Int(k), let data = Data(hexString: v) {
+                    knownKeys[idx] = data
+                }
+            }
+        }
+
+        return SenderKeyState(
+            senderKey: senderKey,
+            chainId: chainId,
+            chainKey: chainKey,
+            n: n,
+            knownKeys: knownKeys
+        )
+    }
+}