foo.swift

import CommonCrypto
import CSSH2
import Foundation

/// Raw libssh2 operations with no orchestration logic.
///
/// Each method performs a single discrete step (TCP connect,
/// handshake, authenticate, open channel). The TransportDriver
/// calls these in the order dictated by the state machine.
public final class SSHConnection: @unchecked Sendable {
    private let host: String
    private let port: Int
    private let user: String

    /// When set, authenticate() tries Secure Enclave auth first.
    public var secureEnclaveKey: SecureEnclaveKey?

    /// Optional log callback; messages also go to stderr.
    public var onLog: (@Sendable (String) -> Void)?

    private var session: OpaquePointer?
    private var channel: OpaquePointer?
    private var sock: Int32 = -1
    private let writeLock = NSLock()

    public init(host: String, port: Int, user: String) {
        self.host = host
        self.port = port
        self.user = user
    }

    deinit {
        close()
    }

    // MARK: - Connection Steps

    /// Step 1: open a TCP socket to the remote host.
    public func tcpConnect() throws {
        var hints = addrinfo()
        hints.ai_family = AF_UNSPEC
        hints.ai_socktype = SOCK_STREAM

        var res: UnsafeMutablePointer<addrinfo>?
        let gaiRc = getaddrinfo(host, String(port), &hints, &res)
        guard gaiRc == 0, let addrList = res else {
            let msg = gaiRc != 0 ? String(cString: gai_strerror(gaiRc)) : "nil result"
            log("getaddrinfo failed: \(msg)")
            throw SSHConnectionError.tcpFailed
        }
        defer { freeaddrinfo(addrList) }

        var info: UnsafeMutablePointer<addrinfo>? = addrList
        while let ai = info {
            let fd = socket(ai.pointee.ai_family, ai.pointee.ai_socktype, ai.pointee.ai_protocol)
            if fd < 0 {
                info = ai.pointee.ai_next
                continue
            }
            if connect(fd, ai.pointee.ai_addr, ai.pointee.ai_addrlen) == 0 {
                sock = fd
                log("TCP connected to \(host):\(port)")
                return
            }
            Darwin.close(fd)
            info = ai.pointee.ai_next
        }

        throw SSHConnectionError.tcpFailed
    }

    /// Step 2: perform the SSH handshake.
    public func handshake() throws {
        guard let sess = libssh2_session_init_ex(nil, nil, nil, nil) else {
            throw SSHConnectionError.handshakeFailed
        }
        session = sess

        let rc = libssh2_session_handshake(sess, sock)
        guard rc == 0 else {
            throw SSHConnectionError.handshakeFailed
        }
        log("SSH handshake complete")
    }

    /// Step 3: authenticate via Secure Enclave, ssh-agent, then key files.
    public func authenticate(delegate: AuthDelegate?) async throws {
        guard let sess = session else {
            throw SSHConnectionError.authFailed
        }

        if trySecureEnclaveAuth(session: sess) {
            return
        }

        if tryAgentAuth(session: sess) {
            return
        }

        let home = NSHomeDirectory()
        let sshDir = "\(home)/.ssh"

        for keyFile in ["id_ed25519_mty", "id_ed25519", "id_rsa"] {
            let privPath = "\(sshDir)/\(keyFile)"
            guard FileManager.default.fileExists(atPath: privPath) else { continue }

            let pubPath = "\(privPath).pub"
            let pubPathOrNil: String? = FileManager.default.fileExists(atPath: pubPath) ? pubPath : nil

            var rc = libssh2_userauth_publickey_fromfile_ex(
                sess, user, UInt32(user.count),
                pubPathOrNil, privPath, nil
            )
            if rc == 0 {
                log("authenticated with \(privPath)")
                return
            }

            if rc == LIBSSH2_ERROR_FILE || rc == LIBSSH2_ERROR_PUBLICKEY_UNVERIFIED {
                if let passphrase = await delegate?.promptPassphrase(for: privPath) {
                    rc = libssh2_userauth_publickey_fromfile_ex(
                        sess, user, UInt32(user.count),
                        pubPathOrNil, privPath, passphrase
                    )
                    if rc == 0 {
                        log("authenticated with \(privPath) (passphrase)")
                        return
                    }
                }
            }

            log("auth failed with \(privPath) (rc=\(rc))")
        }

        log("no usable SSH key found")
        throw SSHConnectionError.authFailed
    }

