qlean 0.3.0

use std::{
    io::ErrorKind,
    os::unix::fs::PermissionsExt,
    path::{Path, PathBuf},
    sync::Arc,
    time::Duration,
};

use anyhow::{Context, Result, bail};
use russh::{
    ChannelMsg, Disconnect,
    keys::{
        PrivateKey, PrivateKeyWithHashAlg, PublicKey,
        ssh_key::{LineEnding, private::Ed25519Keypair, rand_core::OsRng},
    },
};
use russh_sftp::{client::SftpSession, protocol::OpenFlags};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    time::{Instant, sleep},
};
use tokio_util::sync::CancellationToken;
use tokio_vsock::{VsockAddr, VsockStream};
use tracing::{debug, info};

const ERR_EACCES: i32 = 13; // Permission denied
const ERR_ENODEV: i32 = 19; // No such device

#[derive(Clone, Debug)]
pub(crate) struct PersistedSshKeypair {
    pub pubkey_str: String,
    pub _pubkey_path: PathBuf,
    pub privkey_str: String,
    pub privkey_path: PathBuf,
}

impl PersistedSshKeypair {
    // Try to load a keypair from `dir`
    pub(crate) fn from_dir(dir: &Path) -> Result<Self> {
        let privkey_path = dir.join("id_ed25519");
        let pubkey_path = privkey_path.with_extension("pub");
        let privkey_str = std::fs::read_to_string(&privkey_path)?;
        let pubkey_str = std::fs::read_to_string(&pubkey_path)?;

        Ok(Self {
            pubkey_str,
            _pubkey_path: pubkey_path,
            privkey_str,
            privkey_path,
        })
    }
}

pub(crate) fn get_ssh_key(dir: &Path) -> Result<PersistedSshKeypair> {
    // First try reading an existing keypair from disk.
    // If that fails we'll just create a new one.
    if let Ok(existing_keypair) = PersistedSshKeypair::from_dir(dir) {
        return Ok(existing_keypair);
    }

    let privkey_path = dir.join("id_ed25519");
    let pubkey_path = privkey_path.with_extension("pub");

    let ed25519_keypair = Ed25519Keypair::random(&mut OsRng);

    let pubkey_openssh = PublicKey::from(ed25519_keypair.public).to_openssh()?;
    debug!("Writing SSH public key to {pubkey_path:?}");
    std::fs::write(&pubkey_path, &pubkey_openssh)?;

    let privkey_openssh = PrivateKey::from(ed25519_keypair)
        .to_openssh(LineEnding::default())?
        .to_string();
    debug!("Writing SSH private key to {privkey_path:?}");

    std::fs::write(&privkey_path, &privkey_openssh)?;
    let mut perms = std::fs::metadata(&privkey_path)?.permissions();
    perms.set_mode(0o600);
    std::fs::set_permissions(&privkey_path, perms)?;

    let keypair = PersistedSshKeypair {
        pubkey_str: pubkey_openssh,
        _pubkey_path: pubkey_path,
        privkey_str: privkey_openssh,
        privkey_path,
    };
    Ok(keypair)
}

#[derive(Debug, Clone, Default)]
struct SshClient {}

// More SSH event handlers can be defined in this trait
//
// In this example, we're only using Channel, so these aren't needed.
impl russh::client::Handler for SshClient {
    type Error = russh::Error;

    async fn check_server_key(
        &mut self,
        _server_public_key: &russh::keys::PublicKey,
    ) -> Result<bool, Self::Error> {
        Ok(true)
    }
}

/// This struct is a convenience wrapper around a russh client that handles the input/output event
/// loop
pub(crate) struct Session {
    session: russh::client::Handle<SshClient>,
    // Cached SFTP session for reuse; lazily initialized
    sftp: Option<SftpSession>,
}

