session.rs

use crate::host_metadata::StrictHostKeyPolicy;
use crate::ssh_config::HostConfig;
use anyhow::Context;
use portable_pty::{native_pty_system, CommandBuilder, MasterPty, PtyPair, PtySize};
use serde::Serialize;
use std::collections::HashMap;
use std::env;
use std::io::{Read, Write};
use std::path::Path;
use std::sync::{Arc, Mutex};
use tauri::{AppHandle, Emitter};
use uuid::Uuid;

#[derive(Debug, Clone, Serialize)]
pub struct SessionOutputEvent {
    pub session_id: String,
    pub chunk: String,
    pub host_key_prompt: bool,
}

type SharedWriter = Arc<Mutex<Box<dyn Write + Send>>>;
type SharedMaster = Arc<Mutex<Box<dyn MasterPty + Send>>>;
type SharedChild = Arc<Mutex<Box<dyn portable_pty::Child + Send + Sync>>>;

/// Maximum bytes to buffer per IPC event to avoid oversized messages.
const SESSION_OUTPUT_COALESCE_MAX_BYTES: usize = 16 * 1024;
const SESSION_OUTPUT_HOST_KEY_NEEDLE: &str = "Are you sure you want to continue connecting";

pub struct SessionHandle {
    writer: SharedWriter,
    master: SharedMaster,
    child: SharedChild,
}

#[derive(Default)]
pub struct SessionState {
    sessions: Mutex<HashMap<String, SessionHandle>>,
}

/// Expand `~/…` / `~\…` using the real home directory. Windows OpenSSH does not reliably treat
/// `~` in `-i` / `UserKnownHostsFile` like Unix shells, and GUI apps often lack `ssh` on `PATH`.
fn expand_ssh_user_path(raw: &str) -> String {
    let raw = raw.trim();
    if raw.is_empty() {
        return String::new();
    }
    if let Some(rest) = raw.strip_prefix("~/").or_else(|| raw.strip_prefix("~\\")) {
        if let Some(home) = home::home_dir() {
            return home.join(rest).to_string_lossy().into_owned();
        }
    }
    raw.to_string()
}

fn ssh_user_known_hosts_option() -> String {
    let path = crate::ssh_home::effective_ssh_dir()
        .map(|d| d.join("known_hosts").to_string_lossy().into_owned())
        .unwrap_or_else(|_| "~/.ssh/known_hosts".to_string());
    format!("UserKnownHostsFile={path}")
}

fn resolve_ssh_program() -> String {
    #[cfg(target_os = "windows")]
    {
        let windir = env::var("WINDIR").unwrap_or_else(|_| "C:\\Windows".to_string());
        let openssh = Path::new(&windir).join("System32").join("OpenSSH").join("ssh.exe");
        if openssh.is_file() {
            return openssh.to_string_lossy().into_owned();
        }
        let fallback = Path::new(r"C:\Windows\System32\OpenSSH\ssh.exe");
        if fallback.is_file() {
            return fallback.to_string_lossy().into_owned();
        }
        for key in ["ProgramFiles", "ProgramFiles(x86)"] {
            if let Ok(pf) = env::var(key) {
                let git_ssh = Path::new(&pf).join("Git").join("usr").join("bin").join("ssh.exe");
                if git_ssh.is_file() {
                    return git_ssh.to_string_lossy().into_owned();
                }
            }
        }
        "ssh".to_string()
    }
    #[cfg(not(target_os = "windows"))]
    {
        "ssh".to_string()
    }
}

pub fn build_ssh_command(
    host: &HostConfig,
    strict_host_key_checking: &str,
    connect_timeout_secs: u32,
) -> CommandBuilder {
    let ssh = resolve_ssh_program();
    let mut command = CommandBuilder::new(&ssh);
    command.arg("-o");
    command.arg(format!("StrictHostKeyChecking={strict_host_key_checking}"));
    command.arg("-o");
    command.arg(ssh_user_known_hosts_option());
    command.arg("-o");
    command.arg(format!("ConnectTimeout={connect_timeout_secs}"));
    command.arg("-o");
    command.arg("ConnectionAttempts=1");
    command.arg("-p");
    command.arg(host.port.to_string());
    if !host.identity_file.is_empty() {
        let identity = expand_ssh_user_path(&host.identity_file);
        if !identity.is_empty() {
            command.arg("-i");
            command.arg(identity);
        }
    }
    if !host.proxy_jump.is_empty() {
        command.arg("-J");
        command.arg(host.proxy_jump.clone());
    }
    if !host.proxy_command.is_empty() {
        command.arg("-o");
        command.arg(format!("ProxyCommand={}", host.proxy_command));
    }
    if !host.user.is_empty() {
        command.arg(format!("{}@{}", host.user, host.host_name));
    } else {
        command.arg(host.host_name.clone());
    }
    command
}

