window-sand-box 0.1.1

//! Windows 沙盒终端执行工具

use std::collections::HashMap;
use std::path::PathBuf;
use window_sand_box::runner::executor::SandboxRunner;

use clap::{Parser, Subcommand, CommandFactory};

#[derive(Parser)]
#[command(name = "wsbx", version, about = "Windows 沙盒终端执行工具")]
#[command(long_about = "Windows 沙盒终端执行工具。
为 AI 提供受限制的命令执行环境。

说明:
  外部命令（.exe）直接执行:  wsbx exec git status
  CMD 内部命令加 cmd /c:     wsbx exec cmd /c \"mkdir test && dir\"
  .wsbx\\ 目录和配置在首次运行时自动创建
  ACL 在工作目录首次写入时自动配（需管理员弹窗确认）
  'wsbx clean' 可清除自动添加的 ACL 条目")]
struct Cli {
    #[command(subcommand)]
    command: Option<Commands>,
}

#[derive(Subcommand)]
enum Commands {
    /// 在沙盒中执行命令
    Exec {
        /// 工作目录（默认：当前目录）
        #[arg(long)]
        cwd: Option<PathBuf>,

        /// 只读模式（无写入权限）
        #[arg(long)]
        readonly: bool,

        /// 超时时间（秒）
        #[arg(long)]
        timeout: Option<u64>,

        /// 额外可写入的路径（可多次使用）
        #[arg(long = "whitelist", action = clap::ArgAction::Append)]
        whitelist: Option<Vec<PathBuf>>,

        /// 禁用桌面隔离（默认启用，沙盒进程在私有桌面中运行）
        #[arg(long)]
        no_desktop: bool,

        /// 要执行的命令及参数
        #[arg(required = true, trailing_var_arg = true, allow_hyphen_values = true)]
        command: Vec<String>,
    },
    /// 清理当前目录的 ACL 条目（需管理员）
    Clean,
    /// 内部子命令：由 elevated 进程自动调用以配置 ACL
    #[clap(hide = true)]
    InternalSetup,
}

fn main() {
    let exit_code = match run() {
        Ok(code) => code,
        Err(e) => {
            eprintln!("错误: {}", e);
            1
        }
    };
    std::process::exit(exit_code);
}

fn run() -> anyhow::Result<i32> {
    let cli = Cli::parse();

    let Some(command) = cli.command else {
        // 无子命令时自动打印帮助
        let mut cmd = Cli::command();
        cmd.print_help()?;
        println!();
        return Ok(0);
    };

    match command {
        Commands::Exec {
            cwd,
            readonly,
            timeout,
            whitelist,
            no_desktop,
            command,
        } => cmd_exec(ExecArgs {
            cwd,
            readonly,
            timeout,
            whitelist,
            no_desktop,
            command,
        }),
        Commands::InternalSetup => {
            internal_setup(&std::env::current_dir()?)?;
            Ok(0)
        }
        Commands::Clean => {
            cmd_clean()?;
            Ok(0)
        }
    }
}

// ==================== 数据结构 ====================

/// `wsbx exec` 解析后的参数
struct ExecArgs {
    cwd: Option<PathBuf>,
    readonly: bool,
    timeout: Option<u64>,
    whitelist: Option<Vec<PathBuf>>,
    no_desktop: bool,
    command: Vec<String>,
}

// ==================== 子命令实现 ====================

/// 内部：配置 ACL（由 elevated 进程调用，不对外暴露）
fn internal_setup(cwd: &std::path::Path) -> anyhow::Result<()> {
    let runner = SandboxRunner::new()?;
    let session = runner.create_session(cwd.to_path_buf())?;
    runner.setup_session_acls(&session)?;
    println!("✓ ACL 配置完成");
    Ok(())
}

