PicoBot/src/platform/mod.rs

//! Platform abstraction layer for cross-platform compatibility.
//!
//! This module provides unified interfaces for platform-specific operations,
//! making it easy to add support for new platforms by modifying only this file.

use std::env;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};

/// Supported platform types.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Platform {
    Windows,
    Unix,
}

impl Platform {
    /// Detect the current platform.
    pub fn current() -> Self {
        if cfg!(target_os = "windows") {
            Platform::Windows
        } else {
            Platform::Unix
        }
    }

    /// Check if running on Windows.
    pub fn is_windows() -> bool {
        cfg!(target_os = "windows")
    }
}

/// Shell information for command execution.
#[derive(Debug, Clone)]
pub struct ShellInfo {
    /// Tool name exposed to LLM.
    pub name: &'static str,
    /// Shell executable name.
    pub executable: &'static str,
    /// Arguments to pass before the command.
    pub args: &'static [&'static str],
}

impl ShellInfo {
    /// Get the default shell for the current platform.
    pub fn default() -> Self {
        Self::for_platform(Platform::current())
    }

    /// Get shell info for a specific platform.
    pub fn for_platform(platform: Platform) -> Self {
        match platform {
            Platform::Windows => ShellInfo {
                name: "shell",
                executable: "powershell",
                args: &["-Command"],
            },
            Platform::Unix => ShellInfo {
                name: "bash",
                executable: "bash",
                args: &["-c"],
            },
        }
    }

    /// Alternative shells available on the platform.
    pub fn available_shells(platform: Platform) -> Vec<ShellInfo> {
        match platform {
            Platform::Windows => vec![
                ShellInfo {
                    name: "shell",
                    executable: "powershell",
                    args: &["-Command"],
                },
                ShellInfo {
                    name: "shell",
                    executable: "cmd",
                    args: &["/C"],
                },
            ],
            Platform::Unix => vec![
                ShellInfo {
                    name: "bash",
                    executable: "bash",
                    args: &["-c"],
                },
                // Future: could add zsh, fish, sh
                // ShellInfo { name: "zsh", executable: "zsh", args: &["-c"] },
            ],
        }
    }
}

/// Dangerous command patterns for safety guards.
///
/// Returns patterns filtered by the current platform. Platform-specific
/// rules (e.g. `format` on Windows, `rm` on Unix) are only injected on
/// their target platform to avoid false positives.
pub fn dangerous_command_patterns() -> Vec<String> {
    dangerous_command_patterns_for_platform(Platform::current())
}

/// Platform-specific dangerous command patterns.
///
/// Exposed primarily for testing. Callers should prefer
/// [`dangerous_command_patterns`] which auto-detects the platform.
pub fn dangerous_command_patterns_for_platform(platform: Platform) -> Vec<String> {
    let mut patterns: Vec<String> = Vec::new();

    // Cross-platform: fork bomb
    patterns.push(r":\(\)\s*\{.*\};\s*:".to_string());

    match platform {
        Platform::Unix => {
            // Unix dangerous commands
            patterns.push(r"\brm\s+-[rf]{1,2}\b".to_string());
            patterns.push(r"\bchmod\s+-[Rr]".to_string());
            patterns.push(r"\bchown\s+-[Rr]".to_string());
        }
        Platform::Windows => {
            // Windows cmd dangerous commands.
            // `format` requires a drive letter (`[a-z]:`) somewhere after it,
            // so legitimate uses like `dart format lib/` or
            // `pytest --format json` (no drive letter) are not matched.
            patterns.push(r"\bformat\s+.*[a-z]:".to_string());
            patterns.push(r"\bdel\s+/[fq]\b".to_string());
            patterns.push(r"\brmdir\s+/s\b".to_string());
            // PowerShell dangerous commands. Patterns are lowercase because
            // `guard_command` lowercases the command string before matching.
            patterns.push(r"\bremove-item\s+.*-recurse".to_string());
            patterns.push(r"\bremove-item\s+.*-force".to_string());
        }
    }

    patterns
}

