another.im-ios/ConversationsClassic/AppData/Model/OMEMO.swift

import Foundation
import GRDB
import Martin

// MARK: - Session
struct OMEMOSession: DBStorable {
    static let databaseTableName = "omemo_sessions"

    let account: String
    let name: String
    let deviceId: Int
    let key: String

    var id: String {
        "\(account)_\(name)_\(deviceId)"
    }
}

extension OMEMOSession {
    static func keyFor(account: String, name: String, deviceId: Int32) -> String? {
        do {
            return try Database.shared.dbQueue.read { db in
                try OMEMOSession
                    .filter(Column("account") == account)
                    .filter(Column("name") == name)
                    .filter(Column("deviceId") == deviceId)
                    .fetchOne(db)
            }?.key
        } catch {
            return nil
        }
    }

    static func devicesIdsFor(account: String, name: String) -> [Int32] {
        do {
            return try Database.shared.dbQueue.read { db in
                try OMEMOSession
                    .filter(Column("account") == account)
                    .filter(Column("name") == name)
                    .fetchAll(db)
                    .map(\.deviceId)
            }.map { Int32($0) }
        } catch {
            return []
        }
    }

    static func trustedDevicesIdsFor(account: String, name: String) -> [Int32] {
        do {
            let sql =
                """
                SELECT s.device_id
                FROM omemo_sessions s 
                LEFT JOIN omemo_identities i 
                ON s.account = i.account 
                AND s.name = i.name 
                AND s.device_id = i.device_id 
                WHERE s.account = :account 
                AND s.name = :name 
                AND ((i.status >= 0 AND i.status % 2 = 0) OR i.status IS NULL)
                """

            let arguments: StatementArguments = ["account": account, "name": name]

            return try Database.shared.dbQueue.read { db in
                try Int32.fetchAll(db, sql: sql, arguments: arguments)
            }
        } catch {
            return []
        }
    }

    static func wipe(account: String) {
        do {
            _ = try Database.shared.dbQueue.write { db in
                try OMEMOSession
                    .filter(Column("account") == account)
                    .deleteAll(db)
            }
        } catch {
            logIt(.error, "Failed to wipe OMEMO session: \(error)")
        }
    }
}

// MARK: - Identity
struct OMEMOIdentity: DBStorable {
    static let databaseTableName = "omemo_identities"

    let account: String
    let name: String
    let deviceId: Int
    let fingerprint: String
    let key: Data
    let own: Bool
    let status: Int

    var id: String {
        "\(account)_\(name)_\(deviceId)"
    }
}

extension OMEMOIdentity {
    static func wipe(account: String) {
        do {
            _ = try Database.shared.dbQueue.write { db in
                try OMEMOIdentity
                    .filter(Column("account") == account)
                    .deleteAll(db)
            }
        } catch {
            logIt(.error, "Failed to wipe OMEMO identity: \(error)")
        }
    }

    static func getFor(account: String, name: String, deviceId: Int32) -> OMEMOIdentity? {
        do {
            return try Database.shared.dbQueue.read { db in
                try OMEMOIdentity
                    .filter(Column("account") == account)
                    .filter(Column("name") == name)
                    .filter(Column("deviceId") == deviceId)
                    .fetchOne(db)
            }
        } catch {
            return nil
        }
    }

    static func existsFor(account: String, name: String, fingerprint: String) -> Bool {
        do {
            return try Database.shared.dbQueue.read { db in
                try OMEMOIdentity
                    .filter(Column("account") == account)
                    .filter(Column("name") == name)
                    .filter(Column("fingerprint") == fingerprint)
                    .fetchOne(db) != nil
            }
        } catch {
            return false
        }
    }

    func updateStatus(_ status: Int) -> Bool {
        do {
            _ = try Database.shared.dbQueue.write { db in
                try OMEMOIdentity
                    .filter(Column("account") == account)
                    .filter(Column("name") == name)
                    .filter(Column("deviceId") == deviceId)
                    .updateAll(db, Column("status").set(to: status))
            }
            return true
        } catch {
            logIt(.error, "Failed to update OMEMO identity status: \(error)")
            return false
        }
    }

    static func getAllFor(account: String, name: String) -> [OMEMOIdentity] {
        do {
            return try Database.shared.dbQueue.read { db in
                try OMEMOIdentity
                    .filter(Column("account") == account)
                    .filter(Column("name") == name)
                    .fetchAll(db)
            }
        } catch {
            return []
        }
    }
}

// MARK: - PreKey
struct OMEMOPreKey: DBStorable {
    static let databaseTableName = "omemo_pre_keys"

    let account: String
    let id: Int
    let key: Data
    let markForDeletion: Bool
}

extension OMEMOPreKey {
    static func wipe(account: String) {
        do {
            _ = try Database.shared.dbQueue.write { db in
                try OMEMOPreKey
                    .filter(Column("account") == account)
                    .deleteAll(db)
            }
        } catch {
            logIt(.error, "Failed to wipe OMEMO pre key: \(error)")
        }
    }

    static func currentIdFor(account: String) -> Int {
        do {
            return try Database.shared.dbQueue.read { db in
                try OMEMOPreKey
                    .filter(Column("account") == account)
                    .order(Column("id").desc)
                    .fetchOne(db)
                    .map(\.id)
            } ?? 0
        } catch {
            return 0
        }
    }

    static func keyFor(account: String, id: UInt32) -> Data? {
        do {
            return try Database.shared.dbQueue.read { db in
                try OMEMOPreKey
                    .filter(Column("account") == account)
                    .filter(Column("id") == id)
                    .fetchOne(db)
            }?.key
        } catch {
            return nil
        }
    }

    static func contains(account: String, id: UInt32) -> Bool {
        do {
            return try Database.shared.dbQueue.read { db in
                try OMEMOPreKey
                    .filter(Column("account") == account)
                    .filter(Column("id") == id)
                    .fetchOne(db) != nil
            }
        } catch {
            return false
        }
    }

    static func markForDeletion(account: String, id: UInt32) -> Bool {
        do {
            _ = try Database.shared.dbQueue.write { db in
                try OMEMOPreKey
                    .filter(Column("account") == account)
                    .filter(Column("id") == id)
                    .updateAll(db, Column("markForDeletion").set(to: true))
            }
            return true
        } catch {
            logIt(.error, "Failed to mark OMEMO pre key for deletion: \(error)")
            return false
        }
    }

    static func deleteMarked(account: String) -> Bool {
        do {
            _ = try Database.shared.dbQueue.write { db in
                try OMEMOPreKey
                    .filter(Column("account") == account)
                    .filter(Column("markForDeletion") == true)
                    .deleteAll(db)
            }
            return true
        } catch {
            logIt(.error, "Failed to delete marked OMEMO pre keys: \(error)")
            return false
        }
    }
}

// MARK: - SignedPreKey
struct OMEMOSignedPreKey: DBStorable {
    static let databaseTableName = "omemo_signed_pre_keys"

    let account: String
    let id: Int
    let key: Data
}

extension OMEMOSignedPreKey {
    static func wipe(account: String) {
        do {
            _ = try Database.shared.dbQueue.write { db in
                try OMEMOSignedPreKey
                    .filter(Column("account") == account)
                    .deleteAll(db)
            }
        } catch {
            logIt(.error, "Failed to wipe OMEMO signed pre key: \(error)")
        }
    }
}