    /// Execute a command on a new channel and return its stdout.
    /// Must be called after authentication, while the session is
    /// still in blocking mode.
    public func exec(command: String) throws -> String {
        guard let sess = session else {
            throw SSHConnectionError.channelFailed
        }

        libssh2_session_set_blocking(sess, 1)

        guard let chan = libssh2_channel_open_ex(
            sess, "session", UInt32(7),
            UInt32(2 * 1024 * 1024),
            UInt32(32768),
            nil, 0
        ) else {
            throw SSHConnectionError.channelFailed
        }
        defer { libssh2_channel_free(chan) }

        let rc = command.withCString { cmdPtr in
            libssh2_channel_process_startup(
                chan, "exec", 4,
                cmdPtr, UInt32(command.utf8.count)
            )
        }
        guard rc == 0 else {
            throw SSHConnectionError.execFailed
        }

        var output = Data()
        var buf = [CChar](repeating: 0, count: 4096)

        while true {
            let n = libssh2_channel_read_ex(chan, 0, &buf, buf.count)
            if n > 0 {
                buf.withUnsafeBufferPointer { ptr in
                    ptr.baseAddress!.withMemoryRebound(to: UInt8.self, capacity: Int(n)) { base in
                        output.append(base, count: Int(n))
                    }
                }
            } else {
                break
            }
        }

        // Log remote stderr for diagnostics.
        while true {
            let n = libssh2_channel_read_ex(chan, Int32(SSH_EXTENDED_DATA_STDERR), &buf, buf.count)
            if n > 0 {
                let bytes = buf.prefix(Int(n)).map { UInt8(bitPattern: $0) }
                if let str = String(bytes: bytes, encoding: .utf8) {
                    log("remote stderr: \(str)")
                }
            } else {
                break
            }
        }

        guard let result = String(data: output, encoding: .utf8) else {
            throw SSHConnectionError.execFailed
        }

        return result
    }

    /// Step 4: open a channel, request a PTY, start a shell.
    public func openChannel(cols: UInt16, rows: UInt16) throws {
        guard let sess = session else {
            throw SSHConnectionError.channelFailed
        }

        guard let chan = libssh2_channel_open_ex(
            sess, "session", UInt32(7),
            UInt32(2 * 1024 * 1024),
            UInt32(32768),
            nil, 0
        ) else {
            throw SSHConnectionError.channelFailed
        }
        channel = chan

        let term = "xterm-256color"
        var rc = libssh2_channel_request_pty_ex(
            chan, term, UInt32(term.count),
            nil, 0,
            Int32(cols), Int32(rows), 0, 0
        )
        guard rc == 0 else {
            throw SSHConnectionError.ptyFailed
        }

        rc = libssh2_channel_process_startup(chan, "shell", 5, nil, 0)
        guard rc == 0 else {
            throw SSHConnectionError.shellFailed
        }

        log("PTY shell started (\(cols)x\(rows))")
    }

    // MARK: - I/O (hot path)

    /// Switch to non-blocking mode for the read loop.
    public func setNonBlocking() {
        guard let session else { return }
        libssh2_session_set_blocking(session, 0)
    }

    /// Read from the channel. Returns the number of bytes read,
    /// LIBSSH2_ERROR_EAGAIN for would-block, 0 for EOF, or a
    /// negative error code.
    public func channelRead(
        _ buffer: UnsafeMutablePointer<CChar>,
        count: Int,
        stderr useStderr: Bool = false
    ) -> Int {
        guard let channel else { return -1 }
        let stream: Int32 = useStderr ? Int32(SSH_EXTENDED_DATA_STDERR) : 0
        return Int(libssh2_channel_read_ex(channel, stream, buffer, count))
    }