fn resolve_local_shell_path(explicit_shell: Option<&str>) -> String {
    if let Some(shell) = explicit_shell {
        let trimmed = shell.trim();
        if !trimmed.is_empty() {
            return trimmed.to_string();
        }
    }
    #[cfg(not(target_os = "windows"))]
    {
        if let Ok(shell_from_env) = env::var("SHELL") {
            let trimmed = shell_from_env.trim();
            if !trimmed.is_empty() {
                return trimmed.to_string();
            }
        }
        if Path::new("/bin/zsh").exists() {
            return "/bin/zsh".to_string();
        }
        if Path::new("/bin/bash").exists() {
            return "/bin/bash".to_string();
        }
        return "sh".to_string();
    }
    #[cfg(target_os = "windows")]
    {
        // Prefer PowerShell; fall back to cmd.exe.
        if let Ok(windir) = env::var("WINDIR") {
            let ps = Path::new(&windir)
                .join("System32")
                .join("WindowsPowerShell")
                .join("v1.0")
                .join("powershell.exe");
            if ps.exists() {
                return ps.to_string_lossy().into_owned();
            }
            let cmd = Path::new(&windir).join("System32").join("cmd.exe");
            if cmd.exists() {
                return cmd.to_string_lossy().into_owned();
            }
        }
        "powershell.exe".to_string()
    }
}

pub fn build_local_shell_command(explicit_shell: Option<&str>) -> CommandBuilder {
    let shell = resolve_local_shell_path(explicit_shell);
    #[allow(unused_mut)]
    let mut command = CommandBuilder::new(&shell);
    // `-l` (login shell) is a POSIX convention; PowerShell and cmd.exe do not support it.
    #[cfg(not(target_os = "windows"))]
    command.arg("-l");
    command
}

impl SessionState {
    /// `quick_policy` overrides host-metadata resolution (used for Quick Connect).
    pub fn start(
        &self,
        app: AppHandle,
        host: HostConfig,
        quick_policy: Option<StrictHostKeyPolicy>,
    ) -> anyhow::Result<String> {
        let pty_system = native_pty_system();
        let pair = pty_system
            .openpty(PtySize {
                rows: 30,
                cols: 120,
                pixel_width: 0,
                pixel_height: 0,
            })
            .context("failed to allocate pty")?;
        let sk = match quick_policy {
            Some(p) => p.as_ssh_value(),
            None => crate::host_metadata::resolved_strict_host_key_for_alias(&host.host),
        };
        let connect_secs = crate::app_prefs::current_preferences().connect_timeout_secs;
        let command = build_ssh_command(&host, sk, connect_secs);
        self.spawn_and_register_command(app, pair, command)
    }

    pub fn start_local(&self, app: AppHandle) -> anyhow::Result<String> {
        let pty_system = native_pty_system();
        let pair = pty_system
            .openpty(PtySize {
                rows: 30,
                cols: 120,
                pixel_width: 0,
                pixel_height: 0,
            })
            .context("failed to allocate pty")?;
        let mut command = build_local_shell_command(None);
        if let Some(home_dir) = home::home_dir() {
            command.cwd(home_dir);
        }
        self.spawn_and_register_command(app, pair, command)
    }

