import Crypto
import Foundation
import SwiftASN1
import X509
import _CryptoExtras

/// Builds a PKCS#12 (RFC 7292) bundle — the device identity baked into the enrollment profile's
/// `com.apple.security.pkcs12` payload (docs/MACOS_VM_MDM.md §4). swift-certificates ships no PKCS#12
/// serializer, so this is hand-rolled DER:
///
///   • **cert bag** — the device leaf, plaintext (a `certBag` inside a `data` SafeContents);
///   • **shrouded key bag** — the RSA private key as an `EncryptedPrivateKeyInfo` using the RFC 7292
///     v1 scheme **`pbeWithSHAAnd3-KeyTripleDES-CBC`** (PKCS#12 KDF/SHA-1 → 3DES-EDE-CBC);
///   • **MAC** — a `MacData` over the AuthenticatedSafe, **HMAC-SHA-1** keyed by the same PKCS#12 v1
///     KDF (RFC 7292 App. B.2, `id = 3`).
///
/// ## Why the obsolete algorithms are mandatory here
///
/// This encoder emitted modern crypto (PBES2 / PBKDF2-HMAC-SHA256 / AES-256-CBC, SHA-256 `MacData`)
/// for eleven runs, and every enrollment failed with
/// `ConfigProfilePluginDomain -323 "The certificate could not be verified (authentication error)."`
///
/// macOS's configuration-profile certificate plugin imports identities through **`SecItemImport`**,
/// the legacy CDSA-backed API — `nm -u /usr/libexec/mdmclient` and `ManagedClient` list
/// `_SecItemImport` and *not* `SecPKCS12Import`. Measured on macOS 27.0 against the exact bytes from a
/// failing run:
///
/// | `MacData` | shrouded key bag        | `SecItemImport`                       | `SecPKCS12Import` |
/// |-----------|-------------------------|---------------------------------------|-------------------|
/// | SHA-256   | PBES2/AES-256 (old code)| `-25264` errSecPkcs12VerifyFailure    | ok                |
/// | SHA-256   | PBE-SHA1-3DES           | `-25264` errSecPkcs12VerifyFailure    | —                 |
/// | SHA-1     | PBES2/AES-256           | "Unknown format in import"            | ok                |
/// | **SHA-1** | **PBE-SHA1-3DES**       | **ok — 1 identity**                   | ok                |
/// | SHA-1     | *unencrypted* keyBag    | ok — 1 **certificate**, key dropped   | —                 |
///
/// So both must change together, and a plain (unencrypted) key bag is not an escape hatch: the
/// importer silently drops the key and yields a bare certificate, which would leave the MDM payload's
/// `IdentityCertificateUUID` resolving to no identity. The cert bag may stay plaintext.
///
/// Of the v1 PBE algorithms, 3DES is the only one a stock OpenSSL 3 can also read back — RC4 and RC2
/// need the legacy provider, which would make our own interop test unrunnable. Hence `TripleDES`.
///
/// Nothing is lost by using a broken cipher: the import password ships in cleartext in the same
/// `.mobileconfig` (`Password`), so this layer is a format requirement, not a security control.
///
/// The output is byte-verifiable on-host: `openssl pkcs12 -info -nodes` parses it and recovers the
/// key, and `SecItemImport`/`SecPKCS12Import` both return an identity (`PKCS12Tests` asserts all three
/// — the `SecItemImport` assertion is the gate whose absence let `-323` survive eleven runs).
public enum PKCS12 {
    // OIDs (RFC 7292 / NIST).
    private enum OID {
        static let data = "1.2.840.113549.1.7.1"
        static let pkcs8ShroudedKeyBag = "1.2.840.113549.1.12.10.1.2"
        static let certBag = "1.2.840.113549.1.12.10.1.3"
        static let x509Certificate = "1.2.840.113549.1.9.22.1"
        static let localKeyID = "1.2.840.113549.1.9.21"
        static let friendlyName = "1.2.840.113549.1.9.20"
        /// RFC 7292 App. C: PKCS#12 KDF/SHA-1 → 24-byte 3DES key (`id = 1`) + 8-byte IV (`id = 2`).
        static let pbeWithSHAAnd3KeyTripleDESCBC = "1.2.840.113549.1.12.1.3"
        static let sha1 = "1.3.14.3.2.26"
    }

