ProxyIOS/ProxyCore/Sources/ProxyEngine/CertificateManager.swift

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

/// Manages the shared MITM root CA. The app owns generation and writes the CA
/// into the App Group container; the extension only loads that shared identity.
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?
    private var caFingerprintCache: String?
    private var certificateMTime: Date?
    private var keyMTime: Date?

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

    private init() {
        loadOrGenerateCAIfNeeded()
    }

    // MARK: - Public API

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

    public var caFingerprint: String? {
        lock.lock()
        defer { lock.unlock() }
        refreshFromDiskLocked()
        return caFingerprintCache
    }

    public var canGenerateCA: Bool {
        Bundle.main.infoDictionary?["NSExtension"] == nil
    }

    public func reloadSharedCA() {
        lock.lock()
        refreshFromDiskLocked(force: true)
        lock.unlock()
    }

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

        refreshFromDiskLocked()

        if let cached = certCache[domain] {
            cacheOrder.removeAll { $0 == domain }
            cacheOrder.append(domain)
            ProxyLogger.cert.debug("TLS context CACHE HIT for \(domain)")
            return try makeServerContext(cert: cached.0, key: cached.1)
        }

        guard let caKey = rootCAKey, let caCert = rootCACert else {
            ProxyLogger.cert.error("TLS context FAILED for \(domain): no CA loaded. hasKey=\(self.rootCAKey != nil) hasCert=\(self.rootCACert != nil)")
            throw CertificateError.caNotFound
        }

        ProxyLogger.cert.info("TLS: generating leaf cert for \(domain), CA issuer=\(String(describing: caCert.subject))")

        do {
            let (leafCert, leafKey) = try generateLeaf(domain: domain, caKey: caKey, caCert: caCert)
            ProxyLogger.cert.info("TLS: leaf cert generated for \(domain), SAN=\(domain), notBefore=\(leafCert.notValidBefore), notAfter=\(leafCert.notValidAfter)")

            var serializer = DER.Serializer()
            try leafCert.serialize(into: &serializer)
            let leafDER = serializer.serializedBytes
            ProxyLogger.cert.debug("TLS: leaf DER size=\(leafDER.count) bytes")

            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)
            }

            let ctx = try makeServerContext(cert: nioLeafCert, key: nioLeafKey)
            ProxyLogger.cert.info("TLS: server context READY for \(domain)")
            return ctx
        } catch {
            ProxyLogger.cert.error("TLS: leaf cert/context FAILED for \(domain): \(error)")
            throw error
        }
    }

    /// Export the root CA as DER data for user installation.
    public func exportCACertificateDER() -> [UInt8]? {
        lock.lock()
        defer { lock.unlock() }
        refreshFromDiskLocked()
        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() }
        refreshFromDiskLocked()
        guard let cert = rootCACert else { return nil }
        guard let pem = try? cert.serializeAsPEM() else { return nil }
        return pem.pemString
    }

    public var caGeneratedDate: Date? {
        lock.lock()
        defer { lock.unlock() }
        refreshFromDiskLocked()
        return rootCACert?.notValidBefore
    }

    public var caExpirationDate: Date? {
        lock.lock()
        defer { lock.unlock() }
        refreshFromDiskLocked()
        return rootCACert?.notValidAfter
    }

    public func regenerateCA() {
        guard canGenerateCA else {
            ProxyLogger.cert.error("Refusing to regenerate CA from extension context")
            return
        }

        lock.lock()
        clearStateLocked()
        deleteStoredCALocked()

        do {
            try generateAndStoreCALocked()
        } catch {
            ProxyLogger.cert.error("CA regeneration failed: \(error.localizedDescription)")
        }
        lock.unlock()
    }

    // MARK: - CA bootstrap

    private func loadOrGenerateCAIfNeeded() {
        lock.lock()
        defer { lock.unlock() }

        refreshFromDiskLocked(force: true)
        guard rootCACert == nil || rootCAKey == nil else { return }

        guard canGenerateCA else {
            ProxyLogger.cert.info("Shared CA not found; extension will remain passthrough-only")
            return
        }

        do {
            try generateAndStoreCALocked()
        } catch {
            ProxyLogger.cert.error("Failed to generate shared CA: \(error.localizedDescription)")
        }
    }

    private func refreshFromDiskLocked(force: Bool = false) {
        let certURL = AppGroupPaths.caCertificateURL
        let keyURL = AppGroupPaths.caPrivateKeyURL

        let certExists = FileManager.default.fileExists(atPath: certURL.path)
        let keyExists = FileManager.default.fileExists(atPath: keyURL.path)
        guard certExists, keyExists else {
            if rootCACert != nil || rootCAKey != nil {
                ProxyLogger.cert.info("Shared CA files missing; clearing in-memory state")
            }
            clearStateLocked()
            return
        }

        let currentCertMTime = modificationDate(for: certURL)
        let currentKeyMTime = modificationDate(for: keyURL)
        if !force, currentCertMTime == certificateMTime, currentKeyMTime == keyMTime {
            return
        }

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

            rootCAKey = key
            rootCACert = cert
            rootCANIOSSL = nioCert
            certificateMTime = currentCertMTime
            keyMTime = currentKeyMTime
            caFingerprintCache = fingerprint(for: certData)
            certCache.removeAll()
            cacheOrder.removeAll()
            ProxyLogger.cert.info("Loaded shared CA from App Group container")
        } catch {
            ProxyLogger.cert.error("Failed to load shared CA: \(error.localizedDescription)")
            clearStateLocked()
        }
    }

    private func generateAndStoreCALocked() throws {
        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)
        )

        var serializer = DER.Serializer()
        try cert.serialize(into: &serializer)
        let der = serializer.serializedBytes

        try FileManager.default.createDirectory(
            at: AppGroupPaths.certificatesDirectory,
            withIntermediateDirectories: true,
            attributes: nil
        )
        try Data(der).write(to: AppGroupPaths.caCertificateURL, options: .atomic)
        try key.rawRepresentation.write(to: AppGroupPaths.caPrivateKeyURL, options: .atomic)

        rootCAKey = key
        rootCACert = cert
        rootCANIOSSL = try NIOSSLCertificate(bytes: der, format: .der)
        certificateMTime = modificationDate(for: AppGroupPaths.caCertificateURL)
        keyMTime = modificationDate(for: AppGroupPaths.caPrivateKeyURL)
        caFingerprintCache = fingerprint(for: Data(der))
        certCache.removeAll()
        cacheOrder.removeAll()

        ProxyLogger.cert.info("Generated new shared root CA")
    }

    private func clearStateLocked() {
        rootCAKey = nil
        rootCACert = nil
        rootCANIOSSL = nil
        caFingerprintCache = nil
        certificateMTime = nil
        keyMTime = nil
        certCache.removeAll()
        cacheOrder.removeAll()
    }

    private func deleteStoredCALocked() {
        for url in [AppGroupPaths.caCertificateURL, AppGroupPaths.caPrivateKeyURL] {
            try? FileManager.default.removeItem(at: url)
        }
    }

    // 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: - Helpers

    private func modificationDate(for url: URL) -> Date? {
        (try? FileManager.default.attributesOfItem(atPath: url.path)[.modificationDate] as? Date) ?? nil
    }

    private func fingerprint(for data: Data) -> String {
        SHA256.hash(data: data).map { String(format: "%02x", $0) }.joined()
    }

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

    public enum CertificateError: Error {
        case caNotFound
    }
}