/// 清理 ACL
///
/// 移除 `wsbx setup` 添加的所有 ACE 条目。
/// 先尝试直接清理，如果失败且非管理员，自动触发 UAC 提权重试。
fn cmd_clean() -> anyhow::Result<()> {
    let runner = SandboxRunner::new()?;
    let cwd = std::env::current_dir()?;

    let session = runner.create_session(cwd)?;
    println!("正在清理 '{}' 的沙盒 ACL...", session.work_dir.display());

    match runner.clean_session_acls(&session) {
        Ok(_) => {
            println!("✓ ACL 清理完成");
            println!("  工作目录: {}", session.work_dir.display());
            Ok(())
        }
        Err(e) => {
            if is_running_as_admin() {
                eprintln!("✗ ACL 清理失败: {}", e);
                return Err(e);
            }
            println!("⚠ ACL 清理失败: {}", e);
            println!("🔑 正在请求管理员权限...（请在弹出的 UAC 对话框中点击「是」）");
            elevate_and_rerun("clean")
        }
    }
}

/// 在沙盒中执行命令（clap 重构版）
fn cmd_exec(args: ExecArgs) -> anyhow::Result<i32> {
    if args.command.is_empty() {
        eprintln!("错误: exec 需要指定命令");
        return Ok(1);
    }

    // 解析工作目录
    let cwd = args.cwd.unwrap_or_else(|| std::env::current_dir().expect("获取当前目录失败"));

    // 确保工作目录存在
    if !cwd.exists() {
        eprintln!("错误: 工作目录 '{}' 不存在", cwd.display());
        return Ok(1);
    }

    // 创建执行器（提前创建，以便使用 config 检查黑名单）
    let mut runner = SandboxRunner::new()?;

    // 桌面隔离配置（默认启用私有桌面）
    runner.set_no_desktop(args.no_desktop);

    // 检查工作目录是否在黑名单中
    if runner.config().is_blacklisted(&cwd) {
        eprintln!(
            "错误: 工作目录 '{}' 在黑名单中，不允许作为沙盒工作目录",
            cwd.display()
        );
        return Ok(1);
    }

    // 添加 --whitelist 指定的临时白名单路径
    if let Some(ref whitelist) = args.whitelist {
        for p in whitelist {
            // 检查白名单路径是否在黑名单中（白名单不能覆盖黑名单）
            if runner.config().is_blacklisted(p) {
                eprintln!("⚠ 白名单路径 '{}' 在黑名单中，已忽略", p.display());
                continue;
            }
            runner.add_temp_whitelist(p);
        }
    }

    let timeout_ms = args.timeout.map(|s| s * 1000);
    let extra_env = HashMap::new();

    let result = if args.readonly {
        println!("[沙盒:只读] 执行命令: {}", args.command.join(" "));
        println!("[工作目录] {}", cwd.display());
        runner.execute_readonly(&args.command, &cwd, extra_env, timeout_ms)?
    } else {
        let mut session = runner.create_session(cwd.clone())?;
        println!("[沙盒] 执行命令: {}", args.command.join(" "));
        println!("[工作目录] {}", session.work_dir.display());
        println!("[会话 ID] {}", session.session_id);

        // 尝试设置 ACL；未配置时自动提权配置，而不是静默失败
        let acl_ok = if let Err(e) = runner.setup_session_acls(&session) {
            if !is_running_as_admin() {
                println!("⚠ 工作目录未配置沙盒 ACL: {}", e);
                println!("🔑 正在请求管理员权限以自动配置...");
                // 提权运行 setup，在目标工作目录下配置 ACL
                elevate_and_rerun_in("internal-setup", &cwd)?;
                // 提权完成后重新创建 session，确保 SID 与 elevated 进程一致。
                //
                // elevated 进程调用 setup 时加载了持久化的 CapSidStore，可能已为
                // 工作目录生成 SID。重新创建 session 读取 store 中的同一 SID，
                // 避免 TOCTOU 导致的 SID 不匹配。
                session = runner.create_session(cwd.clone())?;
                println!("✓ ACL 配置完成，继续执行命令...");
                true
            } else {
                eprintln!("警告: ACL 设置失败: {}", e);
                false
            }
        } else {
            true
        };

        // 如果 ACL 配置失败（已是管理员但仍失败），发出警告但继续执行
        if !acl_ok {
            eprintln!("⚠ ACL 未正确配置，沙盒隔离可能不完整");
        }

        runner.execute(&session, &args.command, extra_env, timeout_ms)?
    };

    // 输出结果
    if !result.stdout.is_empty() {
        print!("{}", result.stdout);
    }
    if !result.stderr.is_empty() {
        eprint!("{}", result.stderr);
    }

    println!("[退出码] {}", result.exit_code);
    if result.timed_out {
        eprintln!("[超时] 命令执行超时");
    }

    Ok(result.exit_code)
}

