nucleic/Sources/NucleicProtocol/Noise/SecureChannel.swift

import Foundation
import CryptoKit

/// A device's long-term identity (SYNC_PROTOCOL §4.1): an X25519 static keypair (the Noise
/// `s`) plus an Ed25519 signing key. On-device these live in the Secure Enclave / Keychain;
/// this type is the in-memory handle and a raw-bytes codec for that storage.
public struct DeviceIdentity: Sendable {
    public let staticKey: Curve25519.KeyAgreement.PrivateKey
    public let signingKey: Curve25519.Signing.PrivateKey

    public init(
        staticKey: Curve25519.KeyAgreement.PrivateKey = .init(),
        signingKey: Curve25519.Signing.PrivateKey = .init()
    ) {
        self.staticKey = staticKey
        self.signingKey = signingKey
    }

    public var staticPublicKey: Data { staticKey.publicKey.rawRepresentation }
    public var signingPublicKey: Data { signingKey.publicKey.rawRepresentation }

    /// Stable fingerprint of the static key for display / TXT-record advertising
    /// (structural metadata only — safe to log/show).
    public var fingerprint: String { Self.fingerprint(ofStaticKey: staticPublicKey) }

    /// Fingerprint of any static public key — the phone uses this to match a Bonjour result
    /// against the host key it pinned at pairing.
    public static func fingerprint(ofStaticKey key: Data) -> String {
        let digest = SHA256.hash(data: key)
        return digest.prefix(8).map { String(format: "%02x", $0) }.joined()
    }

    // Raw persistence (store these bytes in the Keychain).
    public func exportRaw() -> Data {
        staticKey.rawRepresentation + signingKey.rawRepresentation
    }
    public init(importingRaw data: Data) throws {
        guard data.count == 64 else { throw SecureChannelError.badIdentityBlob }
        staticKey = try .init(rawRepresentation: data.prefix(32))
        signingKey = try .init(rawRepresentation: data.suffix(32))
    }

    public func sign(_ data: Data) throws -> Data { try signingKey.signature(for: data) }
    public static func verify(_ signature: Data, of data: Data, publicKey: Data) -> Bool {
        guard let pk = try? Curve25519.Signing.PublicKey(rawRepresentation: publicKey) else { return false }
        return pk.isValidSignature(signature, for: data)
    }
}

public enum SecureChannelError: Error, Equatable {
    case badIdentityBlob
    case handshakeIncomplete
    case notEstablished
}

/// An established, encrypted application session over a completed Noise handshake. Wraps the
/// two transport `CipherState`s and (de)serializes application messages as CBOR. Drive it
/// from the transport's serial receive/send path — it is not internally synchronized.
public final class SecureSession {
    private var send: CipherState
    private var recv: CipherState
    /// The peer's pinned static public key (who we're talking to).
    public let remoteStaticKey: Data
    /// Channel-binding hash — identical on both ends; basis for an out-of-band verify code.
    public let handshakeHash: Data

    init(send: CipherState, recv: CipherState, remoteStaticKey: Data, handshakeHash: Data) {
        self.send = send
        self.recv = recv
        self.remoteStaticKey = remoteStaticKey
        self.handshakeHash = handshakeHash
    }

    /// Encrypt an application message → the opaque frame payload (caller length-prefixes it).
    public func seal<T: Encodable>(_ message: T) throws -> Data {
        let plaintext = try CBOREncoder().encode(message)
        return try send.encrypt(ad: Data(), plaintext: plaintext)
    }

    /// Decrypt a received frame payload → an application message.
    public func open<T: Decodable>(_ type: T.Type, from frame: Data) throws -> T {
        let plaintext = try recv.decrypt(ad: Data(), ciphertext: frame)
        return try CBORDecoder().decode(T.self, from: plaintext)
    }
}

// MARK: - Pairing QR payload (SYNC §4.2)

/// What the Mac encodes into the pairing QR: its static public key, a one-time pairing
/// secret (the Noise PSK), a human label, and a connection hint for whichever transport the
/// host is listening on — LAN (`lanHost`/`lanPort`) or tailnet (`tailnetHost`/`tailnetPort`).
/// The phone scans this to bootstrap the `xxpsk0` pairing handshake bound to exactly this QR.
///
/// `transport` is a raw string (see `SyncTransportHint`), not an enum, so an old phone
/// scanning a newer QR — or this build scanning a future transport — still decodes the
/// payload and can say "unsupported" instead of failing the scan. All hint fields are
/// optional; a payload from an older host simply decodes them as nil (LAN).
public struct PairingPayload: Sendable, Codable, Equatable {
    public let protocolVersion: Int
    public let hostName: String
    public let hostStaticKey: Data
    public let pairingSecret: Data   // 32 bytes; the Noise PSK
    public let lanHost: String?
    public let lanPort: UInt16?
    /// Which transport the host is listening on (`SyncTransportHint` raw value); nil = LAN.
    public let transport: String?
    /// The host's tailnet IP (v4 preferred) when `transport == .tailnet`.
    public let tailnetHost: String?
    public let tailnetPort: UInt16?

