eric / piclead-client

import CryptoKit
import Foundation

/// Kotlin `ApiCleCryptoException`.
struct ApiCleCryptoException: Error {
    let message: String
    let cause: Error?

    init(_ message: String, cause: Error? = nil) {
        self.message = message
        self.cause = cause
    }
}

/// Decrypts the Projectiaon API key blob produced by the Rust server:
/// KDF Argon2id (crate `argon2` 0.5 defaults: m=19456 KiB, t=2, p=1, 32-byte key)
/// + XChaCha20-Poly1305, blob layout = 24-byte nonce || ciphertext || 16-byte tag.
/// Same on-disk format as the Android port (BouncyCastle); implemented here with a
/// pure-Swift Argon2id/Blake2b + HChaCha20 and CryptoKit `ChaChaPoly` for the AEAD.
enum ApiCleCrypto {
    private static let NONCE_LEN = 24
    private static let TAG_LEN = 16
    private static let KEY_LEN = 32
    private static let SALT_LEN = 16
    private static let ARGON2_M_KB = 19456
    private static let ARGON2_T = 2
    private static let ARGON2_P = 1

    static func decryptApiKey(password: String, saltB64: String, cipherB64: String) throws -> String {
        let salt = try decodeB64(saltB64, "sel")
        guard salt.count == SALT_LEN else {
            throw ApiCleCryptoException("sel invalide")
        }

        let blob = try decodeB64(cipherB64, "cle_chiffree")
        guard blob.count >= NONCE_LEN + TAG_LEN else {
            throw ApiCleCryptoException("blob trop court")
        }

        let key = Argon2id.hash(
            password: [UInt8](password.utf8),
            salt: salt,
            timeCost: ARGON2_T,
            memoryKiB: ARGON2_M_KB,
            parallelism: ARGON2_P,
            outputLength: KEY_LEN
        )
        let nonce = Array(blob[0..<NONCE_LEN])
        let ciphertext = Array(blob[NONCE_LEN...])

        let plaintext: Data
        do {
            plaintext = try XChaCha20Poly1305.open(key: key, nonce24: nonce, ciphertextAndTag: ciphertext)
        } catch {
            throw ApiCleCryptoException("échec déchiffrement clé API", cause: error)
        }

        return String(decoding: plaintext, as: UTF8.self)
    }

    private static func decodeB64(_ value: String, _ label: String) throws -> [UInt8] {
        guard let data = Data(base64Encoded: value) else {
            throw ApiCleCryptoException("\(label) base64 invalide")
        }
        return [UInt8](data)
    }
}

// MARK: - XChaCha20-Poly1305 (24-byte nonce) via HChaCha20 + CryptoKit ChaChaPoly

enum XChaCha20Poly1305 {
    struct AuthenticationFailure: Error {}

    static func open(key: [UInt8], nonce24: [UInt8], ciphertextAndTag: [UInt8]) throws -> Data {
        guard key.count == 32, nonce24.count == 24, ciphertextAndTag.count >= 16 else {
            throw AuthenticationFailure()
        }
        let subkey = hChaCha20(key: key, nonce16: Array(nonce24[0..<16]))
        // IETF ChaCha20-Poly1305 nonce = 4 zero bytes || last 8 bytes of the 24-byte nonce.
        var ietfNonce = [UInt8](repeating: 0, count: 4)
        ietfNonce.append(contentsOf: nonce24[16..<24])

        let ciphertext = Data(ciphertextAndTag.dropLast(16))
        let tag = Data(ciphertextAndTag.suffix(16))
        let box = try ChaChaPoly.SealedBox(
            nonce: ChaChaPoly.Nonce(data: Data(ietfNonce)),
            ciphertext: ciphertext,
            tag: tag
        )
        return try ChaChaPoly.open(box, using: SymmetricKey(data: Data(subkey)))
    }