/// Check whether a child process is blocked waiting for stdin input.
///
/// Uses platform-specific mechanisms to determine if the process is genuinely
/// waiting for user input (as opposed to computing, sleeping, or doing I/O).
///
/// Returns `None` when the platform does not support this check or the
/// information cannot be read.
pub fn is_process_waiting_on_stdin(pid: u32) -> Option<bool> {
    #[cfg(target_os = "linux")]
    {
        let wchan = std::fs::read_to_string(format!("/proc/{}/wchan", pid)).ok()?;
        let wchan = wchan.trim();
        if wchan.is_empty() {
            return None;
        }
        Some(wchan.contains("tty_read") || wchan.contains("n_tty_read") || wchan == "pipe_wait")
    }
    #[cfg(target_os = "macos")]
    {
        use std::mem;
        let mut task_info: libc::proc_taskinfo = unsafe { mem::zeroed() };
        let size = mem::size_of::<libc::proc_taskinfo>() as i32;
        let ret = unsafe {
            libc::proc_pidinfo(
                pid as i32,
                libc::PROC_PIDTASKINFO,
                0,
                &mut task_info as *mut _ as *mut libc::c_void,
                size,
            )
        };
        if ret <= 0 {
            return None;
        }
        // pti_numrunning == 0 means no thread is actively on CPU.
        // Combined with output silence this strongly suggests the process
        // is blocked on I/O (likely a stdin read).
        Some(task_info.pti_numrunning == 0)
    }
    #[cfg(target_os = "windows")]
    {
        windows_is_process_waiting_on_stdin(pid)
    }
    #[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
    {
        let _ = pid;
        None
    }
}

/// Windows implementation: check if a process is waiting for stdin input.
///
/// Uses NtQuerySystemInformation to enumerate process threads and check if
/// all threads are in Wait state with Executive wait reason, which indicates
/// the process is blocked on I/O (likely console input).
#[cfg(target_os = "windows")]
fn windows_is_process_waiting_on_stdin(pid: u32) -> Option<bool> {
    // SystemProcessInformation = 5
    const SYSTEM_PROCESS_INFORMATION: u32 = 5;
    const STATUS_INFO_LENGTH_MISMATCH: i32 = -1073741820; // 0xC0000004

    #[repr(C)]
    #[allow(non_snake_case)]
    struct SystemProcessInfo {
        next_entry_offset: u32,
        number_of_threads: u32,
        working_set_private_size: i64,
        hard_fault_count: u32,
        number_of_threads_high_watermark: u32,
        cycle_time: u64,
        create_time: i64,
        user_time: i64,
        kernel_time: i64,
        image_name_length: u16,
        image_name_max_length: u16,
        image_name: *const u16,
        base_priority: i32,
        unique_process_id: *mut std::ffi::c_void,
        inherited_from_unique_process_id: *mut std::ffi::c_void,
        handle_count: u32,
        session_id: u32,
        unique_process_key: usize,
        peak_virtual_size: usize,
        virtual_size: usize,
        page_fault_count: u32,
        peak_working_set_size: usize,
        working_set_size: usize,
        quota_peak_paged_pool_usage: usize,
        quota_paged_pool_usage: usize,
        quota_peak_non_paged_pool_usage: usize,
        quota_non_paged_pool_usage: usize,
        pagefile_usage: usize,
        peak_pagefile_usage: usize,
        private_page_count: usize,
        read_operation_count: i64,
        write_operation_count: i64,
        other_operation_count: i64,
        read_transfer_count: i64,
        write_transfer_count: i64,
        other_transfer_count: i64,
        // SYSTEM_THREAD_INFORMATION[1] follows in memory
        threads: [SystemThreadInfo; 1],
    }

    #[repr(C)]
    #[derive(Clone, Copy)]
    #[allow(non_snake_case)]
    struct SystemThreadInfo {
        kernel_time: i64,
        user_time: i64,
        create_time: i64,
        wait_time: u32,
        start_address: *mut std::ffi::c_void,
        client_id_unique_process: *mut std::ffi::c_void,
        client_id_unique_thread: *mut std::ffi::c_void,
        priority: i32,
        base_priority: i32,
        context_switches: u32,
        thread_state: u32,
        wait_reason: u32,
    }

    #[allow(non_snake_case)]
    unsafe extern "system" {
        fn NtQuerySystemInformation(
            system_information_class: u32,
            system_information: *mut u8,
            system_information_length: u32,
            return_length: *mut u32,
        ) -> i32;
    }

    unsafe {
        // Query buffer size first
        let mut buf_len: u32 = 0;
        let status = NtQuerySystemInformation(
            SYSTEM_PROCESS_INFORMATION,
            std::ptr::null_mut(),
            0,
            &mut buf_len,
        );
        if status != STATUS_INFO_LENGTH_MISMATCH || buf_len == 0 {
            return None;
        }

        // Allocate buffer with extra space (processes may be created between calls)
        buf_len = buf_len.saturating_mul(2).max(65536);
        let mut buffer: Vec<u8> = vec![0u8; buf_len as usize];

        let status = NtQuerySystemInformation(
            SYSTEM_PROCESS_INFORMATION,
            buffer.as_mut_ptr(),
            buf_len,
            &mut buf_len,
        );
        if status < 0 {
            return None;
        }

        // Walk the linked list of SYSTEM_PROCESS_INFORMATION
        let mut offset: usize = 0;
        loop {
            let info = &*(buffer.as_ptr().add(offset) as *const SystemProcessInfo);
            let proc_id = info.unique_process_id as u32;

            if proc_id == pid {
                let thread_count = info.number_of_threads as usize;
                if thread_count == 0 {
                    return Some(false);
                }

                // Thread states: Running=2, Waiting=5
                // Wait reasons: Executive=0, FreePage=1, PageIn=2, PoolAllocation=3,
                //               DelayExecution=4, Suspended=5, UserRequest=6, ...
                // Executive wait + all threads waiting = likely blocked on I/O
                let all_waiting = (0..thread_count).all(|i| {
                    let thread = &*info.threads.as_ptr().add(i);
                    thread.thread_state == 5 && thread.wait_reason == 0
                });

                return Some(all_waiting);
            }

            if info.next_entry_offset == 0 {
                break;
            }
            offset += info.next_entry_offset as usize;
        }

        None
    }
}

