IronStorage/apple/Sources/Watch/IronStorageWatchApp.swift

import Security
import SwiftUI
import WatchConnectivity

private enum SecureSnapshotStore {
    private static let service = "de.rfc1437.ironstorage.watch.snapshot"
    private static let account = "selected-totp"

    static func load() throws -> Data? {
        var result: CFTypeRef?
        let status = SecItemCopyMatching([
            kSecClass: kSecClassGenericPassword,
            kSecAttrService: service,
            kSecAttrAccount: account,
            kSecReturnData: true,
            kSecMatchLimit: kSecMatchLimitOne,
        ] as CFDictionary, &result)
        if status == errSecItemNotFound { return nil }
        guard status == errSecSuccess, let data = result as? Data else {
            throw SnapshotStoreError(status: status)
        }
        return data
    }

    static func replace(_ snapshot: Data) throws {
        let query = [
            kSecClass: kSecClassGenericPassword,
            kSecAttrService: service,
            kSecAttrAccount: account,
        ] as CFDictionary
        let status = SecItemUpdate(query, [kSecValueData: snapshot] as CFDictionary)
        if status == errSecItemNotFound {
            let added = SecItemAdd([
                kSecClass: kSecClassGenericPassword,
                kSecAttrService: service,
                kSecAttrAccount: account,
                kSecValueData: snapshot,
                kSecAttrAccessible: kSecAttrAccessibleWhenPasscodeSetThisDeviceOnly,
            ] as CFDictionary, nil)
            guard added == errSecSuccess else { throw SnapshotStoreError(status: added) }
        } else if status != errSecSuccess {
            throw SnapshotStoreError(status: status)
        }
    }

    static func delete() throws {
        let status = SecItemDelete([
            kSecClass: kSecClassGenericPassword,
            kSecAttrService: service,
            kSecAttrAccount: account,
        ] as CFDictionary)
        guard status == errSecSuccess || status == errSecItemNotFound else {
            throw SnapshotStoreError(status: status)
        }
    }
}

private struct SnapshotStoreError: Error {
    let status: OSStatus
}

@MainActor
private final class WatchSnapshotTransport: NSObject, ObservableObject, WCSessionDelegate {
    private static let snapshotKey = "de.rfc1437.ironstorage.watch.snapshot"
    private static let receiptKey = "de.rfc1437.ironstorage.watch.delivered"

    @Published private(set) var records: [WatchTotpRecord] = []
    private let core = watchCore()

    override init() {
        super.init()
        restoreSnapshot()
        guard WCSession.isSupported() else { return }
        let session = WCSession.default
        session.delegate = self
        session.activate()
    }

    func sceneBecameActive() {
        restoreSnapshot()
    }

    func sceneBecameInactive() {
        try? core.protectedDataUnavailable()
        records.removeAll(keepingCapacity: false)
    }

    private func restoreSnapshot() {
        do {
            if var snapshot = try SecureSnapshotStore.load() {
                defer { snapshot.resetBytes(in: 0..<snapshot.count) }
                let update = try core.applySnapshot(snapshot: snapshot)
                if update.persistence == .delete { try SecureSnapshotStore.delete() }
            } else {
                try core.noPersistedSnapshot()
            }
            refreshRecords()
        } catch {
            try? core.protectedDataUnavailable()
            records.removeAll(keepingCapacity: false)
        }
    }

    private func receive(_ applicationContext: [String: Any], session: WCSession) {
        guard var snapshot = applicationContext[Self.snapshotKey] as? Data else { return }
        defer { snapshot.resetBytes(in: 0..<snapshot.count) }
        do {
            let update = try core.applySnapshot(snapshot: snapshot)
            switch update.persistence {
            case .keep:
                break
            case .replace:
                try SecureSnapshotStore.replace(snapshot)
            case .delete:
                try SecureSnapshotStore.delete()
            }
            refreshRecords()
            guard !update.receipt.isEmpty else { return }
            let acknowledgement = [Self.receiptKey: update.receipt]
            if session.isReachable {
                session.sendMessage(acknowledgement, replyHandler: nil) { _ in
                    session.transferUserInfo(acknowledgement)
                }
            } else {
                session.transferUserInfo(acknowledgement)
            }
        } catch {
            try? core.noPersistedSnapshot()
            records.removeAll(keepingCapacity: false)
        }
    }

    private func refreshRecords() {
        records = (try? core.recordsAt(
            unixSeconds: UInt64(Date().timeIntervalSince1970)
        )) ?? []
    }

    nonisolated func session(
        _ session: WCSession,
        activationDidCompleteWith activationState: WCSessionActivationState,
        error: Error?
    ) {
        guard activationState == .activated, error == nil else { return }
        Task { @MainActor [weak self] in
            self?.receive(session.receivedApplicationContext, session: session)
        }
    }

    nonisolated func session(
        _ session: WCSession,
        didReceiveApplicationContext applicationContext: [String: Any]
    ) {
        Task { @MainActor [weak self] in
            self?.receive(applicationContext, session: session)
        }
    }
}

@main
struct IronStorageWatchApp: App {
    @Environment(\.scenePhase) private var scenePhase
    @StateObject private var transport = WatchSnapshotTransport()

    var body: some Scene {
        WindowGroup {
            ContentUnavailableView("No TOTP Codes", systemImage: "timer")
                .onChange(of: scenePhase) { _, phase in
                    if phase == .active {
                        transport.sceneBecameActive()
                    } else {
                        transport.sceneBecameInactive()
                    }
                }
        }
    }
}