import Foundation
import Security
import CryptoKit

/// How to evaluate the server's TLS certificate.
///
/// MCMS appliances usually present a self-signed or internal-CA certificate.
/// Pick the *narrowest* policy that works for your deployment. Never ship
/// `.allowSelfSignedForHost` against a host you don't control.
enum ServerTrustPolicy: Equatable, Codable {
    /// Default OS trust evaluation (use this if the box has a real cert or a
    /// CA your devices already trust via MDM/profile — strongly preferred).
    case system

    /// Accept the server's certificate ONLY for this exact host, after a normal
    /// trust evaluation that ignores the "untrusted root" failure. Scoped escape
    /// hatch for self-signed boxes on a controlled mgmt network/VPN.
    case allowSelfSignedForHost(String)

    /// Certificate/public-key pinning: trust only if the leaf's SHA-256
    /// public-key hash (base64) matches one of these. The safest option for a
    /// fixed appliance — capture the hash once and pin it.
    case pinPublicKeySHA256([String], host: String)
}

/// `URLSessionDelegate` that applies a `ServerTrustPolicy`. Retained by the
/// `URLSession` it's attached to.
final class ServerTrustEvaluator: NSObject, URLSessionDelegate {
    private let policy: ServerTrustPolicy

    init(policy: ServerTrustPolicy) {
        self.policy = policy
    }

    func urlSession(
        _ session: URLSession,
        didReceive challenge: URLAuthenticationChallenge,
        completionHandler: @escaping (URLSession.AuthChallengeDisposition, URLCredential?) -> Void
    ) {
        guard challenge.protectionSpace.authenticationMethod == NSURLAuthenticationMethodServerTrust,
              let trust = challenge.protectionSpace.serverTrust else {
            completionHandler(.performDefaultHandling, nil)
            return
        }

        let host = challenge.protectionSpace.host

        switch policy {
        case .system:
            completionHandler(.performDefaultHandling, nil)

        case .allowSelfSignedForHost(let allowedHost):
            // Scoped escape hatch for self-signed appliances on a controlled
            // network. Because a self-signed root fails normal evaluation, this
            // accepts ANY certificate the allowed host presents — including
            // expired or name-mismatched ones. The only gate is the hostname, so
            // use this only on hosts you control; prefer `.pinPublicKeySHA256`
            // once you can capture the appliance's key hash.
            guard host == allowedHost else {
                completionHandler(.cancelAuthenticationChallenge, nil)
                return
            }
            completionHandler(.useCredential, URLCredential(trust: trust))

        case .pinPublicKeySHA256(let pins, let pinnedHost):
            guard host == pinnedHost,
                  let leaf = Self.leafCertificate(from: trust),
                  let hash = Self.publicKeySHA256Base64(for: leaf),
                  pins.contains(hash) else {
                completionHandler(.cancelAuthenticationChallenge, nil)
                return
            }
            completionHandler(.useCredential, URLCredential(trust: trust))
        }
    }

    // MARK: - Pinning helpers

    private static func leafCertificate(from trust: SecTrust) -> SecCertificate? {
        if #available(iOS 15.0, *) {
            return (SecTrustCopyCertificateChain(trust) as? [SecCertificate])?.first
        } else {
            return SecTrustGetCertificateAtIndex(trust, 0)
        }
    }

    /// Base64 of SHA-256 over the DER-encoded SubjectPublicKeyInfo's key bits.
    /// (Capture once with `openssl` and store as a pin; this returns the raw
    /// public-key hash for comparison.)
    static func publicKeySHA256Base64(for certificate: SecCertificate) -> String? {
        guard let key = SecCertificateCopyKey(certificate),
              let data = SecKeyCopyExternalRepresentation(key, nil) as Data? else {
            return nil
        }
        return CryptoSHA256.base64(of: data)
    }
}

/// Minimal SHA-256 wrapper using CryptoKit.
enum CryptoSHA256 {
    static func base64(of data: Data) -> String {
        let digest = SHA256.hash(data: data)
        return Data(digest).base64EncodedString()
    }
}