ProxyIOS/ProxyCore/Sources/ProxyEngine/CertificateManager.swift

import Foundation
import X509
import Crypto
import SwiftASN1
import NIOSSL
#if canImport(UIKit)
import UIKit
#endif

/// Manages root CA generation, leaf certificate signing, and an LRU certificate cache.
public final class CertificateManager: @unchecked Sendable {
    public static let shared = CertificateManager()

    private let lock = NSLock()
    private var rootCAKey: P256.Signing.PrivateKey?
    private var rootCACert: Certificate?
    private var rootCANIOSSL: NIOSSLCertificate?

    // LRU cache for generated leaf certificates
    private var certCache: [String: (NIOSSLCertificate, NIOSSLPrivateKey)] = [:]
    private var cacheOrder: [String] = []

    private let keychainCAKeyTag = "com.treyt.proxyapp.ca.privatekey"
    private let keychainCACertTag = "com.treyt.proxyapp.ca.cert"

    private init() {
        loadOrGenerateCA()
    }

    // MARK: - Public API

    public var hasCA: Bool {
        lock.lock()
        defer { lock.unlock() }
        return rootCACert != nil
    }

    /// Get or generate a leaf certificate + key for the given domain.
    public func tlsServerContext(for domain: String) throws -> NIOSSLContext {
        lock.lock()
        defer { lock.unlock() }

        if let cached = certCache[domain] {
            cacheOrder.removeAll { $0 == domain }
            cacheOrder.append(domain)
            return try makeServerContext(cert: cached.0, key: cached.1)
        }

        guard let caKey = rootCAKey, let caCert = rootCACert else {
            throw CertificateError.caNotFound
        }

        let (leafCert, leafKey) = try generateLeaf(domain: domain, caKey: caKey, caCert: caCert)

        // Serialize to DER/PEM for NIOSSL
        var serializer = DER.Serializer()
        try leafCert.serialize(into: &serializer)
        let leafDER = serializer.serializedBytes
        let nioLeafCert = try NIOSSLCertificate(bytes: leafDER, format: .der)
        let leafKeyPEM = leafKey.pemRepresentation
        let nioLeafKey = try NIOSSLPrivateKey(bytes: [UInt8](leafKeyPEM.utf8), format: .pem)

        certCache[domain] = (nioLeafCert, nioLeafKey)
        cacheOrder.append(domain)

        while cacheOrder.count > ProxyConstants.certificateCacheSize {
            let evicted = cacheOrder.removeFirst()
            certCache.removeValue(forKey: evicted)
        }

        return try makeServerContext(cert: nioLeafCert, key: nioLeafKey)
    }

    /// Export the root CA as DER data for user installation.
    public func exportCACertificateDER() -> [UInt8]? {
        lock.lock()
        defer { lock.unlock() }
        guard let cert = rootCACert else { return nil }
        var serializer = DER.Serializer()
        try? cert.serialize(into: &serializer)
        return serializer.serializedBytes
    }

    /// Export as PEM for display.
    public func exportCACertificatePEM() -> String? {
        lock.lock()
        defer { lock.unlock() }
        guard let cert = rootCACert else { return nil }
        guard let pem = try? cert.serializeAsPEM() else { return nil }
        return pem.pemString
    }

    public var caNotValidAfter: Date? {
        lock.lock()
        defer { lock.unlock() }
        // notValidAfter is a Time, not directly a Date — we stored the date when generating
        return nil // Will be set properly after we store dates
    }

    // MARK: - CA Generation

    private func loadOrGenerateCA() {
        if loadCAFromKeychain() { return }

        do {
            let key = P256.Signing.PrivateKey()
            let name = try DistinguishedName {
                CommonName("Proxy CA (\(deviceName()))")
                OrganizationName("ProxyApp")
            }

            let now = Date()
            let twoYearsLater = now.addingTimeInterval(365 * 24 * 3600 * 2)

            let extensions = try Certificate.Extensions {
                Critical(BasicConstraints.isCertificateAuthority(maxPathLength: 0))
                Critical(KeyUsage(keyCertSign: true, cRLSign: true))
            }

            let cert = try Certificate(
                version: .v3,
                serialNumber: Certificate.SerialNumber(),
                publicKey: .init(key.publicKey),
                notValidBefore: now,
                notValidAfter: twoYearsLater,
                issuer: name,
                subject: name,
                signatureAlgorithm: .ecdsaWithSHA256,
                extensions: extensions,
                issuerPrivateKey: .init(key)
            )

            self.rootCAKey = key
            self.rootCACert = cert

            var serializer = DER.Serializer()
            try cert.serialize(into: &serializer)
            let der = serializer.serializedBytes
            self.rootCANIOSSL = try NIOSSLCertificate(bytes: der, format: .der)

            saveCAToKeychain(key: key, certDER: der)
            print("[CertificateManager] Generated new root CA")
        } catch {
            print("[CertificateManager] Failed to generate CA: \(error)")
        }
    }