    /// Write bytes to the channel (blocking, thread-safe).
    public func write(_ ptr: UnsafePointer<UInt8>, count: Int) {
        guard let chan = channel, let sess = session else { return }
        writeLock.lock()
        defer { writeLock.unlock() }

        libssh2_session_set_blocking(sess, 1)
        defer { libssh2_session_set_blocking(sess, 0) }

        ptr.withMemoryRebound(to: CChar.self, capacity: count) { cptr in
            var offset = 0
            while offset < count {
                let rc = libssh2_channel_write_ex(chan, 0, cptr.advanced(by: offset), count - offset)
                if rc > 0 {
                    offset += Int(rc)
                } else {
                    break
                }
            }
        }
    }

    /// Resize the remote PTY (blocking, thread-safe).
    public func resizePTY(cols: UInt16, rows: UInt16, widthPx: UInt32, heightPx: UInt32) {
        guard let chan = channel, let sess = session else { return }
        writeLock.lock()
        defer { writeLock.unlock() }

        libssh2_session_set_blocking(sess, 1)
        libssh2_channel_request_pty_size_ex(
            chan, Int32(cols), Int32(rows), Int32(widthPx), Int32(heightPx)
        )
        libssh2_session_set_blocking(sess, 0)
    }

    /// Tear down the connection. Idempotent.
    public func close() {
        if let channel {
            libssh2_channel_free(channel)
            self.channel = nil
        }
        if let session {
            libssh2_session_disconnect_ex(session, SSH_DISCONNECT_BY_APPLICATION, "bye", "")
            libssh2_session_free(session)
            self.session = nil
        }
        if sock >= 0 {
            Darwin.close(sock)
            sock = -1
        }
    }

    // MARK: - Logging

    private func log(_ message: String) {
        fputs("MTY: \(message)\n", stderr)
        onLog?(message)
    }

    // MARK: - Auth Helpers

    /// Context passed through the libssh2 abstract pointer so the
    /// C sign callback can access both the key and the log function.
    private final class SignContext {
        let key: SecureEnclaveKey
        let log: (String) -> Void

        init(key: SecureEnclaveKey, log: @escaping (String) -> Void) {
            self.key = key
            self.log = log
        }
    }

    private func trySecureEnclaveAuth(session sess: OpaquePointer) -> Bool {
        guard let seKey = secureEnclaveKey else { return false }

        let pubKeyData = seKey.publicKeySSHWireFormat

        // Box key + log so the C callback can reach them
        // through the abstract pointer.
        let ctx = SignContext(key: seKey, log: { [weak self] msg in
            self?.log(msg)
        })
        let boxed = Unmanaged.passRetained(ctx)
        defer { boxed.release() }

        var abstract: UnsafeMutableRawPointer? = boxed.toOpaque()

        let rc = pubKeyData.withUnsafeBytes { pubBytes -> Int32 in
            withUnsafeMutablePointer(to: &abstract) { abstractPtr in
                libssh2_userauth_publickey(
                    sess,
                    user,
                    pubBytes.baseAddress!.assumingMemoryBound(to: UInt8.self),
                    pubKeyData.count,
                    { _, sig, sigLen, data, dataLen, abstract -> Int32 in
                        guard let abstract, let ctxPtr = abstract.pointee else {
                            return -1
                        }
                        let ctx = Unmanaged<SSHConnection.SignContext>
                            .fromOpaque(ctxPtr)
                            .takeUnretainedValue()

                        guard let dataPtr = data, dataLen > 0 else { return -1 }
                        let challenge = Data(
                            bytes: dataPtr, count: dataLen
                        )

                        let rawSig: Data
                        do {
                            rawSig = try ctx.key.sign(challenge)
                        } catch {
                            ctx.log("SE sign failed: \(error)")
                            return -1
                        }

                        guard let sigBlob = ecdsaSignatureBlob(rawSig) else {
                            ctx.log("SE signature encoding failed")
                            return -1
                        }

                        // Allocate with malloc; libssh2 uses free()
                        // by default (we pass nil allocators).
                        let buf = malloc(sigBlob.count)?
                            .assumingMemoryBound(to: UInt8.self)
                        guard let buf else { return -1 }
                        sigBlob.copyBytes(to: buf, count: sigBlob.count)
                        sig?.pointee = buf
                        sigLen?.pointee = sigBlob.count
                        return 0
                    },
                    abstractPtr
                )
            }
        }

        if rc == 0 {
            log("authenticated via Secure Enclave (ecdsa-sha2-nistp256)")
            return true
        }

        var errmsg: UnsafeMutablePointer<CChar>?
        var errmsg_len: Int32 = 0
        libssh2_session_last_error(sess, &errmsg, &errmsg_len, 0)
        let detail = errmsg.map { String(cString: $0) } ?? "unknown"
        log("Secure Enclave auth failed (rc=\(rc)): \(detail)")
        log("SE pubkey blob: \(pubKeyData.count) bytes")
        return false
    }