    public enum PKCS12Error: Error, CustomStringConvertible {
        case crypto(String)
        public var description: String { if case .crypto(let s) = self { return "PKCS#12 build failed: \(s)" }; return "" }
    }

    /// Export `certificate` + `privateKey` as a password-protected PKCS#12 blob.
    ///
    /// - Parameters:
    ///   - certificate: the device leaf cert.
    ///   - privateKey: its RSA private key (PKCS#8 DER is shrouded).
    ///   - password: the import password (also placed in the profile's `Password` key).
    ///   - iterations: KDF/MAC iteration count.
    ///   - randomBytes: injectable RNG for the two salts (tests pass a deterministic source). The
    ///     cipher IV is *derived* from the key-bag salt, not random — that is what the v1 PBE spec says.
    public static func export(
        certificate: Certificate,
        privateKey: _RSA.Signing.PrivateKey,
        password: String,
        iterations: Int = 2048,
        randomBytes: (Int) -> [UInt8] = PKCS12.secureRandom
    ) throws -> Data {
        // --- cert bag (plaintext) ---
        var certSer = DER.Serializer()
        try certificate.serialize(into: &certSer)
        let certDER = certSer.serializedBytes
        let localKeyID = Array(SHA256.hash(data: certDER).prefix(20))  // link key↔cert

        let certBagValue = DERBuilder.sequence([
            DERBuilder.oid(OID.x509Certificate),
            DERBuilder.explicit(0, DERBuilder.octetString(certDER)),
        ])
        let certBag = safeBag(
            bagID: OID.certBag, bagValue: certBagValue, localKeyID: localKeyID,
            friendlyName: "Nucleic Device Identity")
        let certSafeContents = DERBuilder.sequence([certBag])
        let certContentInfo = dataContentInfo(certSafeContents)

        // --- shrouded key bag (RFC 7292 v1 PBE: 3DES) ---
        // The shrouded bag wraps a PKCS#8 PrivateKeyInfo (swift-crypto's plain `derRepresentation`
        // and PEM are PKCS#1, so use the explicit PKCS#8 accessor).
        let pkcs8 = Array(privateKey.pkcs8DERRepresentation)
        let epki = try shroudPKCS8(
            pkcs8, password: password, iterations: iterations, randomBytes: randomBytes)
        let keyBag = safeBag(
            bagID: OID.pkcs8ShroudedKeyBag, bagValue: epki, localKeyID: localKeyID,
            friendlyName: "Nucleic Device Identity")
        let keySafeContents = DERBuilder.sequence([keyBag])
        let keyContentInfo = dataContentInfo(keySafeContents)

        // --- AuthenticatedSafe = SEQUENCE OF ContentInfo ---
        let authSafe = DERBuilder.sequence([certContentInfo, keyContentInfo])

        // --- MacData over the AuthenticatedSafe bytes ---
        // HMAC-SHA-1, not SHA-256: `SecItemImport` verifies the MAC first and rejects a SHA-256 one
        // with errSecPkcs12VerifyFailure (-25264), which the profile plugin reports as -323. The input
        // is the complete AuthenticatedSafe TLV — i.e. the contents of the OCTET STRING that
        // `dataContentInfo` wraps below (RFC 7292 §4), verified byte-for-byte against openssl.
        let macSalt = randomBytes(8)
        let macKey = pkcs12KDF(
            password: password, salt: macSalt, id: 3, iterations: iterations, keyLength: 20)
        let mac = Array(
            HMAC<Insecure.SHA1>.authenticationCode(
                for: authSafe, using: SymmetricKey(data: macKey)))
        let digestInfo = DERBuilder.sequence([
            DERBuilder.sequence([DERBuilder.oid(OID.sha1), DERBuilder.null()]),
            DERBuilder.octetString(mac),
        ])
        let macData = DERBuilder.sequence([
            digestInfo, DERBuilder.octetString(macSalt), DERBuilder.integer(iterations),
        ])

        // --- PFX ::= SEQUENCE { version(3), authSafe ContentInfo(data), macData } ---
        let authSafeContentInfo = dataContentInfo(authSafe)
        let pfx = DERBuilder.sequence([
            DERBuilder.integer(3), authSafeContentInfo, macData,
        ])
        return Data(pfx)
    }

