import Foundation import GRDB import Martin import MartinOMEMO final class ClientMartinOMEMO { let credentials: Credentials private let queue = DispatchQueue(label: "SignalPreKeyRemovalQueue") private var preKeysMarkedForRemoval: [UInt32] = [] init(_ credentials: Credentials) { self.credentials = credentials print("ClientMartinOMEMO init") } deinit { print("ClientMartinOMEMO deinit") } var signal: (SignalStorage, SignalContext) { let signalStorage = SignalStorage(sessionStore: self, preKeyStore: self, signedPreKeyStore: self, identityKeyStore: self, senderKeyStore: self) // swiftlint:disable:next force_unwrapping let signalContext = SignalContext(withStorage: signalStorage)! signalStorage.setup(withContext: signalContext) _ = regenerateKeys(wipe: false, context: signalContext) return (signalStorage, signalContext) } private func regenerateKeys(wipe: Bool = false, context: SignalContext) -> Bool { if wipe { OMEMOSession.wipe(account: credentials.bareJid) OMEMOPreKey.wipe(account: credentials.bareJid) OMEMOSignedPreKey.wipe(account: credentials.bareJid) OMEMOIdentity.wipe(account: credentials.bareJid) Settings.getFor(credentials.bareJid)?.wipeOmemoRegId() } let hasKeyPair = keyPair() != nil if wipe || localRegistrationId() == 0 || !hasKeyPair { let regId = context.generateRegistrationId() let address = SignalAddress(name: credentials.bareJid, deviceId: Int32(regId)) var settings = Settings.getFor(credentials.bareJid) settings?.omemoRegId = Int(regId) settings?.save() guard let keyPair = SignalIdentityKeyPair.generateKeyPair(context: context), let publicKey = keyPair.publicKey else { return false } let fingerprint = publicKey.map { byte -> String in String(format: "%02x", byte) }.joined() return save(address: address, fingerprint: fingerprint, own: true, data: keyPair.serialized()) } return true } private func save(address: SignalAddress, fingerprint: String, own: Bool, data: Data) -> Bool { guard !OMEMOIdentity.existsFor(account: credentials.bareJid, name: address.name, fingerprint: fingerprint) else { return false } do { _ = try Database.shared.dbQueue.write { db in try OMEMOIdentity( account: credentials.bareJid, name: address.name, deviceId: Int(address.deviceId), fingerprint: fingerprint, key: data, own: own, status: MartinOMEMO.IdentityStatus.trustedActive.rawValue ) .insert(db) } return true } catch { logIt(.error, "Error storing identity key: \(error.localizedDescription)") return false } } } // MARK: - Session extension ClientMartinOMEMO: SignalSessionStoreProtocol { func sessionRecord(forAddress address: MartinOMEMO.SignalAddress) -> Data? { if let key = OMEMOSession.keyFor(account: credentials.bareJid, name: address.name, deviceId: address.deviceId) { return Data(base64Encoded: key) } else { return nil } } func allDevices(for name: String, activeAndTrusted: Bool) -> [Int32] { activeAndTrusted ? OMEMOSession.trustedDevicesIdsFor(account: credentials.bareJid, name: name) : OMEMOSession.devicesIdsFor(account: credentials.bareJid, name: name) } func storeSessionRecord(_ data: Data, forAddress: MartinOMEMO.SignalAddress) -> Bool { do { try Database.shared.dbQueue.write { db in try OMEMOSession( account: credentials.bareJid, name: forAddress.name, deviceId: Int(forAddress.deviceId), key: data.base64EncodedString() ) .insert(db) } return true } catch { logIt(.error, "Error storing session info: \(error.localizedDescription)") return false } } func containsSessionRecord(forAddress: MartinOMEMO.SignalAddress) -> Bool { OMEMOSession.keyFor(account: credentials.bareJid, name: forAddress.name, deviceId: forAddress.deviceId) != nil } func deleteSessionRecord(forAddress: MartinOMEMO.SignalAddress) -> Bool { do { _ = try Database.shared.dbQueue.write { db in try OMEMOSession .filter(Column("account") == credentials.bareJid) .filter(Column("name") == forAddress.name) .filter(Column("deviceId") == forAddress.deviceId) .deleteAll(db) } return true } catch { logIt(.error, "Error fetching chats: \(error.localizedDescription)") return false } } func deleteAllSessions(for name: String) -> Bool { do { _ = try Database.shared.dbQueue.write { db in try OMEMOSession .filter(Column("account") == credentials.bareJid) .filter(Column("name") == name) .deleteAll(db) } return true } catch { logIt(.error, "Error fetching chats: \(error.localizedDescription)") return false } } func sessionsWipe() { do { _ = try Database.shared.dbQueue.write { db in try OMEMOSession .filter(Column("account") == credentials.bareJid) .deleteAll(db) } } catch { logIt(.error, "Error fetching chats: \(error.localizedDescription)") } } } // MARK: - Identity extension ClientMartinOMEMO: SignalIdentityKeyStoreProtocol { func keyPair() -> (any MartinOMEMO.SignalIdentityKeyPairProtocol)? { let deviceId = localRegistrationId() guard deviceId != 0 else { return nil } do { let record = try Database.shared.dbQueue.read { db in try OMEMOIdentity .filter(Column("account") == credentials.bareJid) .filter(Column("name") == credentials.bareJid) .filter(Column("deviceId") == deviceId) .fetchOne(db) } guard let key = record?.key else { return nil } return SignalIdentityKeyPair(fromKeyPairData: key) } catch { return nil } } func localRegistrationId() -> UInt32 { if let settings = Settings.getFor(credentials.bareJid) { return UInt32(settings.omemoRegId) } else { return 0 } } func save(identity: MartinOMEMO.SignalAddress, key: (any MartinOMEMO.SignalIdentityKeyProtocol)?) -> Bool { guard let key = key as SignalIdentityKeyProtocol?, let publicKey = key.publicKey else { return false } let fingerprint = publicKey.map { byte -> String in String(format: "%02x", byte) }.joined() defer { _ = self.setStatus(.verifiedActive, forIdentity: identity) } return save(address: identity, fingerprint: fingerprint, own: true, data: key.serialized()) } func save(identity: MartinOMEMO.SignalAddress, publicKeyData: Data?) -> Bool { guard let publicKeyData = publicKeyData else { return false } let fingerprint = publicKeyData.map { byte -> String in String(format: "%02x", byte) }.joined() return save(address: identity, fingerprint: fingerprint, own: false, data: publicKeyData) } func isTrusted(identity _: MartinOMEMO.SignalAddress, key _: (any MartinOMEMO.SignalIdentityKeyProtocol)?) -> Bool { true } func isTrusted(identity _: MartinOMEMO.SignalAddress, publicKeyData _: Data?) -> Bool { true } func setStatus(_ status: MartinOMEMO.IdentityStatus, forIdentity: MartinOMEMO.SignalAddress) -> Bool { if let identity = OMEMOIdentity.getFor(account: credentials.bareJid, name: forIdentity.name, deviceId: forIdentity.deviceId) { return identity.updateStatus(status.rawValue) } else { return false } } func setStatus(active: Bool, forIdentity: MartinOMEMO.SignalAddress) -> Bool { if let identity = OMEMOIdentity.getFor(account: credentials.bareJid, name: forIdentity.name, deviceId: forIdentity.deviceId) { let status = IdentityStatus(rawValue: identity.status) ?? .undecidedActive return identity.updateStatus(active ? status.toActive().rawValue : status.toInactive().rawValue) } else { return false } } func identities(forName name: String) -> [MartinOMEMO.Identity] { OMEMOIdentity.getAllFor(account: credentials.bareJid, name: name) .compactMap { identity in guard let status = IdentityStatus(rawValue: identity.status) else { return nil } return MartinOMEMO.Identity( address: MartinOMEMO.SignalAddress(name: identity.name, deviceId: Int32(identity.deviceId)), status: status, fingerprint: identity.fingerprint, key: identity.key, own: identity.own ) } } func identityFingerprint(forAddress address: MartinOMEMO.SignalAddress) -> String? { OMEMOIdentity.getFor(account: credentials.bareJid, name: address.name, deviceId: address.deviceId)?.fingerprint } } // MARK: - PreKey extension ClientMartinOMEMO: SignalPreKeyStoreProtocol { func currentPreKeyId() -> UInt32 { do { let data = try Database.shared.dbQueue.read { db in try Row.fetchOne( db, sql: "SELECT max(id) FROM omemo_pre_keys WHERE account = :account", arguments: ["account": credentials.bareJid] ) } return data?["id"] ?? 0 } catch { logIt(.error, "Error fetching chats: \(error.localizedDescription)") return 0 } } func loadPreKey(withId: UInt32) -> Data? { do { let data = try Database.shared.dbQueue.read { db in try Row.fetchOne( db, sql: "SELECT key FROM omemo_pre_keys WHERE account = :account AND id = :id", arguments: ["account": credentials.bareJid, "id": withId] ) } return data?["key"] } catch { logIt(.error, "Error fetching chats: \(error.localizedDescription)") return nil } } func storePreKey(_ data: Data, withId: UInt32) -> Bool { do { try Database.shared.dbQueue.write { db in try db.execute( sql: "INSERT INTO omemo_pre_keys (account, id, key) VALUES (:account, :id, :key)", arguments: ["account": credentials.bareJid, "id": withId, "key": data] ) } return true } catch { logIt(.error, "Error fetching chats: \(error.localizedDescription)") return false } } func containsPreKey(withId: UInt32) -> Bool { do { let rec = try Database.shared.dbQueue.read { db in try Row.fetchOne( db, sql: "SELECT key FROM omemo_pre_keys WHERE account = :account AND id = :id", arguments: ["account": credentials.bareJid, "id": withId] ) } return rec != nil } catch { logIt(.error, "Error fetching chats: \(error.localizedDescription)") return false } } func deletePreKey(withId: UInt32) -> Bool { queue.async { print("queueing prekey with id \(withId) for removal..") self.preKeysMarkedForRemoval.append(withId) } return true } // TODO: Check logic of this function carefully!!! func flushDeletedPreKeys() -> Bool { false // !queue.sync { () -> [UInt32] in // defer { // preKeysMarkedForRemoval.removeAll() // } // print("removing queued prekeys: \(preKeysMarkedForRemoval)") // do { // Database.shared.dbQueue.write { db in // try db.execute( // sql: "DETLETE FROM omemo_pre_keys WHERE account = :account AND id IN (:ids)", // arguments: ["account": credentials.bareJid, "ids": preKeysMarkedForRemoval] // ) // } // } catch { // logIt(.error, "Error fetching chats: \(error.localizedDescription)") // return [0] // } // // // return preKeysMarkedForRemoval.filter { id in DBOMEMOStore.instance.deletePreKey(forAccount: context!.sessionObject.userBareJid!, withId: id) } // }.isEmpty // // // // do { // try Database.shared.dbQueue.write { db in // try db.execute( // sql: """ // DELETE FROM omemo_pre_keys // WHERE account = :account // AND id IN ( // SELECT id // FROM omemo_pre_keys // WHERE account = :account // AND id NOT IN ( // SELECT id // FROM omemo_pre_keys // WHERE account = :account // ORDER BY id DESC // LIMIT 100) // ) // """, // arguments: ["account": credentials.bareJid] // ) // } // return true // } catch { // logIt(.error, "Error fetching chats: \(error.localizedDescription)") // return false // } } func preKeysWipe() { do { try Database.shared.dbQueue.write { db in try db.execute( sql: "DELETE FROM omemo_pre_keys WHERE account = :account", arguments: ["account": credentials.bareJid] ) } } catch { logIt(.error, "Error fetching chats: \(error.localizedDescription)") } } } // MARK: - SignedPreKey extension ClientMartinOMEMO: SignalSignedPreKeyStoreProtocol { func countSignedPreKeys() -> Int { do { let data = try Database.shared.dbQueue.read { db in try Row.fetchOne( db, sql: "SELECT count(1) FROM omemo_signed_pre_keys WHERE account = :account", arguments: ["account": credentials.bareJid] ) } return data?["count(1)"] ?? 0 } catch { logIt(.error, "Error fetching chats: \(error.localizedDescription)") return 0 } } func loadSignedPreKey(withId: UInt32) -> Data? { do { let data = try Database.shared.dbQueue.read { db in try Row.fetchOne( db, sql: "SELECT key FROM omemo_signed_pre_keys WHERE account = :account AND id = :id", arguments: ["account": credentials.bareJid, "id": withId] ) } return data?["key"] } catch { logIt(.error, "Error fetching chats: \(error.localizedDescription)") return nil } } func storeSignedPreKey(_ data: Data, withId: UInt32) -> Bool { do { try Database.shared.dbQueue.write { db in try db.execute( sql: "INSERT INTO omemo_signed_pre_keys (account, id, key) VALUES (:account, :id, :key)", arguments: ["account": credentials.bareJid, "id": withId, "key": data] ) } return true } catch { logIt(.error, "Error fetching chats: \(error.localizedDescription)") return false } } func containsSignedPreKey(withId: UInt32) -> Bool { do { let rec = try Database.shared.dbQueue.read { db in try Row.fetchOne( db, sql: "SELECT key FROM omemo_signed_pre_keys WHERE account = :account AND id = :id", arguments: ["account": credentials.bareJid, "id": withId] ) } return rec != nil } catch { logIt(.error, "Error fetching chats: \(error.localizedDescription)") return false } } func deleteSignedPreKey(withId: UInt32) -> Bool { do { try Database.shared.dbQueue.write { db in try db.execute( sql: "DELETE FROM omemo_signed_pre_keys WHERE account = :account AND id = :id", arguments: ["account": credentials.bareJid, "id": withId] ) } return true } catch { logIt(.error, "Error fetching chats: \(error.localizedDescription)") return false } } func wipeSignedPreKeys() { do { try Database.shared.dbQueue.write { db in try db.execute( sql: "DELETE FROM omemo_signed_pre_keys WHERE account = :account", arguments: ["account": credentials.bareJid] ) } } catch { logIt(.error, "Error fetching chats: \(error.localizedDescription)") } } } // MARK: - SenderKey extension ClientMartinOMEMO: SignalSenderKeyStoreProtocol { func storeSenderKey(_: Data, address _: MartinOMEMO.SignalAddress?, groupId _: String?) -> Bool { false } func loadSenderKey(forAddress _: MartinOMEMO.SignalAddress?, groupId _: String?) -> Data? { nil } }