rust-work/ssh-proxy/src/main.rs

mod cli;
mod config;
mod handler;
mod logger;
mod session;

use axum::{routing::{get, post}, Router};
use clap::{Parser, Subcommand};
use std::sync::Arc;

#[derive(Parser, Debug)]
#[command(name = "ssh-proxy")]
#[command(about = "SSH HTTP proxy with session pooling")]
struct Args {
    #[command(subcommand)]
    command: Option<Commands>,
    #[arg(long, default_value = "ssh-proxy.toml", global = true)]
    config: String,
    #[arg(short = 'S', long, default_value = "http://127.0.0.1:3308", global = true)]
    server: String,
}

#[derive(Debug, Subcommand)]
enum Commands {
    Server {
        #[arg(short = 'P', long)]
        port: Option<u16>,
        #[arg(short = 'H', long)]
        host: Option<String>,
    },
    Exec {
        #[arg(short = 'n', long)]
        name: String,
        #[arg(short, long)]
        command: String,
        /// Output format: text, json
        #[arg(short = 'F', long, default_value = "text")]
        format: String,
    },
    Servers,
    /// Add server dynamically (temporary)
    AddServer {
        #[arg(short = 'n', long)]
        name: String,
        #[arg(short = 'H', long)]
        host: String,
        #[arg(short = 'P', long, default_value = "22")]
        port: u16,
        #[arg(short = 'u', long, default_value = "root")]
        user: String,
        #[arg(short = 'p', long)]
        password: Option<String>,
        #[arg(short = 'k', long)]
        private_key: Option<String>,
    },
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let args = Args::parse();
    match args.command {
        Some(Commands::Server { port, host }) => run_server(&args.config, port, host).await,
        Some(Commands::Exec { name, command, format }) => run_cli_exec(&args.server, &name, &command, &format),
        Some(Commands::Servers) => run_cli_servers(&args.server),
        Some(Commands::AddServer { name, host, port, user, password, private_key }) => {
            run_cli_add_server(&args.server, name, host, port, user, password, private_key)
        }
        None => run_server(&args.config, None, None).await,
    }
}

async fn run_server(config_path: &str, port: Option<u16>, host: Option<String>) -> anyhow::Result<()> {
    println!("SSH HTTP Proxy v0.1.0\n");
    let mut config = config::Config::from_file(config_path)?;
    if let Some(p) = port { config.server.port = p; }
    if let Some(h) = host { config.server.host = h; }

    // 初始化日志
    let log_path = std::env::var("SSH_PROXY_LOG").ok();
    let logger = Arc::new(logger::RequestLogger::new(log_path.as_deref()));
    if logger.is_enabled() {
        println!("Request logging: enabled");
    }

    println!("Initializing SSH sessions...\n");
    let manager = Arc::new(session::SessionManager::new(&config.servers)?);

    let manager_clone = manager.clone();
    let idle_timeout = config.pool.idle_timeout_secs;
    let check_interval = config.pool.check_interval_secs;
    tokio::spawn(async move {
        let mut interval = tokio::time::interval(std::time::Duration::from_secs(check_interval));
        loop {
            interval.tick().await;
            manager_clone.cleanup_idle(idle_timeout).await;
        }
    });

    let app = Router::new()
        .route("/exec", post(handler::exec))
        .route("/servers", get(handler::servers))
        .route("/servers/add", post(handler::add_server))
        .route("/health", get(handler::health))
        .with_state(Arc::new((manager, logger)));

    let addr = format!("{}:{}", config.server.host, config.server.port);
    let listener = tokio::net::TcpListener::bind(&addr).await?;

    println!("\nServer started at http://{}", addr);
    println!("\nAPI Endpoints:");
    println!("  POST /exec         - Execute remote command");
    println!("  POST /servers/add  - Add server (temporary)");
    println!("  GET  /servers      - List all servers");
    println!("  GET  /health       - Health check");
    println!("\nCLI Usage:");
    println!("  ssh-proxy exec -n flux_dev -c \"docker ps\"");
    println!("  ssh-proxy exec -n flux_dev -c \"docker ps\" -F json");
    println!("  ssh-proxy servers");

    axum::serve(listener, app).await?;
    Ok(())
}

fn run_cli_exec(server_url: &str, server: &str, command: &str, format: &str) -> anyhow::Result<()> {
    let cli = cli::Cli::new(Some(server_url.to_string()));
    if !cli.check_server()? {
        eprintln!();
        eprintln!("=== ssh-proxy 未运行，请使用降级方案 ===");
        eprintln!();
        eprintln!("启动代理: ssh-proxy");
        eprintln!();
        eprintln!("或直接使用 ssh 命令:");
        eprintln!("  ssh <user>@<host> \"<command>\"");
        eprintln!();
        eprintln!("配置文件中的服务器信息见: ssh-proxy.toml");
        anyhow::bail!("Proxy server not running at {}", server_url);
    }
    cli.exec(server, command, format)
}

fn run_cli_servers(server_url: &str) -> anyhow::Result<()> {
    let cli = cli::Cli::new(Some(server_url.to_string()));
    if !cli.check_server()? {
        eprintln!();
        eprintln!("=== ssh-proxy 未运行 ===");
        eprintln!("启动代理: ssh-proxy");
        eprintln!("配置文件: ssh-proxy.toml");
        anyhow::bail!("Proxy server not running at {}", server_url);
    }
    cli.list_servers()
}

fn run_cli_add_server(
    server_url: &str,
    name: String,
    host: String,
    port: u16,
    user: String,
    password: Option<String>,
    private_key: Option<String>,
) -> anyhow::Result<()> {
    let cli = cli::Cli::new(Some(server_url.to_string()));
    if !cli.check_server()? {
        eprintln!();
        eprintln!("=== ssh-proxy 未运行 ===");
        eprintln!("启动代理: ssh-proxy");
        anyhow::bail!("Proxy server not running at {}", server_url);
    }
    cli.add_server(name, host, port, user, password, private_key)
}