Closer/iphone/Closer/Crypto/CoupleEncryptionManager.swift

import Foundation
import CryptoKit
import Sodium

/// iOS couple-key material and wrapped-key container.
///
/// Unlike Android, which stores a Tink JSON keyset envelope, iOS stores a raw
/// 32-byte AES-256-GCM key. The wrapped blob is the only cross-platform artifact
/// the server ever sees; each platform unwraps to its own usable AES-256-GCM key.
public struct CoupleKeyMaterial: @unchecked Sendable {
    /// Raw 32-byte AES-256-GCM key.
    public let rawKey: SymmetricKey

    public init(rawKey: SymmetricKey) {
        self.rawKey = rawKey
    }

    public init(rawBytes: Data) {
        self.rawKey = SymmetricKey(data: rawBytes)
    }
}

public struct WrappedCoupleKey: @unchecked Sendable {
    /// AES-GCM ciphertext: nonce (12) || ciphertext || tag (16).
    public let ciphertext: Data
    /// 16-byte Argon2id salt.
    public let kdfSalt: Data
    /// KDF parameter tag, e.g. "argon2id;v=19;m=47104;t=3;p=1".
    public let kdfParams: String

    public init(ciphertext: Data, kdfSalt: Data, kdfParams: String) {
        self.ciphertext = ciphertext
        self.kdfSalt = kdfSalt
        self.kdfParams = kdfParams
    }
}

/// High-level couple-key orchestration.
///
/// Wrap/unwrap uses Argon2id via libsodium (swift-sodium) with:
/// - algorithm: Argon2id v1.3 (ARGON2ID13)
/// - salt length: 16 bytes
/// - output length: 32 bytes
/// - memory: 46 MiB = 47104 KiB
/// - iterations: 3
/// - parallelism: 1
/// - AAD: "closer_couple_key"
///
/// These parameters match Android `RecoveryKeyManager` exactly.
public enum CoupleEncryptionManager {
    public static let kdfParamsTag = "argon2id;v=19;m=47104;t=3;p=1"
    public static let coupleKeyAAD = "closer_couple_key"
    public static let invitePhraseAAD = "closer_invite_phrase"
    public static let saltBytes = 16
    public static let keyBytes = 32
    public static let argon2MemoryKiB = 46 * 1024
    public static let argon2Iterations = 3
    public static let argon2Parallelism = 1

    /// Generates a new random 32-byte AES-256-GCM key.
    public static func generateCoupleKey() throws -> CoupleKeyMaterial {
        var bytes = [UInt8](repeating: 0, count: keyBytes)
        let status = SecRandomCopyBytes(kSecRandomDefault, keyBytes, &bytes)
        guard status == errSecSuccess else {
            throw CoupleEncryptionError.keyGenerationFailed(status)
        }
        return CoupleKeyMaterial(rawBytes: Data(bytes))
    }

    /// Wraps a couple key with a recovery-phrase-derived KEK.
    public static func wrap(_ key: CoupleKeyMaterial, with phrase: String) throws -> WrappedCoupleKey {
        var salt = [UInt8](repeating: 0, count: saltBytes)
        let saltStatus = SecRandomCopyBytes(kSecRandomDefault, saltBytes, &salt)
        guard saltStatus == errSecSuccess else {
            throw CoupleEncryptionError.saltGenerationFailed(saltStatus)
        }
        let saltData = Data(salt)
        let kek = try deriveKEK(phrase: phrase, salt: saltData)
        let rawKeyData = key.rawKey.bytes
        let ct = try FieldEncryptor.encrypt(rawKeyData, key: kek, aad: coupleKeyAAD.data(using: .utf8))
        return WrappedCoupleKey(ciphertext: ct, kdfSalt: saltData, kdfParams: kdfParamsTag)
    }

    /// Unwraps a couple key with a recovery-phrase-derived KEK.
    public static func unwrap(_ wrapped: WrappedCoupleKey, with phrase: String) throws -> CoupleKeyMaterial {
        let kek = try deriveKEK(phrase: phrase, salt: wrapped.kdfSalt)
        let rawKeyData = try FieldEncryptor.decrypt(
            wrapped.ciphertext,
            key: kek,
            aad: coupleKeyAAD.data(using: .utf8)
        )
        return CoupleKeyMaterial(rawBytes: rawKeyData)
    }

    /// Encrypts the recovery phrase with an invite-code-derived KEK.
    /// Output format: base64(salt[16] || AES-GCM ciphertext).
    public static func encryptRecoveryPhrase(_ phrase: String, with inviteCode: String) throws -> String {
        var salt = [UInt8](repeating: 0, count: saltBytes)
        let saltStatus = SecRandomCopyBytes(kSecRandomDefault, saltBytes, &salt)
        guard saltStatus == errSecSuccess else {
            throw CoupleEncryptionError.saltGenerationFailed(saltStatus)
        }
        let saltData = Data(salt)
        let kek = try deriveKEK(phrase: inviteCode, salt: saltData)
        let ct = try FieldEncryptor.encrypt(
            Data(phrase.utf8),
            key: kek,
            aad: invitePhraseAAD.data(using: .utf8)
        )
        return (saltData + ct).base64EncodedString()
    }

    /// Decrypts the recovery phrase that was encrypted with an invite code.
    public static func decryptRecoveryPhrase(_ blob: String, with inviteCode: String) throws -> String {
        guard let payload = Data(base64Encoded: blob) else {
            throw CoupleEncryptionError.invalidBase64
        }
        guard payload.count > saltBytes else {
            throw CoupleEncryptionError.invalidPayload
        }
        let salt = payload.prefix(saltBytes)
        let ct = payload.dropFirst(saltBytes)
        let kek = try deriveKEK(phrase: inviteCode, salt: salt)
        let pt = try FieldEncryptor.decrypt(
            ct,
            key: kek,
            aad: invitePhraseAAD.data(using: .utf8)
        )
        guard let phrase = String(data: pt, encoding: .utf8) else {
            throw CoupleEncryptionError.invalidUTF8
        }
        return phrase
    }

    // MARK: - Internal

    /// Exposed for deterministic known-vector tests. Derives the KEK from a
    /// password + salt using the same Argon2id parameters as Android.
    public static func unwrapKEK(phrase: String, salt: Data) throws -> SymmetricKey {
        return try deriveKEK(phrase: phrase, salt: salt)
    }

    private static func deriveKEK(phrase: String, salt: Data) throws -> SymmetricKey {
        guard salt.count == saltBytes else {
            throw CoupleEncryptionError.invalidSaltLength
        }
        let sodium = Sodium()
        guard let bytes = sodium.pwHash.hash(
            outputLength: keyBytes,
            passwd: Array(phrase.utf8),
            salt: [UInt8](salt),
            opsLimit: argon2Iterations,
            memLimit: argon2MemoryKiB * 1024,
            alg: .Argon2ID13
        ) else {
            throw CoupleEncryptionError.keyDerivationFailed
        }
        return SymmetricKey(data: Data(bytes))
    }

    public enum CoupleEncryptionError: Error {
        case keyGenerationFailed(OSStatus)
        case saltGenerationFailed(OSStatus)
        case keyDerivationFailed
        case invalidSaltLength
        case invalidBase64
        case invalidPayload
        case invalidUTF8
    }
}