    private static func hChaCha20(key: [UInt8], nonce16: [UInt8]) -> [UInt8] {
        var state = [UInt32](repeating: 0, count: 16)
        state[0] = 0x6170_7865
        state[1] = 0x3320_646e
        state[2] = 0x7962_2d32
        state[3] = 0x6b20_6574
        for i in 0..<8 { state[4 + i] = load32(key, 4 * i) }
        for i in 0..<4 { state[12 + i] = load32(nonce16, 4 * i) }

        for _ in 0..<10 {
            quarterRound(&state, 0, 4, 8, 12)
            quarterRound(&state, 1, 5, 9, 13)
            quarterRound(&state, 2, 6, 10, 14)
            quarterRound(&state, 3, 7, 11, 15)
            quarterRound(&state, 0, 5, 10, 15)
            quarterRound(&state, 1, 6, 11, 12)
            quarterRound(&state, 2, 7, 8, 13)
            quarterRound(&state, 3, 4, 9, 14)
        }

        var out = [UInt8]()
        out.reserveCapacity(32)
        for i in [0, 1, 2, 3, 12, 13, 14, 15] {
            var word = state[i].littleEndian
            withUnsafeBytes(of: &word) { out.append(contentsOf: $0) }
        }
        return out
    }

    private static func load32(_ bytes: [UInt8], _ offset: Int) -> UInt32 {
        UInt32(bytes[offset])
            | (UInt32(bytes[offset + 1]) << 8)
            | (UInt32(bytes[offset + 2]) << 16)
            | (UInt32(bytes[offset + 3]) << 24)
    }

    private static func quarterRound(_ s: inout [UInt32], _ a: Int, _ b: Int, _ c: Int, _ d: Int) {
        s[a] = s[a] &+ s[b]; s[d] = rotl(s[d] ^ s[a], 16)
        s[c] = s[c] &+ s[d]; s[b] = rotl(s[b] ^ s[c], 12)
        s[a] = s[a] &+ s[b]; s[d] = rotl(s[d] ^ s[a], 8)
        s[c] = s[c] &+ s[d]; s[b] = rotl(s[b] ^ s[c], 7)
    }

    private static func rotl(_ x: UInt32, _ n: UInt32) -> UInt32 {
        (x << n) | (x >> (32 - n))
    }
}

// MARK: - Blake2b (unkeyed, variable digest length up to 64 bytes)

enum Blake2b {
    private static let iv: [UInt64] = [
        0x6a09_e667_f3bc_c908, 0xbb67_ae85_84ca_a73b,
        0x3c6e_f372_fe94_f82b, 0xa54f_f53a_5f1d_36f1,
        0x510e_527f_ade6_82d1, 0x9b05_688c_2b3e_6c1f,
        0x1f83_d9ab_fb41_bd6b, 0x5be0_cd19_137e_2179,
    ]

    private static let sigma: [[Int]] = [
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
        [14, 10, 4, 8, 9, 15, 13, 6, 1, 12, 0, 2, 11, 7, 5, 3],
        [11, 8, 12, 0, 5, 2, 15, 13, 10, 14, 3, 6, 7, 1, 9, 4],
        [7, 9, 3, 1, 13, 12, 11, 14, 2, 6, 5, 10, 4, 0, 15, 8],
        [9, 0, 5, 7, 2, 4, 10, 15, 14, 1, 11, 12, 6, 8, 3, 13],
        [2, 12, 6, 10, 0, 11, 8, 3, 4, 13, 7, 5, 15, 14, 1, 9],
        [12, 5, 1, 15, 14, 13, 4, 10, 0, 7, 6, 3, 9, 2, 8, 11],
        [13, 11, 7, 14, 12, 1, 3, 9, 5, 0, 15, 4, 8, 6, 2, 10],
        [6, 15, 14, 9, 11, 3, 0, 8, 12, 2, 13, 7, 1, 4, 10, 5],
        [10, 2, 8, 4, 7, 6, 1, 5, 15, 11, 9, 14, 3, 12, 13, 0],
    ]