    fn spawn_and_register_command(
        &self,
        app: AppHandle,
        pair: PtyPair,
        mut command: CommandBuilder,
    ) -> anyhow::Result<String> {
        command.env("TERM", "xterm-256color");

        let child = pair
            .slave
            .spawn_command(command)
            .context("failed to spawn ssh process")?;
        drop(pair.slave);

        let mut reader = pair
            .master
            .try_clone_reader()
            .context("failed to create pty reader")?;
        let writer = pair
            .master
            .take_writer()
            .context("failed to create pty writer")?;

        let session_id = Uuid::new_v4().to_string();
        let child = Arc::new(Mutex::new(child));
        let writer = Arc::new(Mutex::new(writer));
        let master = Arc::new(Mutex::new(pair.master));

        {
            let mut sessions = self
                .sessions
                .lock()
                .map_err(|_| anyhow::anyhow!("session lock poisoned"))?;
            sessions.insert(
                session_id.clone(),
                SessionHandle {
                    writer: writer.clone(),
                    master: master.clone(),
                    child: child.clone(),
                },
            );
        }

        let session_id_for_thread = session_id.clone();
        std::thread::spawn(move || {
            let emit_chunk = |chunk: String| {
                let host_key_prompt = chunk.contains(SESSION_OUTPUT_HOST_KEY_NEEDLE);
                let _ = app.emit(
                    "session-output",
                    SessionOutputEvent {
                        session_id: session_id_for_thread.clone(),
                        chunk,
                        host_key_prompt,
                    },
                );
            };

            let mut buf = [0_u8; 8192];
            let mut pending = String::new();

            let flush_pending = |pending: &mut String, emit: &dyn Fn(String)| {
                if pending.is_empty() {
                    return;
                }
                let chunk = std::mem::take(pending);
                emit(chunk);
            };

            loop {
                match reader.read(&mut buf) {
                    Ok(0) => {
                        flush_pending(&mut pending, &emit_chunk);
                        break;
                    }
                    Ok(read_len) => {
                        let fragment = String::from_utf8_lossy(&buf[..read_len]);
                        pending.push_str(&fragment);

                        // Emit in MAX_BYTES-sized chunks to avoid oversized IPC messages.
                        while pending.len() > SESSION_OUTPUT_COALESCE_MAX_BYTES {
                            let rest = pending.split_off(SESSION_OUTPUT_COALESCE_MAX_BYTES);
                            let chunk = std::mem::replace(&mut pending, rest);
                            emit_chunk(chunk);
                        }

                        // Always flush after each read so prompts and interactive output
                        // are delivered immediately without waiting for a next read.
                        // Under high-throughput bursts, `reader.read` returns ~buf-sized
                        // chunks (8 KB), keeping the IPC event rate reasonable while
                        // ensuring no data is ever held indefinitely when output quiets.
                        flush_pending(&mut pending, &emit_chunk);
                    }
                    Err(_) => {
                        flush_pending(&mut pending, &emit_chunk);
                        break;
                    }
                }
            }
        });

        Ok(session_id)
    }

    pub fn send_input(&self, session_id: &str, data: &str) -> anyhow::Result<()> {
        let writer = {
            let sessions = self
                .sessions
                .lock()
                .map_err(|_| anyhow::anyhow!("session lock poisoned"))?;
            let session = sessions
                .get(session_id)
                .ok_or_else(|| anyhow::anyhow!("unknown session"))?;
            session.writer.clone()
        };
        let mut writer = writer
            .lock()
            .map_err(|_| anyhow::anyhow!("writer lock poisoned"))?;
        writer.write_all(data.as_bytes())?;
        writer.flush()?;
        Ok(())
    }

    pub fn resize(&self, session_id: &str, cols: u16, rows: u16) -> anyhow::Result<()> {
        let master = {
            let sessions = self
                .sessions
                .lock()
                .map_err(|_| anyhow::anyhow!("session lock poisoned"))?;
            let session = sessions
                .get(session_id)
                .ok_or_else(|| anyhow::anyhow!("unknown session"))?;
            session.master.clone()
        };
        let master = master
            .lock()
            .map_err(|_| anyhow::anyhow!("master lock poisoned"))?;
        master
            .resize(PtySize {
                rows,
                cols,
                pixel_width: 0,
                pixel_height: 0,
            })
            .context("failed to resize pty")?;
        Ok(())
    }