    // MARK: - Building blocks

    /// A `ContentInfo` of type `data` wrapping the given DER as an OCTET STRING in `[0] EXPLICIT`.
    private static func dataContentInfo(_ innerDER: [UInt8]) -> [UInt8] {
        DERBuilder.sequence([
            DERBuilder.oid(OID.data),
            DERBuilder.explicit(0, DERBuilder.octetString(innerDER)),
        ])
    }

    /// A `SafeBag ::= SEQUENCE { bagId, [0] EXPLICIT bagValue, SET OF attributes }`.
    private static func safeBag(
        bagID: String, bagValue: [UInt8], localKeyID: [UInt8], friendlyName: String
    ) -> [UInt8] {
        // friendlyName is a BMPString (UTF-16BE); localKeyID an OCTET STRING. Both as PKCS12Attribute
        // { attrId OID, attrValues SET OF }.
        let nameAttr = DERBuilder.sequence([
            DERBuilder.oid(OID.friendlyName),
            DERBuilder.set([bmpString(friendlyName)]),
        ])
        let idAttr = DERBuilder.sequence([
            DERBuilder.oid(OID.localKeyID),
            DERBuilder.set([DERBuilder.octetString(localKeyID)]),
        ])
        return DERBuilder.sequence([
            DERBuilder.oid(bagID),
            DERBuilder.explicit(0, bagValue),
            DERBuilder.set([nameAttr, idAttr]),
        ])
    }

    /// Shroud PKCS#8 bytes into an `EncryptedPrivateKeyInfo` using the RFC 7292 v1 scheme
    /// `pbeWithSHAAnd3-KeyTripleDES-CBC`: the PKCS#12 KDF (SHA-1) derives the 24-byte 3DES key with
    /// `id = 1` and the 8-byte CBC IV with `id = 2` from the *same* salt and iteration count.
    ///
    /// Not PBES2 — see the type's doc comment: `SecItemImport` answers a PBES2/AES-256 key bag with
    /// "Unknown format in import" (measured; other PBES2 ciphers were not tried, since the v1 scheme
    /// is what the importer documents support for).
    private static func shroudPKCS8(
        _ pkcs8: [UInt8], password: String, iterations: Int, randomBytes: (Int) -> [UInt8]
    ) throws -> [UInt8] {
        // 8 bytes is what openssl and Apple's own exporter emit for a v1 PBE salt.
        let salt = randomBytes(8)
        let key = pkcs12KDF(
            password: password, salt: salt, id: 1, iterations: iterations, keyLength: 24)
        let iv = pkcs12KDF(
            password: password, salt: salt, id: 2, iterations: iterations, keyLength: 8)
        let ciphertext: [UInt8]
        do {
            ciphertext = try TripleDES.encryptCBC(pkcs8, key: key, iv: iv)
        } catch { throw PKCS12Error.crypto("3DES-CBC: \(error)") }

        // pkcs-12PbeParams ::= SEQUENCE { salt OCTET STRING, iterations INTEGER }
        let pbeParams = DERBuilder.sequence([
            DERBuilder.octetString(salt), DERBuilder.integer(iterations),
        ])
        let algID = DERBuilder.sequence([
            DERBuilder.oid(OID.pbeWithSHAAnd3KeyTripleDESCBC), pbeParams,
        ])
        // EncryptedPrivateKeyInfo ::= SEQUENCE { algID, encryptedData OCTET STRING }
        return DERBuilder.sequence([algID, DERBuilder.octetString(ciphertext)])
    }