    static func hash(_ input: [UInt8], outLen: Int) -> [UInt8] {
        precondition(outLen >= 1 && outLen <= 64)
        var h = iv
        h[0] ^= 0x0101_0000 ^ UInt64(outLen)

        var t: UInt64 = 0
        var offset = 0
        let count = input.count
        // Process all blocks but the last (which may be partial or empty).
        while count - offset > 128 {
            t &+= 128
            compress(&h, block: input, blockOffset: offset, t: t, last: false)
            offset += 128
        }
        var lastBlock = [UInt8](repeating: 0, count: 128)
        let remaining = count - offset
        if remaining > 0 {
            lastBlock.replaceSubrange(0..<remaining, with: input[offset..<count])
        }
        t &+= UInt64(remaining)
        compress(&h, block: lastBlock, blockOffset: 0, t: t, last: true)

        var out = [UInt8]()
        out.reserveCapacity(outLen)
        outer: for word in h {
            var value = word.littleEndian
            withUnsafeBytes(of: &value) { bytes in
                for byte in bytes {
                    if out.count == outLen { return }
                    out.append(byte)
                }
            }
            if out.count == outLen { break outer }
        }
        return out
    }

    private static func compress(_ h: inout [UInt64], block: [UInt8], blockOffset: Int, t: UInt64, last: Bool) {
        var m = [UInt64](repeating: 0, count: 16)
        for i in 0..<16 {
            let o = blockOffset + 8 * i
            m[i] = UInt64(block[o])
                | (UInt64(block[o + 1]) << 8)
                | (UInt64(block[o + 2]) << 16)
                | (UInt64(block[o + 3]) << 24)
                | (UInt64(block[o + 4]) << 32)
                | (UInt64(block[o + 5]) << 40)
                | (UInt64(block[o + 6]) << 48)
                | (UInt64(block[o + 7]) << 56)
        }

        var v = [UInt64](repeating: 0, count: 16)
        for i in 0..<8 { v[i] = h[i] }
        for i in 0..<8 { v[8 + i] = iv[i] }
        v[12] ^= t
        // t high 64 bits always 0 for our input sizes.
        if last { v[14] = ~v[14] }

        for round in 0..<12 {
            let s = sigma[round % 10]
            g(&v, 0, 4, 8, 12, m[s[0]], m[s[1]])
            g(&v, 1, 5, 9, 13, m[s[2]], m[s[3]])
            g(&v, 2, 6, 10, 14, m[s[4]], m[s[5]])
            g(&v, 3, 7, 11, 15, m[s[6]], m[s[7]])
            g(&v, 0, 5, 10, 15, m[s[8]], m[s[9]])
            g(&v, 1, 6, 11, 12, m[s[10]], m[s[11]])
            g(&v, 2, 7, 8, 13, m[s[12]], m[s[13]])
            g(&v, 3, 4, 9, 14, m[s[14]], m[s[15]])
        }

        for i in 0..<8 {
            h[i] ^= v[i] ^ v[i + 8]
        }
    }

    private static func g(_ v: inout [UInt64], _ a: Int, _ b: Int, _ c: Int, _ d: Int, _ x: UInt64, _ y: UInt64) {
        v[a] = v[a] &+ v[b] &+ x
        v[d] = rotr(v[d] ^ v[a], 32)
        v[c] = v[c] &+ v[d]
        v[b] = rotr(v[b] ^ v[c], 24)
        v[a] = v[a] &+ v[b] &+ y
        v[d] = rotr(v[d] ^ v[a], 16)
        v[c] = v[c] &+ v[d]
        v[b] = rotr(v[b] ^ v[c], 63)
    }

    private static func rotr(_ x: UInt64, _ n: UInt64) -> UInt64 {
        (x >> n) | (x << (64 - n))
    }
}

// MARK: - Argon2id (RFC 9106, version 0x13), single-lane (p = 1)

enum Argon2id {
    private static let blockWords = 128 // 1024 bytes
    private static let syncPoints = 4
    private static let addressesInBlock = 128

