import Crypto
import Foundation
import X509
import XCTest
import _CryptoExtras

@testable import NucleicMDM

final class MDMCertificateAuthorityTests: XCTestCase {
    func testGenerateRootIsSelfSignedCA() throws {
        let ca = try MDMCertificateAuthority.generate()
        XCTAssertEqual(ca.rootCertificate.issuer, ca.rootCertificate.subject)
        let bc = try XCTUnwrap(
            ca.rootCertificate.extensions.first { $0.oid == .X509ExtensionID.basicConstraints })
        let parsed = try BasicConstraints(bc)
        guard case .isCertificateAuthority = parsed else {
            return XCTFail("root is not marked as a CA")
        }
    }

    func testPersistRoundTrip() throws {
        let ca = try MDMCertificateAuthority.generate()
        let reloaded = try MDMCertificateAuthority(persisted: ca.persisted())
        XCTAssertEqual(
            try ca.rootCertificate.serializeAsPEM().pemString,
            try reloaded.rootCertificate.serializeAsPEM().pemString)
        // The reloaded CA still issues (its key survived the PEM round-trip).
        _ = try reloaded.issueServerCertificate(dnsName: "mdm.nucleic.internal")
    }

    func testServerCertificateHasDNSSANAndServerAuth() throws {
        let ca = try MDMCertificateAuthority.generate()
        let issued = try ca.issueServerCertificate(
            dnsName: "mdm.nucleic.internal", ipSANs: ["192.168.66.1"])
        // Chains to the root: issuer DN matches, and the signature verifies.
        XCTAssertEqual(issued.certificate.issuer, ca.rootCertificate.subject)
        let san = try SubjectAlternativeNames(
            XCTUnwrap(
                issued.certificate.extensions.first {
                    $0.oid == .X509ExtensionID.subjectAlternativeName
                }))
        let dnsNames = san.compactMap { name -> String? in
            if case .dnsName(let n) = name { return n }
            return nil
        }
        XCTAssertTrue(dnsNames.contains("mdm.nucleic.internal"))
        let eku = try ExtendedKeyUsage(
            XCTUnwrap(
                issued.certificate.extensions.first {
                    $0.oid == .X509ExtensionID.extendedKeyUsage
                }))
        XCTAssertTrue(eku.contains(.serverAuth))
    }

    func testDeviceCertificateFromCSR() throws {
        let ca = try MDMCertificateAuthority.generate()
        // A device-side keypair + CSR (mimics what mdmclient's SCEP client produces).
        let deviceKey = try _RSA.Signing.PrivateKey(keySize: .bits2048)
        let deviceCertKey = Certificate.PrivateKey(deviceKey)
        let subject = try DistinguishedName { CommonName("Nucleic macOS Guest") }
        let csr = try CertificateSigningRequest(
            version: .v1, subject: subject, privateKey: deviceCertKey,
            attributes: CertificateSigningRequest.Attributes(),
            signatureAlgorithm: .sha256WithRSAEncryption)
        let deviceCert = try ca.issueDeviceCertificate(fromCSR: csr)
        XCTAssertEqual(deviceCert.issuer, ca.rootCertificate.subject)
        XCTAssertEqual(deviceCert.subject, subject)
        let eku = try ExtendedKeyUsage(
            XCTUnwrap(
                deviceCert.extensions.first { $0.oid == .X509ExtensionID.extendedKeyUsage }))
        XCTAssertTrue(eku.contains(.clientAuth))
    }

    /// Serials must be unique per issuer: X.509 identifies a cert by (issuer, serialNumber) and the
    /// keychain keys lookups on that pair. Every leaf used to be stamped serial 2, so the device
    /// identity and the TLS server cert collided — a live macOS 27 run failed profile installation
    /// with "The certificate could not be verified (authentication error)".
    func testEveryIssuedLeafGetsADistinctSerial() throws {
        let ca = try MDMCertificateAuthority.generate()
        var serials: Set<[UInt8]> = []
        for i in 0..<8 {
            serials.insert(Array(try ca.issueServerCertificate(dnsName: "s\(i).host")
                .certificate.serialNumber.bytes))
        }
        let deviceKey = try _RSA.Signing.PrivateKey(keySize: .bits2048)
        let csr = try CertificateSigningRequest(
            version: .v1, subject: try DistinguishedName { CommonName("Device") },
            privateKey: Certificate.PrivateKey(deviceKey),
            attributes: CertificateSigningRequest.Attributes(),
            signatureAlgorithm: .sha256WithRSAEncryption)
        serials.insert(Array(try ca.issueDeviceCertificate(fromCSR: csr).serialNumber.bytes))

        XCTAssertEqual(serials.count, 9, "every issued leaf must carry a distinct serial")
        // The root keeps serial 1, and no leaf may collide with it.
        XCTAssertEqual(Array(ca.rootCertificate.serialNumber.bytes), [1])
        XCTAssertFalse(serials.contains([1]))
        // Positive, minimally-encoded, and wide enough to be collision-free.
        for s in serials {
            XCTAssertEqual(s.count, 16)
            XCTAssertLessThan(s[0], 0x80, "serial must encode as a positive ASN.1 INTEGER")
            XCTAssertNotEqual(s[0], 0, "no redundant leading zero octet")
        }
    }
}