/// Get the user's home directory.
///
/// Supports environment variable overrides for testing:
/// - `HOME` (Unix-style, works on all platforms for testing)
/// - `USERPROFILE` (Windows-specific)
pub fn home_dir() -> Option<PathBuf> {
    // Test scenario: support HOME variable override on all platforms
    env::var_os("HOME")
        .map(PathBuf::from)
        .or_else(|| {
            // Windows: support USERPROFILE
            env::var_os("USERPROFILE").map(PathBuf::from)
        })
        .or_else(dirs::home_dir)
}

/// 返回 PicoBot 主目录，若无法确定则回退到当前目录 `"."`。
///
/// 包装 [`home_dir`] 并内置 fallback，消除各模块重复的
/// `dirs::home_dir().unwrap_or_else(|| PathBuf::from("."))` 模式。
pub fn picobot_home_dir() -> PathBuf {
    home_dir().unwrap_or_else(|| PathBuf::from("."))
}

/// Atomically rename a file, handling platform differences.
///
/// On Windows, `fs::rename` fails if the destination exists, so we need to
/// remove it first. On Unix, rename is atomic and replaces the destination.
pub fn atomic_rename(src: &Path, dst: &Path) -> io::Result<()> {
    if Platform::is_windows() && dst.exists() {
        fs::remove_file(dst)?;
    }
    fs::rename(src, dst)
}

/// Convert a filesystem path to a file:// URI.
///
/// Handles platform-specific path formats:
/// - Unix: `/path/to/file` -> `file:///path/to/file`
/// - Windows: `C:\path\to\file` -> `file:///C:/path/to/file`
pub fn path_to_uri(path: &Path) -> String {
    let path_str = path.display().to_string();
    if Platform::is_windows() {
        // Windows paths use backslashes which must be converted to forward slashes
        let normalized = path_str.replace('\\', "/");
        format!("file:///{}", normalized)
    } else {
        format!("file://{}", path_str)
    }
}