    /// A BMPString (UTF-16BE) TLV — PKCS#12 friendly names use tag 0x1E.
    private static func bmpString(_ s: String) -> [UInt8] {
        var bytes: [UInt8] = []
        for u in s.unicodeScalars {
            let v = UInt16(exactly: u.value) ?? UInt16(0xFFFD)
            bytes.append(UInt8(v >> 8))
            bytes.append(UInt8(v & 0xFF))
        }
        return DERBuilder.primitive(ASN1Identifier(tagWithNumber: 30, tagClass: .universal)) {
            $0.append(contentsOf: bytes)
        }
    }

    // MARK: - PKCS#12 v1 KDF (RFC 7292 Appendix B.2)

    /// Derive `keyLength` bytes per the PKCS#12 KDF with **SHA-1** (u=20, v=64) — the only hash the
    /// v1 PBE algorithms and `SecItemImport` accept. `id` selects the purpose: 1 = cipher key,
    /// 2 = cipher IV, 3 = MAC key. The password is the BMPString encoding **including** the trailing
    /// 0x0000 terminator.
    ///
    /// Verified against openssl: for a `pbeWithSHAAnd3-KeyTripleDES-CBC` bag written by
    /// `openssl pkcs12 -export -legacy`, `id = 1`/`id = 2` reproduce the key and IV that decrypt it,
    /// and `id = 3` reproduces its stored `MacData` HMAC exactly.
    static func pkcs12KDF(
        password: String, salt: [UInt8], id: UInt8, iterations: Int, keyLength: Int
    ) -> [UInt8] {
        let u = 20, v = 64  // SHA-1 digest and block sizes
        // BMP password with a two-byte null terminator.
        var pwd: [UInt8] = []
        for scalar in password.unicodeScalars {
            let value = UInt16(exactly: scalar.value) ?? UInt16(0xFFFD)
            pwd.append(UInt8(value >> 8)); pwd.append(UInt8(value & 0xFF))
        }
        pwd.append(0); pwd.append(0)

        let D = [UInt8](repeating: id, count: v)
        let S = fill(salt, toMultipleOf: v)
        let P = fill(pwd, toMultipleOf: v)
        var I = S + P

        var result: [UInt8] = []
        let blocks = (keyLength + u - 1) / u
        for _ in 0..<blocks {
            var A = Array(Insecure.SHA1.hash(data: D + I))
            for _ in 1..<max(iterations, 1) { A = Array(Insecure.SHA1.hash(data: A)) }
            result.append(contentsOf: A)

            // B = A expanded to v bytes.
            let B = fill(A, toMultipleOf: v).prefix(v)
            // I_j = (I_j + B + 1) mod 2^(8v), for each v-byte block of I.
            var newI: [UInt8] = []
            newI.reserveCapacity(I.count)
            var j = 0
            while j < I.count {
                let block = Array(I[j..<j + v])
                newI.append(contentsOf: addMod(block, Array(B)))
                j += v
            }
            I = newI
        }
        return Array(result.prefix(keyLength))
    }

    /// Repeat `bytes` (a non-empty buffer) to the smallest positive multiple of `n`.
    private static func fill(_ bytes: [UInt8], toMultipleOf n: Int) -> [UInt8] {
        guard !bytes.isEmpty else { return [UInt8](repeating: 0, count: n) }
        let count = ((bytes.count + n - 1) / n) * n
        var out = [UInt8](); out.reserveCapacity(count)
        while out.count < count { out.append(bytes[out.count % bytes.count]) }
        return out
    }

    /// (a + b + 1) mod 2^(8*len), big-endian, both `len` bytes.
    private static func addMod(_ a: [UInt8], _ b: [UInt8]) -> [UInt8] {
        var out = [UInt8](repeating: 0, count: a.count)
        var carry = 1
        var i = a.count - 1
        while i >= 0 {
            let sum = Int(a[i]) + Int(b[i]) + carry
            out[i] = UInt8(sum & 0xFF)
            carry = sum >> 8
            i -= 1
        }
        return out
    }

    /// Cryptographically secure random bytes (default RNG).
    public static func secureRandom(_ count: Int) -> [UInt8] {
        var bytes = [UInt8](repeating: 0, count: count)
        for i in bytes.indices { bytes[i] = UInt8.random(in: 0...255) }
        return bytes
    }
}