    /// Argon2id hash; `parallelism` must be 1 (matches the server parameters).
    static func hash(
        password: [UInt8],
        salt: [UInt8],
        timeCost: Int,
        memoryKiB: Int,
        parallelism: Int,
        outputLength: Int
    ) -> [UInt8] {
        precondition(parallelism == 1, "single-lane implementation")
        precondition(timeCost >= 1 && memoryKiB >= 8)

        let mPrime = (memoryKiB / (4 * parallelism)) * 4 * parallelism
        let laneLength = mPrime / parallelism
        let segmentLength = laneLength / syncPoints

        // H0
        var seed = [UInt8]()
        seed.append(contentsOf: le32(parallelism))
        seed.append(contentsOf: le32(outputLength))
        seed.append(contentsOf: le32(memoryKiB))
        seed.append(contentsOf: le32(timeCost))
        seed.append(contentsOf: le32(0x13)) // version
        seed.append(contentsOf: le32(2)) // type Argon2id
        seed.append(contentsOf: le32(password.count))
        seed.append(contentsOf: password)
        seed.append(contentsOf: le32(salt.count))
        seed.append(contentsOf: salt)
        seed.append(contentsOf: le32(0)) // secret
        seed.append(contentsOf: le32(0)) // associated data
        let h0 = Blake2b.hash(seed, outLen: 64)

        var memory = [UInt64](repeating: 0, count: mPrime * blockWords)

        // First two blocks of the lane.
        storeBlock(hPrime(h0 + le32(0) + le32(0), outLen: 1024), into: &memory, index: 0)
        storeBlock(hPrime(h0 + le32(1) + le32(0), outLen: 1024), into: &memory, index: 1)

        for pass in 0..<timeCost {
            for slice in 0..<syncPoints {
                fillSegment(
                    memory: &memory,
                    pass: pass,
                    slice: slice,
                    timeCost: timeCost,
                    mPrime: mPrime,
                    laneLength: laneLength,
                    segmentLength: segmentLength
                )
            }
        }

        let finalBlock = loadBlockBytes(memory, index: laneLength - 1)
        return hPrime(finalBlock, outLen: outputLength)
    }

    // H': variable-length hash (RFC 9106 §3.3).
    private static func hPrime(_ input: [UInt8], outLen: Int) -> [UInt8] {
        let prefixed = le32(outLen) + input
        if outLen <= 64 {
            return Blake2b.hash(prefixed, outLen: outLen)
        }
        let r = (outLen + 31) / 32 - 2
        var out = [UInt8]()
        out.reserveCapacity(outLen)
        var v = Blake2b.hash(prefixed, outLen: 64)
        out.append(contentsOf: v[0..<32])
        for _ in 1..<r {
            v = Blake2b.hash(v, outLen: 64)
            out.append(contentsOf: v[0..<32])
        }
        v = Blake2b.hash(v, outLen: outLen - 32 * r)
        out.append(contentsOf: v)
        return out
    }

