//! SSH enablement for sailboxes, shared by the SDK and the CLI.
//!
//! A box trusts its org's SSH certificate authority rather than individual keys:
//! enabling SSH installs the org CA public key as `TrustedUserCAKeys`,
//! (re)starts `sshd`, and exposes guest port 22 as TCP ingress once the CA-only
//! daemon verifiably owns it. Anyone in the org then connects with a
//! short-lived certificate the org CA signs for their key (minted via
//! [`Client::issue_user_cert`]) without per-box key setup. The host key is
//! generated once with `ssh-keygen -A`, so re-running is safe and a client's
//! `known_hosts` stays valid. openssh is baked into every base and built image,
//! so setup only regenerates host keys and starts the server.

use std::time::{Duration, Instant};

use crate::error::SailError;
use crate::exec::{ExecParams, ExecProcess, EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS};
use crate::Client;

/// Options for [`Sailbox::enable_ssh`](crate::Sailbox::enable_ssh).
#[derive(Debug, Clone)]
pub struct EnableSshOptions {
    /// Source restriction for the exposed port 22 (CIDRs). Empty means any
    /// source on a first enable; a re-enable keeps an existing restriction.
    pub allowlist: Vec<String>,
    /// Wait for the SSH endpoint to be reachable before returning.
    pub wait: bool,
    /// Bound on the reachability wait; 60 seconds is a good default.
    pub timeout: std::time::Duration,
}

impl Default for EnableSshOptions {
    fn default() -> EnableSshOptions {
        EnableSshOptions {
            allowlist: Vec::new(),
            wait: true,
            timeout: std::time::Duration::from_mins(1),
        }
    }
}

