import Foundation import CryptoKit /// Notification types from the server enum ChatNotification { case newMessage(data: [String: Any]) case messagesRead(data: [String: Any]) case messageDeleted(data: [String: Any]) case conversationCreated(data: [String: Any]) case memberAdded(data: [String: Any]) case memberRemoved(data: [String: Any]) case userOnline(userId: String) case userOffline(userId: String) case onlineUsers(userIds: [String]) case groupInvitation(data: [String: Any]) case conversationRenamed(data: [String: Any]) case sessionReset(data: [String: Any]) case connectionStateChanged(connected: Bool) } /// Main chat client — handles all server communication and crypto operations. /// Thread-safe via Swift actor isolation. /// Port of Python ChatClient class from chat_core.py actor ChatClient { // MARK: - Connection let connectionManager = ConnectionManager() private(set) var isConnected = false private(set) var sessionToken: String? private(set) var userId: String? private(set) var username: String = "" private(set) var email: String = "" private(set) var loginRejected = false // MARK: - Keys private var rsaPrivate: SecKey? private var rsaPublic: SecKey? private(set) var identityPrivate: Curve25519.Signing.PrivateKey? private(set) var identityPublic: Curve25519.Signing.PublicKey? private var spkPrivate: Curve25519.KeyAgreement.PrivateKey? private var spkId: String = "" private var prevSpkPrivate: Curve25519.KeyAgreement.PrivateKey? private var prevSpkId: String = "" private var opkPrivates: [String: Curve25519.KeyAgreement.PrivateKey] = [:] // MARK: - Sessions & Sender Keys private var sessions: [String: DoubleRatchet] = [:] // "userId:deviceId" -> ratchet private var senderKeyStates: [String: SenderKeyState] = [:] // convId -> own sender key private var recvSenderKeys: [String: SenderKeyState] = [:] // "convId:senderId:deviceId" -> their key // MARK: - Derived Keys private var cacheKey: Data? // for encrypting message cache private var localKey: Data? // for encrypting session/sender key files // MARK: - Multi-Device private(set) var deviceId: String? // MARK: - Caches private var userCache: [String: User] = [:] private var deviceBundleCache: [String: (timestamp: Date, bundles: [DeviceBundle])] = [:] // MARK: - Request/Response Tracking private var pendingRequests: [String: CheckedContinuation<[String: Any], Error>] = [:] private var listenerTask: Task? // MARK: - Notification Stream private var notificationContinuation: AsyncStream.Continuation? nonisolated let notifications: AsyncStream // MARK: - Init init() { var continuation: AsyncStream.Continuation! notifications = AsyncStream { cont in continuation = cont } notificationContinuation = continuation } // MARK: - Connection func connect(host: String = Constants.defaultHost, port: UInt16 = Constants.defaultPort, useTLS: Bool = false, tlsInsecure: Bool = false) async throws { try await connectionManager.connect(host: host, port: port, useTLS: useTLS, tlsInsecure: tlsInsecure) isConnected = true notificationContinuation?.yield(.connectionStateChanged(connected: true)) } func disconnect() async { listenerTask?.cancel() listenerTask = nil await connectionManager.disconnect() isConnected = false // Fail all pending requests let pending = pendingRequests pendingRequests.removeAll() for (_, cont) in pending { cont.resume(throwing: NetworkError.notConnected) } notificationContinuation?.yield(.connectionStateChanged(connected: false)) } // MARK: - Send and Receive /// Send a request and wait for the matching response. func sendAndReceive(type: String, timeout: TimeInterval = 30, params: [String: Any] = [:]) async -> [String: Any] { let requestId = ProtocolHandler.newRequestId() do { let response: [String: Any] = try await withCheckedThrowingContinuation { continuation in pendingRequests[requestId] = continuation Task { do { try await connectionManager.sendMessage(type: type, requestId: requestId, params: params) } catch { if let cont = pendingRequests.removeValue(forKey: requestId) { cont.resume(throwing: error) } } } } return response } catch { pendingRequests.removeValue(forKey: requestId) return [ "type": type, "status": "error", "data": ["message": error.localizedDescription] ] } } // MARK: - Background Listener func startBackgroundListener() { listenerTask = Task { [weak self] in guard let self = self else { return } await self.backgroundListenerLoop() } } private func backgroundListenerLoop() async { while !Task.isCancelled { do { guard let msg = try await connectionManager.readMessage() else { // EOF — connection closed await handleDisconnect() break } await routeMessage(msg) } catch { await handleDisconnect() break } } } private func handleDisconnect() { isConnected = false // Fail all pending futures let pending = pendingRequests pendingRequests.removeAll() for (_, cont) in pending { cont.resume(throwing: NetworkError.notConnected) } notificationContinuation?.yield(.connectionStateChanged(connected: false)) } private func routeMessage(_ msg: [String: Any]) { let msgType = msg["type"] as? String ?? "" // Notification types (no request_id expected from client) let notificationTypes = Set([ "new_message", "messages_read", "message_deleted", "conversation_created", "member_added", "member_removed", "user_online", "user_offline", "online_users", "group_invitation", "conversation_renamed", "session_reset" ]) if notificationTypes.contains(msgType) { let data = msg["data"] as? [String: Any] ?? msg switch msgType { case "new_message": notificationContinuation?.yield(.newMessage(data: data)) case "messages_read": notificationContinuation?.yield(.messagesRead(data: data)) case "message_deleted": notificationContinuation?.yield(.messageDeleted(data: data)) case "conversation_created": notificationContinuation?.yield(.conversationCreated(data: data)) case "member_added": notificationContinuation?.yield(.memberAdded(data: data)) case "member_removed": notificationContinuation?.yield(.memberRemoved(data: data)) case "user_online": if let uid = data["user_id"] as? String { notificationContinuation?.yield(.userOnline(userId: uid)) } case "user_offline": if let uid = data["user_id"] as? String { notificationContinuation?.yield(.userOffline(userId: uid)) } case "online_users": if let uids = data["user_ids"] as? [String] { notificationContinuation?.yield(.onlineUsers(userIds: uids)) } case "group_invitation": notificationContinuation?.yield(.groupInvitation(data: data)) case "conversation_renamed": notificationContinuation?.yield(.conversationRenamed(data: data)) case "session_reset": notificationContinuation?.yield(.sessionReset(data: data)) default: break } } else { // Response to a pending request if let requestId = msg["request_id"] as? String, let cont = pendingRequests.removeValue(forKey: requestId) { cont.resume(returning: msg) } } } // MARK: - User Info Cache func getUserInfo(userId: String = "", userEmail: String = "") async -> User? { if !userId.isEmpty, let cached = userCache[userId] { return cached } var params: [String: Any] = [:] if !userId.isEmpty { params["user_id"] = userId } else if !userEmail.isEmpty { params["email"] = userEmail } else { return nil } let resp = await sendAndReceive(type: "get_user_info", params: params) guard resp.string(for: "status") == "ok", let data = resp.dict(for: "data") else { return nil } var ikData: Data? if let ikB64 = data["identity_key"] as? String { ikData = try? ProtocolHandler.decodeBinary(ikB64) } let user = User( id: data.string(for: "user_id") ?? "", username: data.string(for: "username") ?? "", email: data.string(for: "email") ?? "", identityKey: ikData ) userCache[user.id] = user return user } // MARK: - Registration func register(username: String, password: String, email: String) async -> (success: Bool, message: String) { self.username = username self.email = email var pwdBytes = Array(password.utf8) defer { pwdBytes.withUnsafeMutableBytes { ptr in _ = memset(ptr.baseAddress!, 0, ptr.count) } } let pwdData = Data(pwdBytes) do { // RSA keys let (rsaPriv, rsaPub, err) = KeyStorage.loadRSAKeys(email: email, password: pwdData) if let rsaPriv = rsaPriv, let rsaPub = rsaPub { self.rsaPrivate = rsaPriv self.rsaPublic = rsaPub } else { let (newPriv, newPub) = try RSACrypto.generateKeypair() try KeyStorage.saveRSAKeys(email: email, privateKey: newPriv, publicKey: newPub, password: pwdData) self.rsaPrivate = newPriv self.rsaPublic = newPub } // Ed25519 identity keys let (edPriv, edPub) = KeyStorage.loadIdentityKeys(email: email, password: pwdData) if let edPriv = edPriv, let edPub = edPub { self.identityPrivate = edPriv self.identityPublic = edPub } else { let (newPriv, newPub) = Ed25519Crypto.generateKeypair() try KeyStorage.saveIdentityKeys(email: email, privateKey: newPriv, publicKey: newPub, password: pwdData) self.identityPrivate = newPriv self.identityPublic = newPub } self.cacheKey = CryptoUtils.deriveSelfEncryptionKey(identityPrivateRaw: identityPrivate!.rawData) self.localKey = CryptoUtils.deriveLocalStorageKey(identityPrivateRaw: identityPrivate!.rawData) } catch { return (false, "Key generation failed: \(error.localizedDescription)") } // Send registration request let pubPem = String(data: try! RSACrypto.serializePublicKey(rsaPublic!), encoding: .utf8)! let ikB64 = ProtocolHandler.encodeBinary(Ed25519Crypto.serializePublic(identityPublic!)) let resp = await sendAndReceive(type: "register", params: [ "username": username, "public_key": pubPem, "email": email, "identity_key": ikB64, ]) guard resp.string(for: "status") == "ok" else { let msg = resp.dict(for: "data")?.string(for: "message") ?? "Registration failed" return (false, msg) } let data = resp.dict(for: "data") ?? [:] if let code = data.string(for: "code") { return (true, code) } return (true, data.string(for: "message") ?? "Check your email for the code.") } func confirmRegistration(email: String, username: String, code: String) async -> (success: Bool, message: String) { let resp = await sendAndReceive(type: "register_confirm", params: [ "email": email, "code": code, ]) guard resp.string(for: "status") == "ok" else { let msg = resp.dict(for: "data")?.string(for: "message") ?? "Confirmation failed" return (false, msg) } // Upload prekeys await generateAndUploadPrekeys() let uid = resp.dict(for: "data")?.string(for: "user_id") ?? "" return (true, "Registered as '\(username)' (ID: \(uid))") } // MARK: - Prekeys private func generateAndUploadPrekeys(keepSPK: Bool = false) async { guard let identityPrivate = identityPrivate else { return } do { let spkData: [String: Any] if keepSPK, let spkPriv = spkPrivate, !spkId.isEmpty { let spkPubBytes = X25519Crypto.serializePublic(spkPriv.publicKey) let sig = try Ed25519Crypto.sign(identityPrivate, data: spkPubBytes) spkData = [ "id": spkId, "public_key": ProtocolHandler.encodeBinary(spkPubBytes), "signature": ProtocolHandler.encodeBinary(sig), ] } else { // Save current as previous (grace period) if let spkPriv = spkPrivate, !spkId.isEmpty { prevSpkPrivate = spkPriv prevSpkId = spkId try? KeyStorage.savePrevSPK(email: email, privateKey: spkPriv, spkId: spkId) } let spk = try X3DH.generateSignedPrekey(identityPrivate: identityPrivate) self.spkPrivate = spk.privateKey self.spkId = spk.id try? KeyStorage.saveSPK(email: email, privateKey: spk.privateKey, spkId: spk.id) spkData = [ "id": spk.id, "public_key": ProtocolHandler.encodeBinary(X25519Crypto.serializePublic(spk.publicKey)), "signature": ProtocolHandler.encodeBinary(spk.signature), ] } // Generate OPKs let opks = X3DH.generateOneTimePrekeys(count: Constants.opkBatchSize) for opk in opks { opkPrivates[opk.id] = opk.privateKey try? KeyStorage.saveOPKPrivate(email: email, opkId: opk.id, privateKey: opk.privateKey) } let otpData = opks.map { opk -> [String: Any] in [ "id": opk.id, "public_key": ProtocolHandler.encodeBinary(X25519Crypto.serializePublic(opk.publicKey)), ] } _ = await sendAndReceive(type: "upload_prekeys", params: [ "signed_prekey": spkData, "one_time_prekeys": otpData, ]) } catch { // Log error but don't fail print("Prekey generation error: \(error)") } } private func ensurePrekeys() async { let resp = await sendAndReceive(type: "get_prekey_count") guard resp.string(for: "status") == "ok", let data = resp.dict(for: "data") else { return } let count = data.int(for: "count") ?? 0 let spkCreatedAt = data.string(for: "spk_created_at") ?? "" var needNewSPK = false if !spkCreatedAt.isEmpty { let formatter = ISO8601DateFormatter() formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] if let created = formatter.date(from: spkCreatedAt) ?? ISO8601DateFormatter().date(from: spkCreatedAt) { let ageDays = Calendar.current.dateComponents([.day], from: created, to: Date()).day ?? 0 if ageDays >= Constants.spkRotationDays { needNewSPK = true } } } if count < Constants.opkReplenishThreshold || needNewSPK { await generateAndUploadPrekeys() } } // MARK: - Login func login(email: String, password: String) async -> (success: Bool, message: String) { self.email = email var pwdBytes = Array(password.utf8) defer { pwdBytes.withUnsafeMutableBytes { ptr in _ = memset(ptr.baseAddress!, 0, ptr.count) } } let pwdData = Data(pwdBytes) // Load RSA keys let (rsaPriv, rsaPub, err) = KeyStorage.loadRSAKeys(email: email, password: pwdData) guard let rsaPriv = rsaPriv, let rsaPub = rsaPub else { return (false, err ?? "No local keys found. Register first.") } self.rsaPrivate = rsaPriv self.rsaPublic = rsaPub // Load identity keys let (edPriv, edPub) = KeyStorage.loadIdentityKeys(email: email, password: pwdData) if let edPriv = edPriv, let edPub = edPub { self.identityPrivate = edPriv self.identityPublic = edPub self.cacheKey = CryptoUtils.deriveSelfEncryptionKey(identityPrivateRaw: edPriv.rawData) self.localKey = CryptoUtils.deriveLocalStorageKey(identityPrivateRaw: edPriv.rawData) } // Load SPK let (spkP, spkI) = KeyStorage.loadSPK(email: email) if let spkP = spkP { self.spkPrivate = spkP self.spkId = spkI ?? "" } // Load previous SPK (grace period) let (prevP, prevI) = KeyStorage.loadPrevSPK(email: email) if let prevP = prevP { self.prevSpkPrivate = prevP self.prevSpkId = prevI ?? "" } // Load device ID self.deviceId = KeyStorage.loadDeviceId(email: email) // RSA challenge-response login let startResp = await sendAndReceive(type: "login_start", params: ["email": email]) guard startResp.string(for: "status") == "ok", let startData = startResp.dict(for: "data"), let challengeB64 = startData.string(for: "challenge") else { let msg = startResp.dict(for: "data")?.string(for: "message") ?? "Login failed" return (false, msg) } let challengeData: Data do { challengeData = try ProtocolHandler.decodeBinary(challengeB64) } catch { return (false, "Invalid challenge data") } let signature: Data do { signature = try RSACrypto.sign(rsaPriv, data: challengeData) } catch { return (false, "RSA signing failed: \(error.localizedDescription)") } var finishParams: [String: Any] = [ "email": email, "signature": ProtocolHandler.encodeBinary(signature), "client_version": Constants.version, ] if let deviceId = deviceId { finishParams["device_id"] = deviceId } let finishResp = await sendAndReceive(type: "login_finish", params: finishParams) guard finishResp.string(for: "status") == "ok", let finishData = finishResp.dict(for: "data") else { let msg = finishResp.dict(for: "data")?.string(for: "message") ?? "Login failed" loginRejected = true return (false, msg) } self.userId = finishData.string(for: "user_id") self.username = finishData.string(for: "username") ?? "" self.sessionToken = finishData.string(for: "session_token") // Save device ID from server if let newDeviceId = finishData.string(for: "device_id") { self.deviceId = newDeviceId try? KeyStorage.saveDeviceId(email: email, deviceId: newDeviceId) } // Start background listener startBackgroundListener() // Handle online_users if included if let onlineUserIds = finishData["online_user_ids"] as? [String] { notificationContinuation?.yield(.onlineUsers(userIds: onlineUserIds)) } // Ensure prekeys in background Task { await ensurePrekeys() } return (true, "Logged in as \(username)") } // MARK: - Reconnect func reconnect() async -> Bool { guard rsaPrivate != nil else { return false } await disconnect() do { try await connect() } catch { return false } // RSA challenge-response with in-memory keys let startResp = await sendAndReceive(type: "login_start", params: ["email": email]) guard startResp.string(for: "status") == "ok", let startData = startResp.dict(for: "data"), let challengeB64 = startData.string(for: "challenge"), let challengeData = try? ProtocolHandler.decodeBinary(challengeB64), let signature = try? RSACrypto.sign(rsaPrivate!, data: challengeData) else { return false } var finishParams: [String: Any] = [ "email": email, "signature": ProtocolHandler.encodeBinary(signature), "client_version": Constants.version, ] if let deviceId = deviceId { finishParams["device_id"] = deviceId } let finishResp = await sendAndReceive(type: "login_finish", params: finishParams) guard finishResp.string(for: "status") == "ok" else { return false } startBackgroundListener() return true } // MARK: - Device Bundles private func getDeviceBundles(userId: String) async throws -> [DeviceBundle] { // Check cache (5-min TTL) if let cached = deviceBundleCache[userId], Date().timeIntervalSince(cached.timestamp) < Constants.deviceBundleCacheTTL { return cached.bundles } let resp = await sendAndReceive(type: "get_key_bundle", params: ["user_id": userId]) guard resp.string(for: "status") == "ok", let data = resp.dict(for: "data") else { throw ChatError.operationFailed("Failed to get key bundle") } var bundles: [DeviceBundle] = [] // Per-device bundles (new format) if let deviceBundlesRaw = data["device_bundles"] as? [[String: Any]] { for bundleDict in deviceBundlesRaw { if let bundle = try? DeviceBundle.fromDict(bundleDict) { bundles.append(bundle) } } } // Legacy single bundle else if let ikHex = data["identity_key"] as? String { let bundle = try DeviceBundle.fromDict(data) bundles.append(bundle) } deviceBundleCache[userId] = (Date(), bundles) return bundles } // MARK: - Session Management private func getOrCreateSession( peerUserId: String, peerDeviceId: String, bundle: DeviceBundle ) async throws -> DoubleRatchet { let sessionKey = "\(peerUserId):\(peerDeviceId)" // Check memory if let session = sessions[sessionKey] { return session } // Check disk if let session = KeyStorage.loadSession( email: email, peerUserId: peerUserId, localKey: localKey, peerDeviceId: peerDeviceId ) { sessions[sessionKey] = session return session } // Create new via X3DH let remoteIkEd = try Ed25519Crypto.loadPublic(bundle.identityKey) let spkRemote = try X25519Crypto.loadPublic(bundle.spk) var opkRemote: Curve25519.KeyAgreement.PublicKey? if let opkData = bundle.opk { opkRemote = try X25519Crypto.loadPublic(opkData) } let (sharedSecret, ekPriv, ekPub) = try X3DH.initiate( ikPrivateEd: identityPrivate!, ikPublicRemoteEd: remoteIkEd, spkRemote: spkRemote, spkSignature: bundle.spkSignature, opkRemote: opkRemote ) let ratchet = try DoubleRatchet.initAlice(sharedSecret: sharedSecret, bobSpkPub: spkRemote) // Build X3DH header for first message var x3dhHeader: [String: Any] = [ "ik": Ed25519Crypto.serializePublic(identityPublic!).hexString, "ek": X25519Crypto.serializePublic(ekPub).hexString, "spk_id": bundle.spkId, ] if let opkId = bundle.opkId { x3dhHeader["opk_id"] = opkId } ratchet.x3dhHeader = x3dhHeader sessions[sessionKey] = ratchet try? KeyStorage.saveSession(email: email, peerUserId: peerUserId, ratchet: ratchet, localKey: localKey, peerDeviceId: peerDeviceId) return ratchet } // MARK: - X3DH Response (Bob Side) private func processX3DHHeader( senderId: String, x3dhHeader: [String: Any], senderDeviceId: String, spkOverride: Curve25519.KeyAgreement.PrivateKey? = nil ) throws -> DoubleRatchet { guard let ikHex = x3dhHeader["ik"] as? String, let ikData = Data(hexString: ikHex), let ekHex = x3dhHeader["ek"] as? String, let ekData = Data(hexString: ekHex), let spkIdStr = x3dhHeader["spk_id"] as? String else { throw CryptoError.x3dhFailed("Invalid X3DH header") } let remoteIkEd = try Ed25519Crypto.loadPublic(ikData) let ekRemote = try X25519Crypto.loadPublic(ekData) // Determine which SPK to use let spkToUse: Curve25519.KeyAgreement.PrivateKey if let override = spkOverride { spkToUse = override } else if spkIdStr == spkId, let spk = spkPrivate { spkToUse = spk } else if spkIdStr == prevSpkId, let prevSpk = prevSpkPrivate { spkToUse = prevSpk } else { throw CryptoError.x3dhFailed("SPK \(spkIdStr) not found") } // OPK var opkPriv: Curve25519.KeyAgreement.PrivateKey? if let opkIdStr = x3dhHeader["opk_id"] as? String { opkPriv = opkPrivates[opkIdStr] ?? KeyStorage.loadOPKPrivate(email: email, opkId: opkIdStr) if opkPriv != nil { opkPrivates.removeValue(forKey: opkIdStr) KeyStorage.deleteOPKPrivate(email: email, opkId: opkIdStr) } } let sharedSecret = try X3DH.respond( ikPrivateEd: identityPrivate!, spkPrivate: spkToUse, ikRemoteEd: remoteIkEd, ekRemote: ekRemote, opkPrivate: opkPriv ) let ratchet = DoubleRatchet.initBob( sharedSecret: sharedSecret, spkPair: (spkToUse, spkToUse.publicKey) ) let sessionKey = "\(senderId):\(senderDeviceId)" sessions[sessionKey] = ratchet try? KeyStorage.saveSession(email: email, peerUserId: senderId, ratchet: ratchet, localKey: localKey, peerDeviceId: senderDeviceId) return ratchet } // MARK: - Send Message func sendMessage(convId: String, text: String, members: [ConversationMember], replyTo: String? = nil) async -> (success: Bool, message: String) { let isGroup = members.count > 2 if isGroup { return await sendGroupMessage(convId: convId, text: text, members: members, replyTo: replyTo) } else { return await sendDM(convId: convId, text: text, members: members, replyTo: replyTo) } } // MARK: - Send DM private func sendDM(convId: String, text: String, members: [ConversationMember], replyTo: String? = nil) async -> (success: Bool, message: String) { guard let identityPrivate = identityPrivate else { return (false, "Identity key not loaded") } let plaintext = Data(text.utf8) var payload: [String: Any] = ["text": text] if let replyTo = replyTo { payload["reply_to"] = replyTo } var recipients: [[String: Any]] = [] // Encrypt for each member's devices for member in members where member.userId != userId { do { let bundles = try await getDeviceBundles(userId: member.userId) for bundle in bundles { let ratchet = try await getOrCreateSession( peerUserId: member.userId, peerDeviceId: bundle.deviceId, bundle: bundle ) // Consume X3DH header if present (first message only) let x3dhHeader = ratchet.x3dhHeader ratchet.x3dhHeader = nil let encrypted = try ratchet.encrypt(plaintext) try? KeyStorage.saveSession(email: email, peerUserId: member.userId, ratchet: ratchet, localKey: localKey, peerDeviceId: bundle.deviceId) var recipientEntry: [String: Any] = [ "user_id": member.userId, "device_id": bundle.deviceId, "ciphertext": ProtocolHandler.encodeBinary(encrypted.ciphertext), "nonce": ProtocolHandler.encodeBinary(encrypted.nonce), "ratchet_header": encrypted.header, ] if let x3dh = x3dhHeader { recipientEntry["x3dh_header"] = x3dh } recipients.append(recipientEntry) } } catch { return (false, "Encryption failed for \(member.username): \(error.localizedDescription)") } } // Self-encrypted copy let selfKey = CryptoUtils.deriveSelfEncryptionKey(identityPrivateRaw: identityPrivate.rawData) if let (_, nonce, ct, tag) = try? CryptoUtils.aesEncrypt(plaintext, key: selfKey) { let selfCiphertext = ct + tag let dummyHeader: [String: Any] = [ "dh_pub": String(repeating: "00", count: 32), "n": 0, "pn": 0, ] recipients.append([ "user_id": userId!, "device_id": Constants.selfDeviceId, "ciphertext": ProtocolHandler.encodeBinary(selfCiphertext), "nonce": ProtocolHandler.encodeBinary(nonce), "ratchet_header": dummyHeader, ]) } // Build ratchet header for message table (use first recipient's or dummy) let ratchetHeader: [String: Any] if let first = recipients.first { ratchetHeader = first["ratchet_header"] as? [String: Any] ?? [:] } else { ratchetHeader = ["dh_pub": String(repeating: "00", count: 32), "n": 0, "pn": 0] } var params: [String: Any] = [ "conversation_id": convId, "ratchet_header": ratchetHeader, "recipients": recipients, ] if let replyTo = replyTo { params["reply_to"] = replyTo } let resp = await sendAndReceive(type: "send_message", params: params) guard resp.string(for: "status") == "ok" else { let msg = resp.dict(for: "data")?.string(for: "message") ?? "Send failed" return (false, msg) } // Cache the sent message if let msgData = resp.dict(for: "data"), let messageId = msgData.string(for: "message_id") { var cacheEntry = payload cacheEntry["sender_id"] = userId cacheEntry["sender_username"] = username cacheEntry["created_at"] = ISO8601DateFormatter().string(from: Date()) try? MessageCache.save(email: email, convId: convId, messages: [cacheEntry], cacheKey: cacheKey) } return (true, "Message sent") } // MARK: - Send Group Message private func sendGroupMessage(convId: String, text: String, members: [ConversationMember], replyTo: String? = nil) async -> (success: Bool, message: String) { guard let identityPrivate = identityPrivate, let userId = userId, let deviceId = deviceId else { return (false, "Not properly logged in") } // Get or create sender key for this group var senderKeyState = senderKeyStates[convId] if senderKeyState == nil { senderKeyState = KeyStorage.loadSenderKeyState(email: email, convId: convId, localKey: localKey) } var needDistribute = false if senderKeyState == nil { senderKeyState = SenderKeyState() needDistribute = true } senderKeyStates[convId] = senderKeyState // Distribute sender key if new if needDistribute { await distributeSenderKey(convId: convId, members: members) } // Encrypt with sender key let plaintext = Data(text.utf8) do { let encrypted = try senderKeyState!.encrypt(plaintext) try? KeyStorage.saveSenderKeyState(email: email, convId: convId, state: senderKeyState!, localKey: localKey) // Build recipients (same ciphertext for all) var recipients: [[String: Any]] = [] for member in members where member.userId != userId { recipients.append([ "user_id": member.userId, "device_id": Constants.selfDeviceId, // group messages use sentinel "ciphertext": ProtocolHandler.encodeBinary(encrypted.ciphertext), "nonce": ProtocolHandler.encodeBinary(encrypted.nonce), ]) } // Self copy let selfKey = CryptoUtils.deriveSelfEncryptionKey(identityPrivateRaw: identityPrivate.rawData) if let (_, nonce, ct, tag) = try? CryptoUtils.aesEncrypt(plaintext, key: selfKey) { recipients.append([ "user_id": userId, "device_id": Constants.selfDeviceId, "ciphertext": ProtocolHandler.encodeBinary(ct + tag), "nonce": ProtocolHandler.encodeBinary(nonce), ]) } let dummyHeader: [String: Any] = [ "dh_pub": String(repeating: "00", count: 32), "n": 0, "pn": 0, ] var params: [String: Any] = [ "conversation_id": convId, "ratchet_header": dummyHeader, "recipients": recipients, "sender_chain_id": ProtocolHandler.encodeBinary(encrypted.ciphertext.prefix(0)), // placeholder ] // Include sender key metadata for group routing params["sender_chain_id"] = encrypted.chainIdHex params["sender_chain_n"] = encrypted.n if let replyTo = replyTo { params["reply_to"] = replyTo } let resp = await sendAndReceive(type: "send_message", params: params) guard resp.string(for: "status") == "ok" else { let msg = resp.dict(for: "data")?.string(for: "message") ?? "Send failed" return (false, msg) } return (true, "Message sent") } catch { return (false, "Encryption failed: \(error.localizedDescription)") } } // MARK: - Distribute Sender Key private func distributeSenderKey(convId: String, members: [ConversationMember]) async { guard let senderKeyState = senderKeyStates[convId], let userId = userId, let deviceId = deviceId else { return } let exportedKey = senderKeyState.exportKey() for member in members where member.userId != userId { do { let bundles = try await getDeviceBundles(userId: member.userId) for bundle in bundles { let ratchet = try await getOrCreateSession( peerUserId: member.userId, peerDeviceId: bundle.deviceId, bundle: bundle ) let x3dhHeader = ratchet.x3dhHeader ratchet.x3dhHeader = nil // Payload includes sender key + metadata let controlPayload: [String: Any] = [ "_sender_key": [ "conv_id": convId, "key": ProtocolHandler.encodeBinary(exportedKey), "sender_device_id": deviceId, ] ] let controlData = try JSONSerialization.data(withJSONObject: controlPayload) let encrypted = try ratchet.encrypt(controlData) try? KeyStorage.saveSession(email: email, peerUserId: member.userId, ratchet: ratchet, localKey: localKey, peerDeviceId: bundle.deviceId) var recipientEntry: [String: Any] = [ "user_id": member.userId, "device_id": bundle.deviceId, "ciphertext": ProtocolHandler.encodeBinary(encrypted.ciphertext), "nonce": ProtocolHandler.encodeBinary(encrypted.nonce), "ratchet_header": encrypted.header, ] if let x3dh = x3dhHeader { recipientEntry["x3dh_header"] = x3dh } let dummyHeader: [String: Any] = [ "dh_pub": String(repeating: "00", count: 32), "n": 0, "pn": 0, ] _ = await sendAndReceive(type: "send_message", params: [ "conversation_id": convId, "ratchet_header": dummyHeader, "recipients": [recipientEntry], ]) } } catch { print("Failed to distribute sender key to \(member.userId): \(error)") } } } // MARK: - Decrypt func decryptDMRecipientData( senderData: [String: Any], senderId: String, senderDeviceId: String ) -> Data? { guard let ctB64 = senderData["ciphertext"] as? String, let nonceB64 = senderData["nonce"] as? String, let ct = try? ProtocolHandler.decodeBinary(ctB64), let nonce = try? ProtocolHandler.decodeBinary(nonceB64) else { return nil } // Self-encrypted copy if senderDeviceId == Constants.selfDeviceId || senderId == userId { if let cacheKey = cacheKey { // ct = ciphertext + tag(16) guard ct.count >= 16 else { return nil } let ciphertext = ct.prefix(ct.count - 16) let tag = ct.suffix(16) return try? CryptoUtils.aesDecrypt(key: cacheKey, nonce: nonce, ciphertext: Data(ciphertext), tag: Data(tag)) } return nil } // Regular DM decryption let headerDict = senderData["ratchet_header"] as? [String: Any] let x3dhHeader = senderData["x3dh_header"] as? [String: Any] let sessionKey = "\(senderId):\(senderDeviceId)" var ratchet = sessions[sessionKey] ?? KeyStorage.loadSession(email: email, peerUserId: senderId, localKey: localKey, peerDeviceId: senderDeviceId) // Handle X3DH header (new session) if let x3dh = x3dhHeader { do { ratchet = try processX3DHHeader( senderId: senderId, x3dhHeader: x3dh, senderDeviceId: senderDeviceId ) } catch { // Try with previous SPK (grace period) if let prevSpk = prevSpkPrivate { ratchet = try? processX3DHHeader( senderId: senderId, x3dhHeader: x3dh, senderDeviceId: senderDeviceId, spkOverride: prevSpk ) } if ratchet == nil { return nil } } } guard let ratchet = ratchet, let header = headerDict else { return nil } do { let plaintext = try ratchet.decrypt(headerDict: header, ciphertext: ct, nonce: nonce) sessions[sessionKey] = ratchet try? KeyStorage.saveSession(email: email, peerUserId: senderId, ratchet: ratchet, localKey: localKey, peerDeviceId: senderDeviceId) // Check for sender key distribution (control message) if let jsonObj = try? JSONSerialization.jsonObject(with: plaintext) as? [String: Any], let senderKeyInfo = jsonObj["_sender_key"] as? [String: Any] { handleSenderKeyDistribution(senderKeyInfo, senderId: senderId) return nil // Control message } return plaintext } catch { return nil } } private func handleSenderKeyDistribution(_ info: [String: Any], senderId: String) { guard let convId = info["conv_id"] as? String, let keyB64 = info["key"] as? String, let keyData = try? ProtocolHandler.decodeBinary(keyB64) else { return } let senderDeviceId = info["sender_device_id"] as? String ?? Constants.selfDeviceId do { let senderKey = try SenderKeyState.fromKey(keyData) let stateKey = "\(convId):\(senderId):\(senderDeviceId)" recvSenderKeys[stateKey] = senderKey try? KeyStorage.saveRecvSenderKey( email: email, convId: convId, senderId: senderId, senderDeviceId: senderDeviceId, state: senderKey, localKey: localKey ) } catch { print("Failed to import sender key: \(error)") } } func decryptNotification(_ data: [String: Any]) -> Message? { guard let senderId = data.string(for: "sender_id"), let conversationId = data.string(for: "conversation_id"), let messageId = data.string(for: "message_id") else { return nil } let senderDeviceId = data.string(for: "sender_device_id") ?? Constants.selfDeviceId // Find our device's entry var recipientData: [String: Any]? if let deviceEntries = data["device_entries"] as? [[String: Any]] { recipientData = deviceEntries.first(where: { ($0["device_id"] as? String) == deviceId || ($0["device_id"] as? String) == Constants.selfDeviceId }) } // Fallback: use data directly if it has ciphertext if recipientData == nil, data["ciphertext"] != nil { recipientData = data } guard let recipientData = recipientData else { return nil } // Try DM decryption if let plaintext = decryptDMRecipientData( senderData: recipientData, senderId: senderId, senderDeviceId: senderDeviceId ) { let text = String(data: plaintext, encoding: .utf8) // Parse JSON payload var messageText = text var replyTo: String? var file: FileInfo? if let jsonObj = try? JSONSerialization.jsonObject(with: plaintext) as? [String: Any] { messageText = jsonObj["text"] as? String replyTo = jsonObj["reply_to"] as? String if let fileDict = jsonObj["file"] as? [String: Any] { file = FileInfo( fileId: fileDict["file_id"] as? String ?? "", aesKey: fileDict["aes_key"] as? String ?? "", iv: fileDict["iv"] as? String ?? "", filename: fileDict["filename"] as? String ?? "", size: fileDict["size"] as? Int ?? 0, mimeType: fileDict["mime_type"] as? String ?? "" ) } } let createdAt = data.string(for: "created_at").flatMap { ISO8601DateFormatter().date(from: $0) } ?? Date() let senderUsername = data.string(for: "sender_username") ?? userCache[senderId]?.username ?? "Unknown" return Message( id: messageId, conversationId: conversationId, senderId: senderId, senderUsername: senderUsername, createdAt: createdAt, text: messageText, replyTo: replyTo, imageFileId: data.string(for: "image_file_id"), file: file, isDeleted: false, readBy: [] ) } return nil } // MARK: - Conversations func listConversations() async -> [Conversation] { let resp = await sendAndReceive(type: "list_conversations") guard resp.string(for: "status") == "ok", let data = resp.dict(for: "data"), let convList = data["conversations"] as? [[String: Any]] else { return [] } return convList.compactMap { dict -> Conversation? in guard let id = dict.string(for: "id") else { return nil } let membersRaw = dict["members"] as? [[String: Any]] ?? [] let members = membersRaw.compactMap { m -> ConversationMember? in guard let uid = m.string(for: "user_id"), let uname = m.string(for: "username"), let uemail = m.string(for: "email") else { return nil } return ConversationMember(userId: uid, username: uname, email: uemail) } let unreadCount = dict.int(for: "unread_count") ?? 0 return Conversation( id: id, name: dict.string(for: "name"), members: members, createdBy: dict.string(for: "created_by"), avatarFile: dict.string(for: "avatar_file"), unreadCount: unreadCount, isFavorite: false, lastMessageTime: nil ) } } func createConversation(emails: [String], name: String? = nil) async -> (convId: String?, message: String) { var params: [String: Any] = ["emails": emails] if let name = name { params["name"] = name } let resp = await sendAndReceive(type: "create_conversation", params: params) guard resp.string(for: "status") == "ok", let data = resp.dict(for: "data"), let convId = data.string(for: "conversation_id") else { let msg = resp.dict(for: "data")?.string(for: "message") ?? "Failed to create conversation" return (nil, msg) } return (convId, "Conversation created") } func findConversation(email: String) async -> String? { let resp = await sendAndReceive(type: "find_conversation", params: ["email": email]) guard resp.string(for: "status") == "ok", let data = resp.dict(for: "data") else { return nil } return data.string(for: "conversation_id") } // MARK: - Messages func getMessages(convId: String, limit: Int = 50, offset: Int = 0) async -> [Message] { let resp = await sendAndReceive(type: "get_messages", params: [ "conversation_id": convId, "limit": limit, "offset": offset, ]) guard resp.string(for: "status") == "ok", let data = resp.dict(for: "data"), let messagesRaw = data["messages"] as? [[String: Any]] else { return [] } var messages: [Message] = [] for msgDict in messagesRaw { guard let msgId = msgDict.string(for: "id"), let senderId = msgDict.string(for: "sender_id") else { continue } let senderDeviceId = msgDict.string(for: "sender_device_id") ?? Constants.selfDeviceId let isDeleted = msgDict["deleted_at"] != nil && !(msgDict["deleted_at"] is NSNull) if isDeleted { let createdAt = msgDict.string(for: "created_at").flatMap { ISO8601DateFormatter().date(from: $0) } ?? Date() messages.append(Message( id: msgId, conversationId: convId, senderId: senderId, senderUsername: msgDict.string(for: "sender_username") ?? "", createdAt: createdAt, text: nil, isDeleted: true, readBy: [] )) continue } // Try to decrypt if let plaintext = decryptDMRecipientData( senderData: msgDict, senderId: senderId, senderDeviceId: senderDeviceId ) { let text = String(data: plaintext, encoding: .utf8) var messageText = text var replyTo: String? var file: FileInfo? if let jsonObj = try? JSONSerialization.jsonObject(with: plaintext) as? [String: Any] { messageText = jsonObj["text"] as? String replyTo = jsonObj["reply_to"] as? String if let fileDict = jsonObj["file"] as? [String: Any] { file = FileInfo( fileId: fileDict["file_id"] as? String ?? "", aesKey: fileDict["aes_key"] as? String ?? "", iv: fileDict["iv"] as? String ?? "", filename: fileDict["filename"] as? String ?? "", size: fileDict["size"] as? Int ?? 0, mimeType: fileDict["mime_type"] as? String ?? "" ) } } if messageText == nil && file == nil { continue } // Control message let createdAt = msgDict.string(for: "created_at").flatMap { ISO8601DateFormatter().date(from: $0) } ?? Date() messages.append(Message( id: msgId, conversationId: convId, senderId: senderId, senderUsername: msgDict.string(for: "sender_username") ?? "", createdAt: createdAt, text: messageText, replyTo: replyTo, imageFileId: msgDict.string(for: "image_file_id"), file: file, isDeleted: false, readBy: [] )) } } return messages } func markRead(convId: String, messageIds: [String]) async { _ = await sendAndReceive(type: "mark_read", params: [ "conversation_id": convId, "message_ids": messageIds, ]) } func deleteMessage(messageId: String, convId: String) async -> Bool { let resp = await sendAndReceive(type: "delete_message", params: [ "message_id": messageId, "conversation_id": convId, ]) return resp.string(for: "status") == "ok" } // MARK: - Group Operations func addMember(convId: String, email: String) async -> (success: Bool, message: String) { let resp = await sendAndReceive(type: "add_member", params: [ "conversation_id": convId, "email": email, ]) let msg = resp.dict(for: "data")?.string(for: "message") ?? "" return (resp.string(for: "status") == "ok", msg) } func removeMember(convId: String, userId: String) async -> (success: Bool, message: String) { let resp = await sendAndReceive(type: "remove_member", params: [ "conversation_id": convId, "user_id": userId, ]) let msg = resp.dict(for: "data")?.string(for: "message") ?? "" return (resp.string(for: "status") == "ok", msg) } func leaveGroup(convId: String) async -> (success: Bool, message: String) { let resp = await sendAndReceive(type: "leave_group", params: [ "conversation_id": convId, ]) if resp.string(for: "status") == "ok" { // Clean up local sender keys senderKeyStates.removeValue(forKey: convId) KeyStorage.deleteSenderKeyState(email: email, convId: convId) KeyStorage.deleteRecvSenderKeys(email: email, convId: convId) return (true, "Left group") } let msg = resp.dict(for: "data")?.string(for: "message") ?? "Failed" return (false, msg) } func renameConversation(convId: String, name: String) async -> (success: Bool, message: String) { let resp = await sendAndReceive(type: "rename_conversation", params: [ "conversation_id": convId, "name": name, ]) let msg = resp.dict(for: "data")?.string(for: "message") ?? "" return (resp.string(for: "status") == "ok", msg) } func deleteConversation(convId: String) async -> (success: Bool, message: String) { let resp = await sendAndReceive(type: "delete_conversation", params: [ "conversation_id": convId, ]) if resp.string(for: "status") == "ok" { senderKeyStates.removeValue(forKey: convId) KeyStorage.deleteSenderKeyState(email: email, convId: convId) KeyStorage.deleteRecvSenderKeys(email: email, convId: convId) return (true, "Deleted") } let msg = resp.dict(for: "data")?.string(for: "message") ?? "Failed" return (false, msg) } // MARK: - Invitations func acceptInvitation(convId: String) async -> (success: Bool, message: String) { let resp = await sendAndReceive(type: "accept_invitation", params: [ "conversation_id": convId, ]) let msg = resp.dict(for: "data")?.string(for: "message") ?? "" return (resp.string(for: "status") == "ok", msg) } func declineInvitation(convId: String) async -> (success: Bool, message: String) { let resp = await sendAndReceive(type: "decline_invitation", params: [ "conversation_id": convId, ]) let msg = resp.dict(for: "data")?.string(for: "message") ?? "" return (resp.string(for: "status") == "ok", msg) } func listInvitations() async -> [Invitation] { let resp = await sendAndReceive(type: "list_invitations") guard resp.string(for: "status") == "ok", let data = resp.dict(for: "data"), let invList = data["invitations"] as? [[String: Any]] else { return [] } return invList.compactMap { dict -> Invitation? in guard let convId = dict.string(for: "conversation_id") else { return nil } return Invitation( id: dict.string(for: "id") ?? convId, conversationId: convId, conversationName: dict.string(for: "conversation_name") ?? "Group", invitedBy: dict.string(for: "invited_by") ?? "", invitedByUsername: dict.string(for: "invited_by_username") ?? "" ) } } // MARK: - Profile func getProfile(userId: String? = nil) async -> UserProfile? { var params: [String: Any] = [:] if let userId = userId { params["user_id"] = userId } let resp = await sendAndReceive(type: "get_profile", params: params) guard resp.string(for: "status") == "ok", let data = resp.dict(for: "data") else { return nil } return UserProfile( userId: data.string(for: "user_id") ?? userId ?? self.userId ?? "", username: data.string(for: "username"), email: data.string(for: "email"), phone: data.string(for: "phone"), phoneVisible: data.bool(for: "phone_visible") ?? false, location: data.string(for: "location"), locationVisible: data.bool(for: "location_visible") ?? false, avatarFile: data.string(for: "avatar_file") ) } func updateProfile(phone: String? = nil, phoneVisible: Bool? = nil, location: String? = nil, locationVisible: Bool? = nil) async -> Bool { var params: [String: Any] = [:] if let phone = phone { params["phone"] = phone } if let phoneVisible = phoneVisible { params["phone_visible"] = phoneVisible } if let location = location { params["location"] = location } if let locationVisible = locationVisible { params["location_visible"] = locationVisible } let resp = await sendAndReceive(type: "update_profile", params: params) return resp.string(for: "status") == "ok" } func updateAvatar(imageData: Data) async -> Bool { let resp = await sendAndReceive(type: "update_avatar", params: [ "avatar_data": ProtocolHandler.encodeBinary(imageData), ]) return resp.string(for: "status") == "ok" } func getAvatar(userId: String) async -> Data? { let resp = await sendAndReceive(type: "get_avatar", params: ["user_id": userId]) guard resp.string(for: "status") == "ok", let data = resp.dict(for: "data"), let avatarB64 = data.string(for: "avatar_data"), let avatarData = try? ProtocolHandler.decodeBinary(avatarB64) else { return nil } return avatarData } // MARK: - Group Avatar func updateGroupAvatar(convId: String, imageData: Data) async -> Bool { let resp = await sendAndReceive(type: "update_group_avatar", params: [ "conversation_id": convId, "avatar_data": ProtocolHandler.encodeBinary(imageData), ]) return resp.string(for: "status") == "ok" } func getGroupAvatar(convId: String) async -> Data? { let resp = await sendAndReceive(type: "get_group_avatar", params: ["conversation_id": convId]) guard resp.string(for: "status") == "ok", let data = resp.dict(for: "data"), let avatarB64 = data.string(for: "avatar_data"), let avatarData = try? ProtocolHandler.decodeBinary(avatarB64) else { return nil } return avatarData } // MARK: - File Sharing func sendFile(convId: String, fileData: Data, filename: String, mimeType: String, members: [ConversationMember], replyTo: String? = nil) async -> (success: Bool, message: String) { // Encrypt file with AES-GCM guard let (aesKey, nonce, ct, tag) = try? CryptoUtils.aesEncrypt(fileData) else { return (false, "File encryption failed") } let encryptedData = ct + tag let fileType = mimeType.hasPrefix("image/") ? "image" : "file" // Start upload let startResp = await sendAndReceive(type: "upload_image_start", params: [ "conversation_id": convId, "file_size": encryptedData.count, "file_type": fileType, ]) guard startResp.string(for: "status") == "ok", let startData = startResp.dict(for: "data"), let fileId = startData.string(for: "file_id") else { let msg = startResp.dict(for: "data")?.string(for: "message") ?? "Upload start failed" return (false, msg) } // Upload chunks var offset = 0 while offset < encryptedData.count { let end = min(offset + Constants.imageChunkSize, encryptedData.count) let chunk = encryptedData[offset.. Data? { var allData = Data() var offset = 0 while true { let resp = await sendAndReceive(type: "download_image", params: [ "file_id": fileId, "offset": offset, ]) guard resp.string(for: "status") == "ok", let data = resp.dict(for: "data"), let chunkB64 = data.string(for: "chunk_data"), let chunk = try? ProtocolHandler.decodeBinary(chunkB64) else { break } if chunk.isEmpty { break } allData.append(chunk) offset += chunk.count if data.bool(for: "complete") == true { break } } guard !allData.isEmpty else { return nil } // Decrypt: allData = ciphertext + tag(16) guard allData.count >= 16 else { return nil } let ct = allData.prefix(allData.count - 16) let tag = allData.suffix(16) return try? CryptoUtils.aesDecrypt(key: aesKey, nonce: iv, ciphertext: Data(ct), tag: Data(tag)) } // MARK: - Devices func listDevices() async -> [[String: Any]] { let resp = await sendAndReceive(type: "list_devices") guard resp.string(for: "status") == "ok", let data = resp.dict(for: "data") else { return [] } return data["devices"] as? [[String: Any]] ?? [] } func removeDevice(deviceIdToRemove: String) async -> Bool { let resp = await sendAndReceive(type: "remove_device", params: [ "device_id": deviceIdToRemove, ]) return resp.string(for: "status") == "ok" } // MARK: - Session Reset func resetSession(peerUserId: String, peerDeviceId: String? = nil) async { if let peerDeviceId = peerDeviceId { let sessionKey = "\(peerUserId):\(peerDeviceId)" sessions.removeValue(forKey: sessionKey) KeyStorage.deleteSession(email: email, peerUserId: peerUserId, peerDeviceId: peerDeviceId) } else { // Delete all sessions for this user for key in sessions.keys where key.hasPrefix(peerUserId) { sessions.removeValue(forKey: key) } KeyStorage.deleteSession(email: email, peerUserId: peerUserId) } _ = await sendAndReceive(type: "session_reset", params: [ "peer_user_id": peerUserId, "peer_device_id": peerDeviceId ?? "", ]) } func handleSessionResetNotification(fromUserId: String, fromDeviceId: String?) { if let deviceId = fromDeviceId { let sessionKey = "\(fromUserId):\(deviceId)" sessions.removeValue(forKey: sessionKey) KeyStorage.deleteSession(email: email, peerUserId: fromUserId, peerDeviceId: deviceId) } else { for key in sessions.keys where key.hasPrefix(fromUserId) { sessions.removeValue(forKey: key) } KeyStorage.deleteSession(email: email, peerUserId: fromUserId) } } // MARK: - Search func searchMessages(convId: String, query: String) -> [[String: Any]] { MessageCache.search(email: email, convId: convId, query: query, cacheKey: cacheKey) } }