    private static func fillSegment(
        memory: inout [UInt64],
        pass: Int,
        slice: Int,
        timeCost: Int,
        mPrime: Int,
        laneLength: Int,
        segmentLength: Int
    ) {
        // Argon2id: data-independent addressing for the first two slices of pass 0.
        let dataIndependent = pass == 0 && slice < 2

        var addressBlock = [UInt64](repeating: 0, count: blockWords)
        var inputBlock = [UInt64](repeating: 0, count: blockWords)
        let zeroBlock = [UInt64](repeating: 0, count: blockWords)
        if dataIndependent {
            inputBlock[0] = UInt64(pass)
            inputBlock[1] = 0 // lane
            inputBlock[2] = UInt64(slice)
            inputBlock[3] = UInt64(mPrime)
            inputBlock[4] = UInt64(timeCost)
            inputBlock[5] = 2 // Argon2id
        }

        var startingIndex = 0
        if pass == 0 && slice == 0 {
            startingIndex = 2
            if dataIndependent {
                nextAddresses(address: &addressBlock, input: &inputBlock, zero: zeroBlock)
            }
        }

        var currOffset = slice * segmentLength + startingIndex
        var prevOffset = currOffset == 0 ? laneLength - 1 : currOffset - 1

        for i in startingIndex..<segmentLength {
            if currOffset % laneLength == 0 {
                prevOffset = currOffset + laneLength - 1
            }

            let pseudoRand: UInt64
            if dataIndependent {
                if i % addressesInBlock == 0 {
                    nextAddresses(address: &addressBlock, input: &inputBlock, zero: zeroBlock)
                }
                pseudoRand = addressBlock[i % addressesInBlock]
            } else {
                pseudoRand = memory[prevOffset * blockWords]
            }
            // Single lane: ref lane is always our lane.
            let j1 = pseudoRand & 0xFFFF_FFFF
            let refIndex = indexAlpha(
                pass: pass,
                slice: slice,
                index: i,
                j1: j1,
                laneLength: laneLength,
                segmentLength: segmentLength
            )

            let prevBlock = readBlock(memory, index: prevOffset)
            let refBlock = readBlock(memory, index: refIndex)
            var current = readBlock(memory, index: currOffset)
            fillBlock(prev: prevBlock, ref: refBlock, next: &current, withXor: pass > 0)
            writeBlock(current, into: &memory, index: currOffset)

            currOffset += 1
            prevOffset = currOffset - 1
        }
    }

    private static func indexAlpha(
        pass: Int,
        slice: Int,
        index: Int,
        j1: UInt64,
        laneLength: Int,
        segmentLength: Int
    ) -> Int {
        let refAreaSize: UInt64
        if pass == 0 {
            if slice == 0 {
                refAreaSize = UInt64(index - 1)
            } else {
                refAreaSize = UInt64(slice * segmentLength + index - 1)
            }
        } else {
            refAreaSize = UInt64(laneLength - segmentLength + index - 1)
        }

        var relativePosition = j1
        relativePosition = (relativePosition &* relativePosition) >> 32
        relativePosition = refAreaSize - 1 - ((refAreaSize &* relativePosition) >> 32)

        let startPosition: Int
        if pass == 0 || slice == syncPoints - 1 {
            startPosition = 0
        } else {
            startPosition = (slice + 1) * segmentLength
        }
        return (startPosition + Int(relativePosition)) % laneLength
    }

    private static func nextAddresses(address: inout [UInt64], input: inout [UInt64], zero: [UInt64]) {
        input[6] &+= 1
        fillBlock(prev: zero, ref: input, next: &address, withXor: false)
        let copy = address
        fillBlock(prev: zero, ref: copy, next: &address, withXor: false)
    }

    /// Argon2 compression G: next = P(prev ^ ref) ^ (prev ^ ref) [^ next when withXor].
    private static func fillBlock(prev: [UInt64], ref: [UInt64], next: inout [UInt64], withXor: Bool) {
        var blockR = [UInt64](repeating: 0, count: blockWords)
        for i in 0..<blockWords { blockR[i] = prev[i] ^ ref[i] }
        var blockTmp = blockR
        if withXor {
            for i in 0..<blockWords { blockTmp[i] ^= next[i] }
        }

        // Apply P row-wise: 8 rows of 16 consecutive u64.
        for i in 0..<8 {
            let o = 16 * i
            permutation(
                &blockR,
                o, o + 1, o + 2, o + 3, o + 4, o + 5, o + 6, o + 7,
                o + 8, o + 9, o + 10, o + 11, o + 12, o + 13, o + 14, o + 15
            )
        }
        // Apply P column-wise: 8 columns of 2-u64 pairs.
        for i in 0..<8 {
            let o = 2 * i
            permutation(
                &blockR,
                o, o + 1, o + 16, o + 17, o + 32, o + 33, o + 48, o + 49,
                o + 64, o + 65, o + 80, o + 81, o + 96, o + 97, o + 112, o + 113
            )
        }

        for i in 0..<blockWords { next[i] = blockTmp[i] ^ blockR[i] }
    }