impl Session {
    /// Connect to an SSH server via vsock
    async fn connect(
        privkey: PrivateKey,
        username: &str,
        cid: u32,
        port: u32,
        timeout: Duration,
        cancel_token: CancellationToken,
    ) -> Result<Self> {
        let config = russh::client::Config {
            keepalive_interval: Some(Duration::from_secs(5)),
            ..<_>::default()
        };

        let config = Arc::new(config);
        let sh = SshClient {};

        let vsock_addr = VsockAddr::new(cid, port);
        let now = Instant::now();
        let mut session =
            loop {
                // Check for cancellation
                if cancel_token.is_cancelled() {
                    bail!("SSH connection cancelled");
                }

                tokio::time::sleep(Duration::from_millis(100)).await;

                // Establish vsock connection
                let connect_budget = timeout
                    .saturating_sub(now.elapsed())
                    .max(Duration::from_millis(1));
                let stream =
                    match tokio::time::timeout(connect_budget, VsockStream::connect(vsock_addr))
                        .await
                    {
                        Ok(Ok(stream)) => stream,
                        Err(_) => {
                            bail!("Timeout while connecting to VM via vsock.");
                        }
                        Ok(Err(ref e)) if e.raw_os_error() == Some(ERR_EACCES) => {
                            bail!("Permission denied while connecting via vsock: {e}\n");
                        }
                        Ok(Err(ref e)) if e.raw_os_error() == Some(ERR_ENODEV) => {
                            // ENODEV is commonly observed while QEMU is still booting/initializing the vsock
                            // transport (e.g. the guest CID isn't ready yet). Treat it as transient and retry
                            // until the overall timeout is reached.
                            debug!("SSH vsock connect not ready yet (ENODEV): {e} (will retry)");
                            if now.elapsed() > timeout {
                                bail!("Timeout while connecting to VM via vsock.\n");
                            }
                            continue;
                        }
                        Ok(Err(ref e)) => match e.kind() {
                            ErrorKind::TimedOut
                            | ErrorKind::ConnectionRefused
                            | ErrorKind::ConnectionReset
                            | ErrorKind::NetworkUnreachable
                            | ErrorKind::AddrNotAvailable => {
                                if now.elapsed() > timeout {
                                    bail!("Timeout while connecting to VM via vsock.");
                                }
                                continue;
                            }
                            e => {
                                bail!("SSH vsock connect error: {e}");
                            }
                        },
                    };

                // Connect to SSH via vsock stream
                let handshake_budget = timeout
                    .saturating_sub(now.elapsed())
                    .max(Duration::from_millis(1));
                match tokio::time::timeout(
                    handshake_budget,
                    russh::client::connect_stream(config.clone(), stream, sh.clone()),
                )
                .await
                {
                    Ok(Ok(x)) => break x,
                    Err(_) => {
                        bail!("Timeout establishing SSH handshake over vsock.");
                    }
                    Ok(Err(russh::Error::IO(ref e))) => {
                        match e.kind() {
                            // The VM is still booting at this point so we're just ignoring these errors
                            // for some time.
                            ErrorKind::ConnectionRefused
                            | ErrorKind::ConnectionReset
                            | ErrorKind::UnexpectedEof => {
                                if now.elapsed() > timeout {
                                    bail!("Timeout establishing SSH handshake over vsock.");
                                }
                            }
                            e => {
                                bail!("SSH handshake error: {e}");
                            }
                        }
                    }
                    Ok(Err(russh::Error::Disconnect)) => {
                        if now.elapsed() > timeout {
                            bail!("Timeout establishing SSH handshake over vsock.");
                        }
                    }
                    Ok(Err(e)) => {
                        bail!("SSH client error: {e}");
                    }
                }
            };
        debug!("Authenticating via SSH");
        let auth_budget = timeout
            .saturating_sub(now.elapsed())
            .max(Duration::from_secs(1));
        let auth_res = tokio::time::timeout(
            auth_budget,
            session.authenticate_publickey(
                username,
                PrivateKeyWithHashAlg::new(Arc::new(privkey), None),
            ),
        )
        .await
        .with_context(|| format!("SSH authentication timed out for user {username}"))??;
        if !auth_res.success() {
            bail!("Authentication (with publickey) failed");
        }
        Ok(Self {
            session,
            sftp: None,
        })
    }

    // NOTE: TCP-based SSH fallback intentionally not supported.
    // Reviewer feedback: avoid architecture changes that introduce non-vsock transports.
    /// Open an SFTP session over the existing SSH connection.
    async fn open_sftp(&mut self) -> Result<SftpSession> {
        let channel = self.session.channel_open_session().await?;
        channel.request_subsystem(true, "sftp").await?;
        let sftp = SftpSession::new(channel.into_stream()).await?;
        Ok(sftp)
    }

    /// Get a cached SFTP session, opening one if needed.
    pub(crate) async fn get_sftp(&mut self) -> Result<&mut SftpSession> {
        if self.sftp.is_none() {
            let sftp = self.open_sftp().await?;
            self.sftp = Some(sftp);
        }
        Ok(self.sftp.as_mut().expect("SFTP session must exist"))
    }