// ==================== UAC 提权 ====================

/// 使用 ShellExecuteExW 以管理员权限重新运行 `wsbx <subcommand>`（触发 UAC 提权）
fn elevate_and_rerun(subcommand: &str) -> anyhow::Result<()> {
    elevate_and_rerun_in(subcommand, &std::env::current_dir()?)
}

/// 在指定目录下以管理员权限重新运行 `wsbx <subcommand>`（触发 UAC 提权）
fn elevate_and_rerun_in(subcommand: &str, cwd: &std::path::Path) -> anyhow::Result<()> {
    use windows_sys::Win32::Foundation::{CloseHandle, GetLastError, HWND};
    use windows_sys::Win32::System::Threading::WaitForSingleObject;
    use windows_sys::Win32::UI::Shell::{ShellExecuteExW, SHELLEXECUTEINFOW, SEE_MASK_NOCLOSEPROCESS};
    use windows_sys::Win32::UI::WindowsAndMessaging::SW_SHOWNORMAL;

    let exe_path = std::env::current_exe()?;

    let verb = window_sand_box::winutil::to_wide("runas");
    let file = window_sand_box::winutil::to_wide(exe_path.as_os_str());
    let params = window_sand_box::winutil::to_wide(subcommand);
    let dir = window_sand_box::winutil::to_wide(cwd.as_os_str());

    let mut info: SHELLEXECUTEINFOW = unsafe { std::mem::zeroed() };
    info.cbSize = std::mem::size_of::<SHELLEXECUTEINFOW>() as u32;
    info.fMask = SEE_MASK_NOCLOSEPROCESS;
    info.hwnd = 0 as HWND;
    info.lpVerb = verb.as_ptr();
    info.lpFile = file.as_ptr();
    info.lpParameters = params.as_ptr();
    info.lpDirectory = dir.as_ptr();
    info.nShow = SW_SHOWNORMAL;

    let ok = unsafe { ShellExecuteExW(&mut info) };
    if ok == 0 {
        let err = unsafe { GetLastError() };
        return Err(anyhow::anyhow!(
            "提权失败 (错误码: {})。请手动以管理员身份运行:\n  cd /d \"{}\" && {} {}",
            err,
            cwd.display(),
            exe_path.display(),
            subcommand
        ));
    }

    // 等待 elevated 进程完成（最多 5 分钟，防止 UAC 弹窗不响应导致永久阻塞）
    const UAC_TIMEOUT_MS: u32 = 300_000;
    let h_process = info.hProcess;
    if h_process != 0 {
        unsafe {
            let wait_result = WaitForSingleObject(h_process, UAC_TIMEOUT_MS);
            if wait_result == 0x0000_0102 {
                eprintln!("⚠ 管理员操作超时（5分钟），进程仍在后台运行");
                eprintln!("  请手动检查并关闭 elevated 进程");
            }
            CloseHandle(h_process);
        }
    }

    println!("✓ 管理员操作完成");
    Ok(())
}

/// 检查当前进程是否以管理员权限运行
fn is_running_as_admin() -> bool {
    unsafe {
        use windows_sys::Win32::Security::{
            AllocateAndInitializeSid, CheckTokenMembership, FreeSid,
            SECURITY_NT_AUTHORITY,
        };

        let mut admin_sid: *mut std::ffi::c_void = std::ptr::null_mut();
        let mut is_member: i32 = 0;

        // SECURITY_BUILTIN_DOMAIN_RID = 0x20, DOMAIN_ALIAS_RID_ADMINS = 0x220
        let ok = AllocateAndInitializeSid(
            &SECURITY_NT_AUTHORITY,
            2,
            0x20,    // SECURITY_BUILTIN_DOMAIN_RID
            0x220,   // DOMAIN_ALIAS_RID_ADMINS
            0, 0, 0, 0, 0, 0,
            &mut admin_sid,
        );
        if ok == 0 {
            return false;
        }

        let ok2 = CheckTokenMembership(0, admin_sid, &mut is_member);
        FreeSid(admin_sid);

        ok2 != 0 && is_member != 0
    }
}