    pub fn close(&self, session_id: &str) -> anyhow::Result<()> {
        let session = {
            let mut sessions = self
                .sessions
                .lock()
                .map_err(|_| anyhow::anyhow!("session lock poisoned"))?;
            sessions.remove(session_id)
        }
        .ok_or_else(|| anyhow::anyhow!("unknown session"))?;
        let mut child = session
            .child
            .lock()
            .map_err(|_| anyhow::anyhow!("child lock poisoned"))?;
        child.kill().context("failed to kill session")?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::{
        build_local_shell_command, build_ssh_command, expand_ssh_user_path, resolve_ssh_program,
        ssh_user_known_hosts_option,
    };
    use crate::ssh_config::HostConfig;

    #[test]
    fn builds_expected_ssh_command_with_proxy_and_identity() {
        let host = HostConfig {
            host: "prod".to_string(),
            host_name: "10.0.0.5".to_string(),
            user: "deploy".to_string(),
            port: 2201,
            identity_file: "~/.ssh/id_prod".to_string(),
            proxy_jump: "bastion".to_string(),
            proxy_command: String::new(),
        };

        let cmd = build_ssh_command(&host, "ask", 3);
        let rendered = format!("{cmd:?}");
        assert!(rendered.to_lowercase().contains("ssh"));
        assert!(rendered.contains("-p"));
        assert!(rendered.contains("2201"));
        assert!(rendered.contains("-i"));
        assert!(rendered.contains("id_prod"));
        assert!(rendered.contains("UserKnownHostsFile="));
        assert!(rendered.contains("StrictHostKeyChecking=ask"));
        assert!(rendered.contains("-J"));
        assert!(rendered.contains("bastion"));
        assert!(rendered.contains("deploy@10.0.0.5"));
        assert!(rendered.contains("ConnectTimeout=3"));
        assert!(rendered.contains("ConnectionAttempts=1"));
    }

    #[test]
    fn ssh_command_respects_strict_host_key_mode() {
        let host = HostConfig {
            host: "h".to_string(),
            host_name: "example.com".to_string(),
            user: "u".to_string(),
            port: 22,
            identity_file: String::new(),
            proxy_jump: String::new(),
            proxy_command: String::new(),
        };
        let cmd = build_ssh_command(&host, "accept-new", 3);
        let rendered = format!("{cmd:?}");
        assert!(rendered.contains("StrictHostKeyChecking=accept-new"));
    }

    // Separate tests per OS: CodeQL Rust extractor warns on `assert!` under mixed `cfg` branches.
    #[test]
    #[cfg(not(target_os = "windows"))]
    fn builds_local_shell_command_from_explicit_shell_posix() {
        let cmd = build_local_shell_command(Some("/usr/bin/fish"));
        let rendered = format!("{cmd:?}");
        assert!(rendered.contains("/usr/bin/fish"));
        assert!(rendered.contains("-l"));
    }

    #[test]
    #[cfg(target_os = "windows")]
    fn builds_local_shell_command_from_explicit_shell_windows() {
        let cmd = build_local_shell_command(Some("/usr/bin/fish"));
        let rendered = format!("{cmd:?}");
        assert!(rendered.contains("/usr/bin/fish"));
        assert!(!rendered.contains("-l"));
    }

    #[test]
    #[cfg(target_os = "windows")]
    fn builds_local_shell_command_default_uses_windows_shell() {
        let cmd = build_local_shell_command(None);
        let rendered = format!("{cmd:?}").to_lowercase();
        // On Windows the default must resolve to PowerShell or cmd.exe, never Unix paths.
        assert!(
            rendered.contains("powershell") || rendered.contains("cmd"),
            "expected PowerShell or cmd.exe as default shell on Windows, got: {rendered}"
        );
        // Login flag must not be passed to Windows shells.
        assert!(!rendered.contains("-l"));
    }

    #[test]
    fn expand_ssh_user_path_handles_blank_and_passthrough() {
        assert_eq!(expand_ssh_user_path(""), "");
        assert_eq!(expand_ssh_user_path("   "), "");
        // Absolute / non-tilde paths pass through unchanged.
        let p = "/etc/ssh/key";
        assert_eq!(expand_ssh_user_path(p), p);
        let q = "C:\\ssh\\key";
        assert_eq!(expand_ssh_user_path(q), q);
    }

    #[test]
    fn expand_ssh_user_path_expands_tilde_when_home_known() {
        if let Some(home) = home::home_dir() {
            let expanded = expand_ssh_user_path("~/.ssh/id_rsa");
            let expected = home.join(".ssh/id_rsa").to_string_lossy().into_owned();
            assert_eq!(expanded, expected);
        }
    }

    #[test]
    fn ssh_user_known_hosts_option_uses_user_known_hosts_file_form() {
        let opt = ssh_user_known_hosts_option();
        assert!(
            opt.starts_with("UserKnownHostsFile="),
            "unexpected option: {opt}"
        );
        assert!(
            opt.to_lowercase().ends_with("known_hosts"),
            "unexpected option suffix: {opt}"
        );
    }

    #[test]
    #[cfg(not(target_os = "windows"))]
    fn resolve_ssh_program_returns_unix_default() {
        assert_eq!(resolve_ssh_program(), "ssh");
    }

    #[test]
    #[cfg(target_os = "windows")]
    fn resolve_ssh_program_returns_a_visible_program_on_windows() {
        let p = resolve_ssh_program().to_lowercase();
        assert!(p.contains("ssh"), "unexpected ssh program path: {p}");
    }
}