/// XML escape utility.
pub fn xml_escape(value: &str) -> String {
    value
        .replace('&', "&amp;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
}

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

    #[test]
    fn test_platform_detect() {
        let platform = Platform::current();
        if cfg!(target_os = "windows") {
            assert_eq!(platform, Platform::Windows);
        } else {
            assert_eq!(platform, Platform::Unix);
        }
    }

    #[test]
    fn test_shell_info_default() {
        let shell = ShellInfo::default();
        if cfg!(target_os = "windows") {
            assert_eq!(shell.executable, "powershell");
            assert_eq!(shell.args, &["-Command"]);
        } else {
            assert_eq!(shell.executable, "bash");
            assert_eq!(shell.args, &["-c"]);
        }
    }

    #[test]
    fn test_shell_info_for_platform() {
        let win_shell = ShellInfo::for_platform(Platform::Windows);
        assert_eq!(win_shell.executable, "powershell");

        let unix_shell = ShellInfo::for_platform(Platform::Unix);
        assert_eq!(unix_shell.executable, "bash");
    }

    #[test]
    fn test_path_to_uri() {
        let temp_dir = tempfile::tempdir().unwrap();
        let test_path = temp_dir.path().join("test.txt");
        let uri = path_to_uri(&test_path);

        assert!(uri.starts_with("file://"));
        assert!(uri.contains("test.txt"));
        assert!(!uri.contains('\\')); // No backslashes
    }

    #[test]
    fn test_path_to_uri_windows_format() {
        if cfg!(target_os = "windows") {
            let win_path = PathBuf::from("C:\\Users\\test\\file.txt");
            let uri = path_to_uri(&win_path);
            assert!(uri.starts_with("file:///C:/"));
            assert_eq!(uri, "file:///C:/Users/test/file.txt");
        }
    }

    #[test]
    fn test_atomic_rename() {
        let temp_dir = tempfile::tempdir().unwrap();
        let src = temp_dir.path().join("source.txt");
        let dst = temp_dir.path().join("dest.txt");

        fs::write(&src, "content").unwrap();
        fs::write(&dst, "old content").unwrap();

        atomic_rename(&src, &dst).unwrap();

        assert!(!src.exists());
        assert!(dst.exists());
        assert_eq!(fs::read_to_string(&dst).unwrap(), "content");
    }

    #[test]
    fn test_dangerous_patterns() {
        let patterns = dangerous_command_patterns();
        assert!(!patterns.is_empty());
        // Cross-platform fork bomb rule is always present
        assert!(patterns.iter().any(|p| p.contains(r":\(\)")));
    }

    #[test]
    fn test_dangerous_patterns_unix() {
        let patterns = dangerous_command_patterns_for_platform(Platform::Unix);
        // Unix-specific rules
        assert!(patterns.iter().any(|p| p.contains("rm")));
        assert!(patterns.iter().any(|p| p.contains("chmod")));
        assert!(patterns.iter().any(|p| p.contains("chown")));
        // Windows-specific rules must NOT be present on Unix
        assert!(!patterns.iter().any(|p| p.contains("format")));
        assert!(!patterns.iter().any(|p| p.contains("del")));
        assert!(!patterns.iter().any(|p| p.contains("remove-item")));
    }

    #[test]
    fn test_dangerous_patterns_windows() {
        let patterns = dangerous_command_patterns_for_platform(Platform::Windows);
        // Windows-specific rules
        assert!(patterns.iter().any(|p| p.contains("del")));
        assert!(patterns.iter().any(|p| p.contains("format")));
        assert!(patterns.iter().any(|p| p.contains("remove-item")));
        // Unix-specific rules must NOT be present on Windows.
        // Use r"\brm\s" to match rm-as-command without matching rmdir.
        assert!(!patterns.iter().any(|p| p.contains(r"\brm\s")));
        assert!(!patterns.iter().any(|p| p.contains("chmod")));
        assert!(!patterns.iter().any(|p| p.contains("chown")));
    }

    #[test]
    fn test_format_pattern_precision() {
        let patterns = dangerous_command_patterns_for_platform(Platform::Windows);
        let format_pat = patterns
            .iter()
            .find(|p| p.contains("format"))
            .expect("format pattern should exist on Windows");
        let re = regex::Regex::new(format_pat).unwrap();

        // Helper: guard_command lowercases before matching, so tests must too.
        let m = |cmd: &str| re.is_match(&cmd.to_lowercase());

        // Truly dangerous commands — should match
        assert!(m("format c:"));
        assert!(m("format /q c:"));
        assert!(m("format d: /fs:ntfs"));
        assert!(m("echo ok | format c:"));
        assert!(m("echo ok; format c:"));
        // Sub-shell invocation must also be caught
        assert!(m(r#"cmd /c "format C:""#));

        // Legitimate commands containing literal "format " — should NOT match
        assert!(!m("dart format lib/"));
        assert!(!m("buf format -w"));
        assert!(!m("pytest --format json"));
        assert!(!m(r#"echo "please format the disk""#));
        assert!(!m(r#"git log --pretty=format:"%h""#));
    }

    #[test]
    fn test_xml_escape() {
        assert_eq!(xml_escape("a & b"), "a &amp; b");
        assert_eq!(xml_escape("<tag>"), "&lt;tag&gt;");
    }
}