    private func tryAgentAuth(session sess: OpaquePointer) -> Bool {
        guard let agent = libssh2_agent_init(sess) else {
            log("ssh-agent init failed")
            return false
        }
        defer { libssh2_agent_free(agent) }

        guard libssh2_agent_connect(agent) == 0 else {
            log("ssh-agent not available")
            return false
        }
        defer { libssh2_agent_disconnect(agent) }

        guard libssh2_agent_list_identities(agent) == 0 else {
            log("ssh-agent has no identities")
            return false
        }

        var prev: UnsafeMutablePointer<libssh2_agent_publickey>? = nil
        var cur: UnsafeMutablePointer<libssh2_agent_publickey>? = nil

        while libssh2_agent_get_identity(agent, &cur, prev) == 0, let key = cur {
            let rc = libssh2_agent_userauth(agent, user, key)
            if rc == 0 {
                let comment = key.pointee.comment.map { String(cString: $0) } ?? "unknown"
                let keyType = Self.keyTypeFromBlob(
                    blob: key.pointee.blob, len: key.pointee.blob_len
                )
                let fingerprint = Self.sha256Fingerprint(
                    blob: key.pointee.blob, len: key.pointee.blob_len
                )
                log("authenticated via ssh-agent (\(keyType) \(fingerprint) \(comment))")
                return true
            }
            prev = cur
        }

        log("ssh-agent had no accepted key")
        return false
    }

    private static func sha256Fingerprint(blob: UnsafeMutablePointer<UInt8>?, len: Int) -> String {
        guard let blob, len > 0 else { return "SHA256:?" }
        var hash = [UInt8](repeating: 0, count: Int(CC_SHA256_DIGEST_LENGTH))
        CC_SHA256(blob, CC_LONG(len), &hash)
        let b64 = Data(hash).base64EncodedString()
            .trimmingCharacters(in: CharacterSet(charactersIn: "="))
        return "SHA256:\(b64)"
    }

    private static func keyTypeFromBlob(blob: UnsafeMutablePointer<UInt8>?, len: Int) -> String {
        guard let blob, len > 4 else { return "unknown" }
        let typeLen = Int(blob[0]) << 24 | Int(blob[1]) << 16 | Int(blob[2]) << 8 | Int(blob[3])
        guard typeLen > 0, typeLen + 4 <= len else { return "unknown" }
        return String(
            bytes: UnsafeBufferPointer(start: blob.advanced(by: 4), count: typeLen),
            encoding: .utf8
        ) ?? "unknown"
    }
}

public enum SSHConnectionError: Error, Sendable {
    case tcpFailed
    case handshakeFailed
    case authFailed
    case channelFailed
    case ptyFailed
    case shellFailed
    case execFailed
}