/// Where the org CA public key is installed in the guest. sshd accepts any
/// certificate this CA signs; no `authorized_keys` is used.
const SSH_USER_CA_PATH: &str = "/etc/ssh/sail_user_ca.pub";
/// Guest prep before starting sshd: make the config/runtime dirs, generate host
/// keys once (`ssh-keygen -A` never rotates an existing key), and clear root's
/// password (base images ship root locked) so cert login as `root` isn't
/// refused. Password login stays disabled by [`SSHD_START`]'s `-o` flags.
const SSHD_SETUP: &str = "mkdir -p /etc/ssh /run/sshd && ssh-keygen -A && passwd -d root";
const SSHD_SETUP_TIMEOUT_SECONDS: u32 = 60;
/// Kill whichever sshd holds the port-22 listening socket (so the CA-only daemon
/// below takes over), then start it detached. The listener is identified by the
/// socket it owns, not by parentage or process title: `/proc/net/tcp{,6}` gives
/// the inode of the port-22 `LISTEN` socket and `/proc/<pid>/fd` reveals which
/// sshd holds it. Per-connection session children own established sockets, not
/// the listening one, so they are preserved and re-running does not drop
/// connected users; a pre-existing master started any way (our `-D` daemon, or a
/// custom image's `service ssh start`) is replaced, so a leftover
/// password/`authorized_keys` daemon is never left serving. CA-only policy is
/// passed as `-o` options rather than written into `sshd_config` so it cannot be
/// overridden by an existing config or a `Match` block: certificates are the
/// only accepted credential. The `-o` set also forces the auth path a cert needs
/// (`PubkeyAuthentication yes`, `AuthenticationMethods publickey`) so an image
/// config that disabled public-key auth or required a multi-step chain cannot
/// make cert logins silently fail.
const SSHD_START: &str = "ino=$(for f in /proc/net/tcp /proc/net/tcp6; do \
[ -r \"$f\" ] && awk '$4==\"0A\" && $2 ~ /:0016$/ {print $10}' \"$f\"; done); \
pids=''; for d in /proc/[0-9]*; do \
[ \"$(cat \"$d/comm\" 2>/dev/null)\" = sshd ] || continue; \
for fd in \"$d\"/fd/*; do \
l=$(readlink \"$fd\" 2>/dev/null) || continue; \
for i in $ino; do [ \"$l\" = \"socket:[$i]\" ] || continue; \
p=${d#/proc/}; case \" $pids \" in *\" $p \"*) ;; *) pids=\"$pids $p\";; esac; \
done; done; done; \
[ -n \"$pids\" ] && kill $pids 2>/dev/null; sleep 1; \
nohup /usr/sbin/sshd -D -e \
-o 'PermitRootLogin prohibit-password' \
-o 'PasswordAuthentication no' \
-o 'PubkeyAuthentication yes' \
-o 'AuthenticationMethods publickey' \
-o 'TrustedUserCAKeys /etc/ssh/sail_user_ca.pub' \
-o 'AuthorizedKeysFile none' \
-o 'AuthorizedKeysCommand none' </dev/null >/dev/null 2>&1 &";
const VERIFY_CA_SSHD_TIMEOUT_SECONDS: u32 = 30;
/// Confirm the CA-only daemon actually owns port 22 before reporting success.
/// The daemon is backgrounded, so a failed bind (a non-`sshd` service already on
/// port 22, or a killed master slow to release the socket) would otherwise go
/// unnoticed and the box would keep serving a pre-existing password/
/// `authorized_keys` daemon. Poll until the process holding the port-22 listen
/// socket is an sshd whose command line carries our `TrustedUserCAKeys` option;
/// exit non-zero (failing the enable) if it never does.
const VERIFY_CA_SSHD: &str = "for _ in 1 2 3 4 5 6 7 8 9 10; do sleep 1; \
ino=$(for f in /proc/net/tcp /proc/net/tcp6; do \
[ -r \"$f\" ] && awk '$4==\"0A\" && $2 ~ /:0016$/ {print $10}' \"$f\"; done); \
[ -n \"$ino\" ] || continue; \
for d in /proc/[0-9]*; do \
[ \"$(cat \"$d/comm\" 2>/dev/null)\" = sshd ] || continue; \
case \"$(tr '\\0' ' ' < \"$d/cmdline\" 2>/dev/null)\" in *TrustedUserCAKeys*) ;; *) continue;; esac; \
for fd in \"$d\"/fd/*; do l=$(readlink \"$fd\" 2>/dev/null) || continue; \
for i in $ino; do [ \"$l\" = \"socket:[$i]\" ] && exit 0; done; done; \
done; done; \
echo 'CA-only sshd did not take over port 22' >&2; exit 1";
/// The public TCP endpoint a sailbox's SSH listener is reachable at.
#[derive(Debug, Clone)]
pub struct SshEndpoint {
    /// Hostname to dial.
    pub host: String,
    /// Port to dial.
    pub port: u32,
}

fn ssh_exec_params(
    exec_endpoint: &str,
    sailbox_id: &str,
    argv: Vec<String>,
    timeout_seconds: u32,
) -> ExecParams {
    ExecParams {
        sailbox_id: sailbox_id.to_string(),
        exec_endpoint: exec_endpoint.to_string(),
        argv,
        timeout_seconds,
        idempotency_key: String::new(),
        open_stdin: false,
        pty: false,
        term: String::new(),
        cols: 0,
        rows: 0,
        retry_timeout: EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS,
        extra_metadata: Vec::new(),
    }
}