    // MARK: - Leaf Certificate Generation

    private func generateLeaf(
        domain: String,
        caKey: P256.Signing.PrivateKey,
        caCert: Certificate
    ) throws -> (Certificate, P256.Signing.PrivateKey) {
        let leafKey = P256.Signing.PrivateKey()
        let now = Date()
        let oneYearLater = now.addingTimeInterval(365 * 24 * 3600)

        let extensions = try Certificate.Extensions {
            Critical(BasicConstraints.notCertificateAuthority)
            Critical(KeyUsage(digitalSignature: true))
            try ExtendedKeyUsage([.serverAuth])
            SubjectAlternativeNames([.dnsName(domain)])
        }

        let leafName = try DistinguishedName {
            CommonName(domain)
            OrganizationName("ProxyApp")
        }

        let cert = try Certificate(
            version: .v3,
            serialNumber: Certificate.SerialNumber(),
            publicKey: .init(leafKey.publicKey),
            notValidBefore: now,
            notValidAfter: oneYearLater,
            issuer: caCert.subject,
            subject: leafName,
            signatureAlgorithm: .ecdsaWithSHA256,
            extensions: extensions,
            issuerPrivateKey: .init(caKey)
        )

        return (cert, leafKey)
    }

    // MARK: - TLS Context

    private func makeServerContext(cert: NIOSSLCertificate, key: NIOSSLPrivateKey) throws -> NIOSSLContext {
        var certs = [cert]
        if let caCert = rootCANIOSSL {
            certs.append(caCert)
        }
        var config = TLSConfiguration.makeServerConfiguration(
            certificateChain: certs.map { .certificate($0) },
            privateKey: .privateKey(key)
        )
        config.applicationProtocols = ["http/1.1"]
        return try NIOSSLContext(configuration: config)
    }

    // MARK: - Keychain

    private func loadCAFromKeychain() -> Bool {
        let keyQuery: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrService as String: keychainCAKeyTag,
            kSecAttrAccessGroup as String: ProxyConstants.appGroupIdentifier,
            kSecReturnData as String: true
        ]
        var keyResult: AnyObject?
        guard SecItemCopyMatching(keyQuery as CFDictionary, &keyResult) == errSecSuccess,
              let keyData = keyResult as? Data else { return false }

        let certQuery: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrService as String: keychainCACertTag,
            kSecAttrAccessGroup as String: ProxyConstants.appGroupIdentifier,
            kSecReturnData as String: true
        ]
        var certResult: AnyObject?
        guard SecItemCopyMatching(certQuery as CFDictionary, &certResult) == errSecSuccess,
              let certData = certResult as? Data else { return false }

        do {
            let key = try P256.Signing.PrivateKey(rawRepresentation: keyData)
            let cert = try Certificate(derEncoded: [UInt8](certData))
            let nioCert = try NIOSSLCertificate(bytes: [UInt8](certData), format: .der)

            self.rootCAKey = key
            self.rootCACert = cert
            self.rootCANIOSSL = nioCert
            print("[CertificateManager] Loaded CA from Keychain")
            return true
        } catch {
            print("[CertificateManager] Failed to load CA from Keychain: \(error)")
            return false
        }
    }

    private func saveCAToKeychain(key: P256.Signing.PrivateKey, certDER: [UInt8]) {
        let keyData = key.rawRepresentation

        // Delete existing entries
        for tag in [keychainCAKeyTag, keychainCACertTag] {
            let deleteQuery: [String: Any] = [
                kSecClass as String: kSecClassGenericPassword,
                kSecAttrService as String: tag,
                kSecAttrAccessGroup as String: ProxyConstants.appGroupIdentifier
            ]
            SecItemDelete(deleteQuery as CFDictionary)
        }

        // Save key
        let addKeyQuery: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrService as String: keychainCAKeyTag,
            kSecAttrAccessGroup as String: ProxyConstants.appGroupIdentifier,
            kSecValueData as String: keyData,
            kSecAttrAccessible as String: kSecAttrAccessibleAfterFirstUnlock
        ]
        SecItemAdd(addKeyQuery as CFDictionary, nil)

        // Save cert
        let addCertQuery: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrService as String: keychainCACertTag,
            kSecAttrAccessGroup as String: ProxyConstants.appGroupIdentifier,
            kSecValueData as String: Data(certDER),
            kSecAttrAccessible as String: kSecAttrAccessibleAfterFirstUnlock
        ]
        SecItemAdd(addCertQuery as CFDictionary, nil)
    }

    // MARK: - Helpers

    private func deviceName() -> String {
        #if canImport(UIKit)
        return UIDevice.current.name
        #else
        return Host.current().localizedName ?? "Unknown"
        #endif
    }

    public enum CertificateError: Error {
        case notImplemented
        case generationFailed
        case caNotFound
    }
}