    private static func permutation(
        _ v: inout [UInt64],
        _ v0: Int, _ v1: Int, _ v2: Int, _ v3: Int,
        _ v4: Int, _ v5: Int, _ v6: Int, _ v7: Int,
        _ v8: Int, _ v9: Int, _ v10: Int, _ v11: Int,
        _ v12: Int, _ v13: Int, _ v14: Int, _ v15: Int
    ) {
        blamkaG(&v, v0, v4, v8, v12)
        blamkaG(&v, v1, v5, v9, v13)
        blamkaG(&v, v2, v6, v10, v14)
        blamkaG(&v, v3, v7, v11, v15)
        blamkaG(&v, v0, v5, v10, v15)
        blamkaG(&v, v1, v6, v11, v12)
        blamkaG(&v, v2, v7, v8, v13)
        blamkaG(&v, v3, v4, v9, v14)
    }

    @inline(__always)
    private static func blamkaG(_ v: inout [UInt64], _ a: Int, _ b: Int, _ c: Int, _ d: Int) {
        v[a] = fBlaMka(v[a], v[b]); v[d] = rotr(v[d] ^ v[a], 32)
        v[c] = fBlaMka(v[c], v[d]); v[b] = rotr(v[b] ^ v[c], 24)
        v[a] = fBlaMka(v[a], v[b]); v[d] = rotr(v[d] ^ v[a], 16)
        v[c] = fBlaMka(v[c], v[d]); v[b] = rotr(v[b] ^ v[c], 63)
    }

    @inline(__always)
    private static func fBlaMka(_ x: UInt64, _ y: UInt64) -> UInt64 {
        let m = (x & 0xFFFF_FFFF) &* (y & 0xFFFF_FFFF)
        return x &+ y &+ (m &<< 1)
    }

    @inline(__always)
    private static func rotr(_ x: UInt64, _ n: UInt64) -> UInt64 {
        (x >> n) | (x << (64 - n))
    }

    // ---- block/bytes helpers ----

    private static func readBlock(_ memory: [UInt64], index: Int) -> [UInt64] {
        Array(memory[(index * blockWords)..<((index + 1) * blockWords)])
    }

    private static func writeBlock(_ block: [UInt64], into memory: inout [UInt64], index: Int) {
        memory.replaceSubrange((index * blockWords)..<((index + 1) * blockWords), with: block)
    }

    private static func storeBlock(_ bytes: [UInt8], into memory: inout [UInt64], index: Int) {
        precondition(bytes.count == 1024)
        for i in 0..<blockWords {
            let o = 8 * i
            memory[index * blockWords + i] = UInt64(bytes[o])
                | (UInt64(bytes[o + 1]) << 8)
                | (UInt64(bytes[o + 2]) << 16)
                | (UInt64(bytes[o + 3]) << 24)
                | (UInt64(bytes[o + 4]) << 32)
                | (UInt64(bytes[o + 5]) << 40)
                | (UInt64(bytes[o + 6]) << 48)
                | (UInt64(bytes[o + 7]) << 56)
        }
    }

    private static func loadBlockBytes(_ memory: [UInt64], index: Int) -> [UInt8] {
        var out = [UInt8]()
        out.reserveCapacity(1024)
        for i in 0..<blockWords {
            var word = memory[index * blockWords + i].littleEndian
            withUnsafeBytes(of: &word) { out.append(contentsOf: $0) }
        }
        return out
    }

    private static func le32(_ value: Int) -> [UInt8] {
        let v = UInt32(truncatingIfNeeded: value)
        return [
            UInt8(v & 0xFF),
            UInt8((v >> 8) & 0xFF),
            UInt8((v >> 16) & 0xFF),
            UInt8((v >> 24) & 0xFF),
        ]
    }
}