rsh/crates/rshc/src/main.rs

mod auth;
mod cmd;
mod config;
mod ui;

use anyhow::Result;
use auth::AuthedClient;
use clap::{Args, Parser, Subcommand};
use config::Config;

#[derive(Parser, Debug)]
#[command(name = "rshc", version, about = "rsh operator client")]
struct Cli {
    #[command(subcommand)]
    cmd: Cmd,
}

#[derive(Subcommand, Debug)]
enum Cmd {
    Watch(WatchArgs),
    #[command(alias = "sessions", alias = "s", alias = "sess")]
    Session(SessionCmd),
    #[command(alias = "connections", alias = "conn")]
    Connection(ConnectionCmd),
    #[command(alias = "c")]
    Connect(ConnectArgs),
    #[command(alias = "sh")]
    Shell(ShellArgs),
    ServeSsh(ServeSshArgs),
    Keys(KeysCmd),
}

#[derive(Args, Debug)]
struct WatchArgs {
    session: Option<String>,
}

#[derive(Args, Debug)]
struct SessionCmd {
    #[command(subcommand)]
    sub: SessionSub,
}

#[derive(Subcommand, Debug)]
enum SessionSub {
    #[command(alias = "c")]
    Create { name: String },
    #[command(alias = "del", alias = "d", alias = "rm")]
    Delete {
        name: String,
        #[arg(short = 'y', long)]
        yes: bool,
        #[arg(long)]
        disconnect: bool,
    },
    Update {
        name: String,
        #[arg(long = "pw", num_args = 0..=1, default_missing_value = "")]
        pw: Option<String>,
        #[arg(long)]
        disconnect: bool,
    },
    #[command(alias = "l", alias = "ls")]
    List,
}

#[derive(Args, Debug)]
struct ConnectionCmd {
    #[command(subcommand)]
    sub: ConnectionSub,
}

#[derive(Subcommand, Debug)]
enum ConnectionSub {
    #[command(alias = "l", alias = "ls")]
    List {
        #[arg(long)]
        session: Option<String>,
    },
}

#[derive(Args, Debug)]
struct ConnectArgs {
    session: String,
    connection_id: Option<u64>,
    #[arg(long)]
    no_pty: bool,
}

#[derive(Args, Debug)]
struct ShellArgs {
    session: String,
    #[arg(long)]
    connection: Option<u64>,
    #[arg(long)]
    shell: Option<String>,
    #[arg(long)]
    no_pty: bool,
}

#[derive(Args, Debug)]
struct ServeSshArgs {
    session: String,
    #[arg(short = 'H', long, default_value = "127.0.0.1")]
    listen_host: String,
    #[arg(short = 'p', long, default_value_t = 2222)]
    listen_port: u16,
    #[arg(long)]
    connection: Option<u64>,
    #[arg(long)]
    shell: Option<String>,
}


#[derive(Args, Debug)]
struct KeysCmd {
    #[command(subcommand)]
    sub: KeysSub,
}

#[derive(Subcommand, Debug)]
enum KeysSub {
    Append {
        key: Option<String>,
        #[arg(long)]
        file: Option<String>,
        #[arg(long)]
        url: Option<String>,
    },
    Rm {
        key: Option<String>,
        #[arg(long)]
        file: Option<String>,
        #[arg(long)]
        url: Option<String>,
    },
    #[command(alias = "l", alias = "ls")]
    List,
    Edit,
}

#[tokio::main]
async fn main() {
    let _ = tracing_subscriber::fmt()
        .with_writer(std::io::stderr)
        .with_env_filter(tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn")))
        .try_init();
    if let Err(e) = run().await {
        ui::print_err(&e);
        std::process::exit(1);
    }
}

async fn run() -> Result<()> {
    let cli = Cli::parse();
    let cfg = Config::load()?;
    if let Cmd::Watch(a) = cli.cmd {
        return cmd::watch::run_with_reconnect(&cfg, a.session).await;
    }
    let client = AuthedClient::connect(&cfg).await?;
    match cli.cmd {
        Cmd::Watch(_) => unreachable!(),
        Cmd::Session(s) => match s.sub {
            SessionSub::Create { name } => cmd::session::create(&client, name).await,
            SessionSub::Delete { name, yes, disconnect } => cmd::session::delete(&client, name, yes, disconnect).await,
            SessionSub::Update { name, pw, disconnect } => {
                let pw_flag = match pw {
                    None => None,
                    Some(s) if s.is_empty() => Some(None),
                    Some(s) => Some(Some(s)),
                };
                cmd::session::update(&client, name, pw_flag, disconnect).await
            }
            SessionSub::List => cmd::session::list(&client).await,
        },
        Cmd::Connection(c) => match c.sub {
            ConnectionSub::List { session } => cmd::connection::list(&client, session).await,
        },
        Cmd::Connect(a) => cmd::connect::run(&client, a.session, a.connection_id, a.no_pty).await,
        Cmd::Shell(a) => cmd::shell::run(&client, a.session, a.connection, a.shell, a.no_pty).await,
        Cmd::ServeSsh(a) => cmd::serve_ssh::run(cfg, client, a.session, a.listen_host, a.listen_port, a.connection, a.shell).await,

        Cmd::Keys(k) => match k.sub {
            KeysSub::Append { key, file, url } => cmd::keys::append(&client, key, file, url).await,
            KeysSub::Rm { key, file, url } => cmd::keys::remove(&client, key, file, url).await,
            KeysSub::List => cmd::keys::list(&client).await,
            KeysSub::Edit => cmd::keys::edit(&client).await,
        },
    }
}