    public init(
        protocolVersion: Int = SyncProtocol.version,
        hostName: String, hostStaticKey: Data, pairingSecret: Data,
        lanHost: String? = nil, lanPort: UInt16? = nil,
        transport: String? = nil, tailnetHost: String? = nil, tailnetPort: UInt16? = nil
    ) {
        self.protocolVersion = protocolVersion
        self.hostName = hostName
        self.hostStaticKey = hostStaticKey
        self.pairingSecret = pairingSecret
        self.lanHost = lanHost
        self.lanPort = lanPort
        self.transport = transport
        self.tailnetHost = tailnetHost
        self.tailnetPort = tailnetPort
    }

    /// The transport hint: `.lan` when absent (pre-transport hosts), nil when the QR names
    /// a transport this build doesn't know — callers surface "update the app" rather than
    /// silently dialing the wrong way.
    public var transportHint: SyncTransportHint? {
        guard let transport else { return .lan }
        return SyncTransportHint(rawValue: transport)
    }

    public static func freshSecret() -> Data {
        Data((0..<32).map { _ in UInt8.random(in: .min ... .max) })
    }

    /// Compact text for a QR code: `nucleic://pair?d=<base64url(cbor)>`.
    public func qrString() throws -> String {
        let cbor = try CBOREncoder().encode(self)
        return "nucleic://pair?d=" + base64URL(cbor)
    }

    public init(qrString: String) throws {
        guard let range = qrString.range(of: "d=") else { throw SecureChannelError.badIdentityBlob }
        let encoded = String(qrString[range.upperBound...])
        guard let cbor = Self.dataFromBase64URL(encoded) else { throw SecureChannelError.badIdentityBlob }
        self = try CBORDecoder().decode(PairingPayload.self, from: cbor)
    }
}

private func base64URL(_ data: Data) -> String {
    data.base64EncodedString()
        .replacingOccurrences(of: "+", with: "-")
        .replacingOccurrences(of: "/", with: "_")
        .replacingOccurrences(of: "=", with: "")
}

extension PairingPayload {
    fileprivate static func dataFromBase64URL(_ s: String) -> Data? {
        var b64 = s.replacingOccurrences(of: "-", with: "+").replacingOccurrences(of: "_", with: "/")
        while b64.count % 4 != 0 { b64.append("=") }
        return Data(base64Encoded: b64)
    }
}

// MARK: - Handshake convenience (host = responder, phone = initiator)

public enum NoiseSession {
    /// Phone side. Pairing uses `xxpsk0` with the QR secret; reconnect uses `IK` against the
    /// pinned host static key. Returns a driver and a finisher producing the `SecureSession`.
    public static func makeInitiator(
        identity: DeviceIdentity, hostStaticKey: Data, psk: Data?
    ) -> HandshakeDriver {
        let pattern: NoiseHandshake.Pattern = psk != nil ? .xxpsk0 : .ik
        let hs = NoiseHandshake(
            pattern: pattern, initiator: true, staticKey: identity.staticKey,
            remoteStatic: hostStaticKey, psk: psk)
        return HandshakeDriver(handshake: hs)
    }

    /// Mac side. Mirrors the phone: `xxpsk0` during pairing, `IK` on reconnect.
    public static func makeResponder(
        identity: DeviceIdentity, psk: Data?
    ) -> HandshakeDriver {
        let pattern: NoiseHandshake.Pattern = psk != nil ? .xxpsk0 : .ik
        let hs = NoiseHandshake(
            pattern: pattern, initiator: false, staticKey: identity.staticKey, psk: psk)
        return HandshakeDriver(handshake: hs)
    }
}

/// Thin stateful wrapper that steps a handshake and, when complete, mints a `SecureSession`.
public final class HandshakeDriver {
    private var handshake: NoiseHandshake

    init(handshake: NoiseHandshake) { self.handshake = handshake }

    public var isComplete: Bool { handshake.isComplete }
    public var isMyTurnToWrite: Bool { handshake.isMyTurnToWrite }

    public func write(payload: Data = Data()) throws -> Data {
        try handshake.writeMessage(payload: payload)
    }
    public func read(_ message: Data) throws -> Data {
        try handshake.readMessage(message)
    }

    /// Available once `isComplete` — the established encrypted session.
    public func session() throws -> SecureSession {
        guard handshake.isComplete, let rs = handshake.remoteStaticKey else {
            throw SecureChannelError.handshakeIncomplete
        }
        let ciphers = try handshake.transportCiphers()
        return SecureSession(
            send: ciphers.send, recv: ciphers.recv,
            remoteStaticKey: rs, handshakeHash: handshake.handshakeHash)
    }
}