    /// Call a command via SSH, streaming its output to stdout/stderr.
    pub(crate) async fn call(
        &mut self,
        // env: HashMap<String, String>,
        command: &str,
        cancel_token: CancellationToken,
    ) -> Result<u32> {
        let mut channel = self.session.channel_open_session().await?;

        // for (key, value) in env {
        //     channel.set_env(true, &key, &value).await?;
        // }

        //channel.request_shell(true).await?;
        channel.exec(true, command).await?;

        let code;
        let mut stdout = tokio::io::stdout();
        let mut stderr = tokio::io::stderr();

        loop {
            // Check for cancellation
            if cancel_token.is_cancelled() {
                info!("SSH call cancelled during execution");
                bail!("SSH call cancelled");
            }

            let Some(msg) = channel.wait().await else {
                bail!(
                    "SSH channel closed before exit status for command: {}",
                    command
                );
            };
            match msg {
                // Write data to the terminal
                ChannelMsg::Data { ref data } => {
                    stdout.write_all(data).await?;
                    stdout.flush().await?;
                }
                ChannelMsg::ExtendedData { ref data, ext: 1 } => {
                    // ext == 1 means it's stderr content
                    // https://github.com/Eugeny/russh/discussions/258
                    stderr.write_all(data).await?;
                    stderr.flush().await?;
                }

                // The command has returned an exit code
                ChannelMsg::ExitStatus { exit_status } => {
                    code = exit_status;
                    channel.eof().await?;
                    break;
                }
                _ => {}
            }
        }
        Ok(code)
    }

    /// Call a command via SSH and capture its output.
    pub(crate) async fn call_with_output(
        &mut self,
        command: &str,
        cancel_token: CancellationToken,
    ) -> Result<(u32, Vec<u8>, Vec<u8>)> {
        let mut channel = self.session.channel_open_session().await?;
        channel.exec(true, command).await?;

        let code;
        let mut stdout = Vec::new();
        let mut stderr = Vec::new();

        loop {
            // Check for cancellation
            if cancel_token.is_cancelled() {
                info!("SSH call cancelled during execution");
                bail!("SSH call cancelled");
            }

            let Some(msg) = channel.wait().await else {
                bail!(
                    "SSH channel closed before exit status for command: {}",
                    command
                );
            };
            match msg {
                // Write data to the buffer
                ChannelMsg::Data { ref data } => {
                    stdout.extend_from_slice(data);
                }
                ChannelMsg::ExtendedData { ref data, ext: 1 } => {
                    // ext == 1 means it's stderr content
                    // https://github.com/Eugeny/russh/discussions/258
                    stderr.extend_from_slice(data);
                }
                // The command has returned an exit code
                ChannelMsg::ExitStatus { exit_status } => {
                    code = exit_status;
                    channel.eof().await?;
                    break;
                }
                _ => {}
            }
        }
        Ok((code, stdout, stderr))
    }

    pub(crate) async fn close(&mut self) -> Result<()> {
        self.session
            .disconnect(Disconnect::ByApplication, "", "English")
            .await?;
        Ok(())
    }
}

/// Connect SSH and run a command that checks whether the system is ready for operation.
pub(crate) async fn connect_ssh(
    cid: u32,
    timeout: Duration,
    keypair: PersistedSshKeypair,
    cancel_token: CancellationToken,
    _mac_address: String,
) -> Result<Session> {
    if !std::path::Path::new("/dev/vhost-vsock").exists() {
        bail!("/dev/vhost-vsock is missing. Qlean requires vhost-vsock for SSH (no TCP fallback).")
    }

    let privkey = PrivateKey::from_openssh(&keypair.privkey_str)?;
    let deadline = Instant::now() + timeout;
    let mut last_err: Option<anyhow::Error> = None;

    while Instant::now() < deadline {
        if cancel_token.is_cancelled() {
            bail!("SSH connection cancelled");
        }

        let remaining = deadline.saturating_duration_since(Instant::now());
        let per_attempt_timeout = Duration::from_secs(12)
            .min(remaining.max(Duration::from_secs(1)))
            .min(Duration::from_secs(25));

        match Session::connect(
            privkey.clone(),
            "root",
            cid,
            22,
            per_attempt_timeout,
            cancel_token.clone(),
        )
        .await
        {
            Ok(mut session) => {
                info!("✅ Connected via vsock as root");

                let ready_budget = deadline.saturating_duration_since(Instant::now());
                if ready_budget == Duration::ZERO {
                    bail!("SSH connection timed out");
                }

                let _ =
                    tokio::time::timeout(ready_budget, session.call("true", cancel_token.clone()))
                        .await
                        .context("SSH readiness probe timed out")??;

                debug!("SSH command channel is ready");
                return Ok(session);
            }
            Err(e) => {
                last_err = Some(e);
                sleep(Duration::from_millis(250)).await;
            }
        }
    }

    Err(last_err.unwrap_or_else(|| anyhow::anyhow!("Timeout establishing SSH over vsock.")))
}