final class MDMEnrollmentProfileTests: XCTestCase {
    private func makeProfile(uuids: MDMEnrollmentProfile.UUIDs = MDMEnrollmentProfile.UUIDs())
        throws -> (MDMConfiguration, MDMDeviceIdentity, MDMCertificateAuthority, [String: Any])
    {
        let ca = try MDMCertificateAuthority.generate()
        let config = MDMConfiguration.generate()
        let identity = try MDMDeviceIdentity.issue(from: ca)
        let data = try MDMEnrollmentProfile.make(
            config: config, identity: identity,
            caCertificateDER: Data(try ca.rootCertificateDER()), uuids: uuids)
        return (config, identity, ca, try MDMPlist.parse(data))
    }

    func testProfileCarriesRootIdentityAndMDMPayloads() throws {
        let (_, _, _, dict) = try makeProfile()
        XCTAssertEqual(dict["PayloadType"] as? String, "Configuration")
        let payloads = try XCTUnwrap(dict["PayloadContent"] as? [[String: Any]])
        let types = payloads.compactMap { $0["PayloadType"] as? String }
        // The CA root rides in the profile so installing it establishes trust (no add-trusted-cert),
        // and it is delivered FIRST so the root is present before the MDM check-in TLS runs.
        XCTAssertEqual(types.first, "com.apple.security.root")
        XCTAssertTrue(types.contains("com.apple.security.pkcs12"))
        XCTAssertTrue(types.contains("com.apple.mdm"))
    }

    func testRootPayloadHoldsTheCADER() throws {
        let (_, _, ca, dict) = try makeProfile()
        let payloads = try XCTUnwrap(dict["PayloadContent"] as? [[String: Any]])
        let root = try XCTUnwrap(
            payloads.first { $0["PayloadType"] as? String == "com.apple.security.root" })
        let der = try XCTUnwrap(root["PayloadContent"] as? Data)
        XCTAssertEqual(Array(der), try ca.rootCertificateDER())
    }

    func testMDMPayloadReferencesPKCS12IdentityAndAdvertisesServerCapabilities() throws {
        let uuids = MDMEnrollmentProfile.UUIDs()
        let (config, identity, _, dict) = try makeProfile(uuids: uuids)
        let payloads = try XCTUnwrap(dict["PayloadContent"] as? [[String: Any]])
        let mdm = try XCTUnwrap(payloads.first { $0["PayloadType"] as? String == "com.apple.mdm" })
        XCTAssertEqual(mdm["IdentityCertificateUUID"] as? String, uuids.identity.uuidString)
        XCTAssertEqual(mdm["ServerURL"] as? String, config.serverURL)
        XCTAssertEqual(mdm["Topic"] as? String, config.topic)
        XCTAssertEqual(mdm["SignMessage"] as? Bool, true)
        // ServerCapabilities is MANDATORY on macOS 27 — the install fails without it. The server acks
        // the resulting per-user check-ins it does not model (Workstream 2.2).
        XCTAssertEqual(
            mdm["ServerCapabilities"] as? [String], ["com.apple.mdm.per-user-connections"])

        // Certificate pinning is what makes this whole pipeline viable: `mdmclient` evaluates our
        // server against the ROOT payload carried in this same profile, not against the guest's trust
        // store — which is the only option left, because macOS 27 has no non-interactive way to put a
        // self-issued root into that store (five routes closed; debug log runs 10–18). Both URLs must
        // be pinned: pinning only the command channel would leave the check-in handshake — the FIRST
        // connection enrollment makes — validating against a store that does not trust us.
        XCTAssertEqual(
            mdm["CheckInURLPinningCertificateUUIDs"] as? [String], [uuids.root.uuidString])
        XCTAssertEqual(
            mdm["ServerURLPinningCertificateUUIDs"] as? [String], [uuids.root.uuidString])
        // The pinned UUID must actually name a certificate payload that ships in this profile —
        // a dangling UUID pins to nothing and fails closed.
        XCTAssertTrue(
            payloads.contains {
                $0["PayloadUUID"] as? String == uuids.root.uuidString
                    && ($0["PayloadType"] as? String)?.hasPrefix("com.apple.security") == true
            },
            "the pinned UUID must reference a certificate payload inside this profile")
        // A private CA publishes no CRL/OCSP responder, so a required revocation check would fail
        // every connection.
        XCTAssertEqual(mdm["PinningRevocationCheckRequired"] as? Bool, false)

        let p12 = try XCTUnwrap(
            payloads.first { $0["PayloadType"] as? String == "com.apple.security.pkcs12" })
        XCTAssertEqual(p12["Password"] as? String, identity.password)
        let blob = try XCTUnwrap(p12["PayloadContent"] as? Data)
        // A PKCS#12 PFX is a DER SEQUENCE (0x30) and non-trivially sized.
        XCTAssertEqual(blob.first, 0x30)
        XCTAssertGreaterThan(blob.count, 500)
    }
}