impl Client {
    /// Id-form of [`Sailbox::enable_ssh`](crate::Sailbox::enable_ssh), which
    /// documents the full contract.
    #[doc(hidden)]
    pub async fn enable_ssh(
        &self,
        sailbox_id: &str,
        allowlist: &[String],
        wait: bool,
        timeout: Duration,
    ) -> Result<Option<SshEndpoint>, SailError> {
        // Drop blank allowlist entries (the scheduler does the same before
        // storing) so effectively-empty input counts as omitted rather than
        // taking the replace branch below and clearing an existing restriction.
        let allowlist: Vec<String> = allowlist
            .iter()
            .map(|entry| entry.trim())
            .filter(|entry| !entry.is_empty())
            .map(String::from)
            .collect();

        // Fetch the org CA (read-only; creates the org's CA on first use) before
        // touching the box, so a CA outage fails without resuming, exposing a
        // port, or mutating the guest.
        let ca_public_key = self.org_ssh_ca_public_key().await?;

        let exec_endpoint = self.exec_endpoint(sailbox_id).await?;

        self.ssh_exec_check(
            &exec_endpoint,
            sailbox_id,
            SSHD_SETUP,
            SSHD_SETUP_TIMEOUT_SECONDS,
            "sshd setup",
        )
        .await?;

        // Install the org CA public key the guest sshd trusts (the dir now
        // exists from sshd setup).
        let mut writer = self.worker().write_file(
            &exec_endpoint,
            sailbox_id,
            SSH_USER_CA_PATH,
            /* create_parents */ true,
            Some(0o644),
        );
        writer
            .write_chunk(format!("{}\n", ca_public_key.trim()).into_bytes())
            .await?;
        writer.finish().await?;

        // Start sshd detached; the daemon outlives this exec.
        let proc = ExecProcess::start(
            self.worker(),
            ssh_exec_params(
                &exec_endpoint,
                sailbox_id,
                vec![
                    "/bin/sh".to_string(),
                    "-c".to_string(),
                    SSHD_START.to_string(),
                ],
                30,
            ),
        )
        .await?;
        proc.wait().await?;

        // Confirm the CA-only daemon, not a leftover one, is serving port 22.
        self.ssh_exec_check(
            &exec_endpoint,
            sailbox_id,
            VERIFY_CA_SSHD,
            VERIFY_CA_SSHD_TIMEOUT_SECONDS,
            "sshd ownership check",
        )
        .await?;

        // Expose guest port 22 only now that the CA-only sshd is verified to own
        // it. Expose is declarative — it replaces the stored allowlist — so a
        // non-empty allowlist exposes unconditionally (creating or updating the
        // listener), while an empty one must leave an existing listener
        // untouched: a plain re-enable must not open a restricted port.
        if allowlist.is_empty() {
            match self.get_listener(sailbox_id, 22).await {
                Ok(_) => {}
                Err(SailError::NotFound { .. }) => {
                    self.expose_listener(
                        sailbox_id,
                        22,
                        crate::sailbox::types::IngressProtocol::Tcp,
                        &[],
                    )
                    .await?;
                }
                Err(err) => return Err(err),
            }
        } else {
            self.expose_listener(
                sailbox_id,
                22,
                crate::sailbox::types::IngressProtocol::Tcp,
                &allowlist,
            )
            .await?;
        }

        if !wait {
            return Ok(None);
        }
        self.wait_for_ssh_listener(sailbox_id, timeout)
            .await
            .map(Some)
    }

    /// Run a single shell command in the guest and fail on a non-zero exit.
    async fn ssh_exec_check(
        &self,
        exec_endpoint: &str,
        sailbox_id: &str,
        command: &str,
        timeout_seconds: u32,
        label: &str,
    ) -> Result<(), SailError> {
        let proc = ExecProcess::start(
            self.worker(),
            ssh_exec_params(
                exec_endpoint,
                sailbox_id,
                vec!["/bin/sh".to_string(), "-c".to_string(), command.to_string()],
                timeout_seconds,
            ),
        )
        .await?;
        let result = proc.wait().await?;
        if result.exit_code != 0 {
            let detail = if result.stderr.trim().is_empty() {
                result.stdout.trim()
            } else {
                result.stderr.trim()
            };
            return Err(SailError::Internal {
                message: format!("{label} failed (exit {}): {detail}", result.exit_code),
            });
        }
        Ok(())
    }