impl Session {
    /// Recursively create a directory and all of its parent components if they are missing.
    pub(crate) async fn create_dir_all<P: AsRef<Path>>(&mut self, path: P) -> Result<()> {
        let path = path.as_ref();
        // Build path incrementally like mkdir -p
        let mut cur = PathBuf::new();
        for comp in path.components() {
            cur.push(comp);
            if cur.as_os_str().is_empty() {
                continue;
            }
            // Limit SFTP borrow scope to avoid conflicts with self in recursion
            let create_res = {
                let sftp = self.get_sftp().await?;
                sftp.create_dir(cur.to_string_lossy()).await
            };
            match create_res {
                Ok(_) => {}
                Err(e) => {
                    let meta_res = {
                        let sftp = self.get_sftp().await?;
                        sftp.metadata(cur.to_string_lossy()).await
                    };
                    if let Ok(attr) = meta_res {
                        if !attr.is_dir() {
                            bail!("Remote path exists and is not a directory: {:?}", cur);
                        }
                    } else {
                        bail!("Failed to create remote directory {:?}: {}", cur, e);
                    }
                }
            }
        }
        Ok(())
    }

    /// Upload a single file via SFTP.
    pub(crate) async fn upload_file<P: AsRef<Path>, Q: AsRef<Path>>(
        &mut self,
        local: P,
        remote: Q,
        cancel_token: CancellationToken,
    ) -> anyhow::Result<()> {
        let local = local.as_ref();
        let remote = remote.as_ref();
        let mut src = tokio::fs::File::open(local).await?;
        // Scope SFTP borrow
        let mut dst = {
            let sftp = self.get_sftp().await?;
            sftp.open_with_flags(
                remote.to_string_lossy(),
                OpenFlags::CREATE | OpenFlags::TRUNCATE | OpenFlags::WRITE,
            )
            .await?
        };

        let mut buf = vec![0u8; 128 * 1024];
        loop {
            if cancel_token.is_cancelled() {
                bail!("Upload cancelled");
            }
            let n = AsyncReadExt::read(&mut src, &mut buf).await?;
            if n == 0 {
                break;
            }
            AsyncWriteExt::write_all(&mut dst, &buf[..n]).await?;
        }
        let _ = AsyncWriteExt::flush(&mut dst).await;
        let _ = AsyncWriteExt::shutdown(&mut dst).await;
        Ok(())
    }

    /// Download a single file via SFTP.
    pub(crate) async fn download_file<P: AsRef<Path>, Q: AsRef<Path>>(
        &mut self,
        remote: P,
        local: Q,
        cancel_token: CancellationToken,
    ) -> anyhow::Result<()> {
        let remote = remote.as_ref();
        let local = local.as_ref();
        let mut src = {
            let sftp = self.get_sftp().await?;
            sftp.open(remote.to_string_lossy()).await?
        };
        let mut dst = tokio::fs::File::create(local).await?;

        let mut buf = vec![0u8; 128 * 1024];
        loop {
            if cancel_token.is_cancelled() {
                bail!("Download cancelled");
            }
            let n = AsyncReadExt::read(&mut src, &mut buf).await?;
            if n == 0 {
                break;
            }
            AsyncWriteExt::write_all(&mut dst, &buf[..n]).await?;
        }
        let _ = AsyncWriteExt::flush(&mut dst).await;
        Ok(())
    }

