another.im-ios/ConversationsClassic/AppData/Client/Client+MartinOMEMO.swift

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 deleting session: \(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 deleting all sessions: \(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 wiping sessions: \(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 {
        let id = OMEMOPreKey.currentIdFor(account: credentials.bareJid)
        return UInt32(id)
    }

    func loadPreKey(withId: UInt32) -> Data? {
        OMEMOPreKey.keyFor(account: credentials.bareJid, id: withId)
    }

    func storePreKey(_ data: Data, withId: UInt32) -> Bool {
        do {
            _ = try Database.shared.dbQueue.write { db in
                try OMEMOPreKey(
                    account: credentials.bareJid,
                    id: Int(withId),
                    key: data,
                    markForDeletion: false
                )
                .insert(db)
            }
            return true
        } catch {
            logIt(.error, "Error pre key store: \(error.localizedDescription)")
            return false
        }
    }

    func containsPreKey(withId: UInt32) -> Bool {
        OMEMOPreKey.contains(account: credentials.bareJid, id: withId)
    }

    func deletePreKey(withId: UInt32) -> Bool {
        OMEMOPreKey.markForDeletion(account: credentials.bareJid, id: withId)
    }

    func flushDeletedPreKeys() -> Bool {
        OMEMOPreKey.deleteMarked(account: credentials.bareJid)
    }

    func preKeysWipe() {
        OMEMOPreKey.wipe(account: credentials.bareJid)
    }
}

// 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 signed pre keys counting: \(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
    }
}