    /// Poll the port-22 listener until the SSH endpoint actually accepts, up to
    /// `timeout`. The route flips ACTIVE off the sailbox's running status, not a
    /// guest-port probe, so it can report a host:port before the freshly-started
    /// `sshd` is listening. Connecting and reading the SSH banner is the signal
    /// that `ssh` will work on first use.
    async fn wait_for_ssh_listener(
        &self,
        sailbox_id: &str,
        timeout: Duration,
    ) -> Result<SshEndpoint, SailError> {
        // A saturated "no bound" timeout (Duration::MAX from the bindings)
        // would overflow Instant addition; no deadline means wait indefinitely.
        let deadline = Instant::now().checked_add(timeout);
        loop {
            if let Ok(listener) = self.get_listener(sailbox_id, 22).await {
                if !listener.public_host.is_empty()
                    && listener.public_port != 0
                    && listener.is_active()
                    && ssh_endpoint_accepts(&listener.public_host, listener.public_port).await
                {
                    return Ok(SshEndpoint {
                        host: listener.public_host,
                        port: listener.public_port,
                    });
                }
            }
            if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
                return Err(SailError::Transport {
                    kind: crate::error::TransportKind::Timeout,
                    message: "timed out waiting for the SSH port to become reachable".to_string(),
                    source: None,
                });
            }
            tokio::time::sleep(Duration::from_secs(1)).await;
        }
    }
}

/// Whether the public endpoint answers with an SSH identification banner. A
/// fresh sshd (or the relay before the guest dial succeeds) accepts the TCP
/// connection but sends nothing, so the banner is the signal that ssh will work.
async fn ssh_endpoint_accepts(host: &str, port: u32) -> bool {
    use tokio::io::AsyncReadExt;
    use tokio::net::TcpStream;

    let probe = Duration::from_secs(5);
    let addr = format!("{host}:{port}");
    let Ok(Ok(mut stream)) = tokio::time::timeout(probe, TcpStream::connect(&addr)).await else {
        return false;
    };
    // The banner ("SSH-2.0-...") can arrive split across reads, so accumulate up
    // to its 4-byte prefix before deciding rather than rejecting a partial read.
    let mut buf = [0u8; 4];
    let mut filled = 0;
    while filled < 4 {
        match tokio::time::timeout(probe, stream.read(&mut buf[filled..])).await {
            Ok(Ok(0)) | Err(_) => break,
            Ok(Ok(n)) => filled += n,
            Ok(Err(_)) => break,
        }
    }
    buf[..filled].starts_with(b"SSH-")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn sshd_start_enforces_ca_only_policy() {
        assert!(SSHD_START.contains("TrustedUserCAKeys /etc/ssh/sail_user_ca.pub"));
        assert!(SSHD_START.contains("AuthorizedKeysFile none"));
        assert!(SSHD_START.contains("AuthorizedKeysCommand none"));
        assert!(SSHD_START.contains("PasswordAuthentication no"));
        // A CA cert alone must authenticate, regardless of the image's config.
        assert!(SSHD_START.contains("PubkeyAuthentication yes"));
        assert!(SSHD_START.contains("AuthenticationMethods publickey"));
    }

    /// Guard the string escaping: the literals must be valid `/bin/sh`.
    #[test]
    fn embedded_shell_snippets_are_valid() {
        for (name, snippet) in [
            ("SSHD_START", SSHD_START),
            ("VERIFY_CA_SSHD", VERIFY_CA_SSHD),
        ] {
            let status = std::process::Command::new("sh")
                .args(["-n", "-c", snippet])
                .status()
                .expect("run sh -n");
            assert!(status.success(), "{name} is not valid shell");
        }
    }

    #[test]
    fn verify_matches_the_ca_only_daemon() {
        assert!(VERIFY_CA_SSHD.contains("TrustedUserCAKeys"));
        assert!(SSHD_START.contains("TrustedUserCAKeys"));
    }
}