    /// Walk a remote directory tree over SFTP, similar to walkdir.
    /// Returns a depth-first list of entries including the root.
    pub(crate) async fn walk_remote_dir<P: AsRef<Path>>(
        &mut self,
        root: P,
        follow_links: bool,
        cancel_token: CancellationToken,
    ) -> Result<Vec<RemoteDirEntry>> {
        let root = root.as_ref();
        let mut out = Vec::new();

        // Stat root
        let root_meta = {
            let sftp = self.get_sftp().await?;
            sftp.metadata(root.to_string_lossy()).await?
        };
        let root_type = RemoteFileType::from_attrs(&root_meta);
        out.push(RemoteDirEntry::new(root.to_path_buf(), root_type));

        // If root is not a dir, nothing more to traverse
        if !out[0].file_type.is_dir() {
            return Ok(out);
        }

        // DFS stack of directories to visit
        let mut stack = vec![root.to_path_buf()];
        while let Some(dir) = stack.pop() {
            if cancel_token.is_cancelled() {
                bail!("Walk cancelled");
            }

            let entries = {
                let sftp = self.get_sftp().await?;
                match sftp.read_dir(dir.to_string_lossy()).await {
                    Ok(e) => e,
                    Err(e) => {
                        // If directory can't be read, skip (best-effort)
                        debug!("Failed to read_dir {:?}: {}", dir, e);
                        continue;
                    }
                }
            };

            for entry in entries {
                let name = entry.file_name();
                if name == "." || name == ".." {
                    continue;
                }

                let child_path = dir.join(&name);
                let attrs = {
                    let sftp = self.get_sftp().await?;
                    match sftp.metadata(child_path.to_string_lossy()).await {
                        Ok(a) => a,
                        Err(e) => {
                            debug!("Failed to stat {:?}: {}", child_path, e);
                            continue;
                        }
                    }
                };
                let ftype = RemoteFileType::from_attrs(&attrs);
                out.push(RemoteDirEntry::new(child_path.clone(), ftype.clone()));

                if ftype.is_dir() {
                    stack.push(child_path);
                } else if ftype.is_symlink() && follow_links {
                    // If it's a symlink and we're following links, stat the target
                    let target_path = {
                        let sftp = self.get_sftp().await?;
                        match sftp.read_link(child_path.to_string_lossy()).await {
                            Ok(tp) => PathBuf::from(tp),
                            Err(e) => {
                                bail!("Failed to read_link {:?}: {}", child_path, e);
                            }
                        }
                    };
                    let target_attrs = {
                        let sftp = self.get_sftp().await?;
                        match sftp.metadata(target_path.to_string_lossy()).await {
                            Ok(a) => a,
                            Err(e) => {
                                bail!("Failed to stat symlink target {:?}: {}", target_path, e);
                            }
                        }
                    };
                    let target_type = RemoteFileType::from_attrs(&target_attrs);
                    if target_type.is_dir() {
                        stack.push(target_path);
                    }
                }
            }
        }

        Ok(out)
    }

    /// Get the primary IP address of the remote machine.
    pub(crate) async fn get_remote_ip(&mut self) -> Result<String> {
        let (code, stdout, _stderr) = self
            .call_with_output("hostname -I | awk '{print $1}'", CancellationToken::new())
            .await?;
        if code != 0 {
            bail!("Failed to get remote IP address, exit code {}", code);
        }
        let ip = String::from_utf8(stdout)?.trim().to_string();
        Ok(ip)
    }
}

#[derive(Clone, Debug)]
pub(crate) struct RemoteFileType {
    is_dir: bool,
    is_file: bool,
    is_symlink: bool,
}

impl RemoteFileType {
    fn from_attrs(attrs: &russh_sftp::protocol::FileAttributes) -> Self {
        Self {
            is_dir: attrs.is_dir(),
            is_file: attrs.file_type().is_file(),
            is_symlink: attrs.file_type().is_symlink(),
        }
    }
    pub(crate) fn is_dir(&self) -> bool {
        self.is_dir
    }
    pub(crate) fn is_file(&self) -> bool {
        self.is_file
    }
    pub(crate) fn is_symlink(&self) -> bool {
        self.is_symlink
    }
}

#[derive(Clone, Debug)]
pub(crate) struct RemoteDirEntry {
    path: PathBuf,
    file_type: RemoteFileType,
}

impl RemoteDirEntry {
    fn new(path: PathBuf, file_type: RemoteFileType) -> Self {
        Self { path, file_type }
    }
    pub(crate) fn path(&self) -> &Path {
        &self.path
    }

    pub(crate) fn file_type(&self) -> &RemoteFileType {
        &self.file_type
    }
}