diff --git a/Cargo.lock b/Cargo.lock index fbbc040..ffde7de 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -493,6 +493,7 @@ dependencies = [ "bytes", "chrono", "clap", + "clap_complete", "colored", "console-subscriber", "dashmap", @@ -722,6 +723,15 @@ dependencies = [ "strsim", ] +[[package]] +name = "clap_complete" +version = "4.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3ae8ba90b9d8b007efe66e55e48fb936272f5ca00349b5b0e89877520d35ea7" +dependencies = [ + "clap", +] + [[package]] name = "clap_derive" version = "4.4.2" diff --git a/boltconn/Cargo.toml b/boltconn/Cargo.toml index 9ce9138..72cac13 100644 --- a/boltconn/Cargo.toml +++ b/boltconn/Cargo.toml @@ -7,6 +7,7 @@ readme = "README.md" [features] tokio-console = ["dep:console-subscriber", "tokio/tracing"] +internal-test = [] [dependencies] # Core @@ -76,6 +77,7 @@ shadowsocks = { version = "1.16.0", default-features = false } smoltcp = "0.9.0" # Command line clap = { version = "4.4.6", features = ["derive"] } +clap_complete = "4.4.3" colored = "2.0.0" tabular = "0.2.0" diff --git a/boltconn/src/cli/mod.rs b/boltconn/src/cli/mod.rs index 8da13f8..2b68a0b 100644 --- a/boltconn/src/cli/mod.rs +++ b/boltconn/src/cli/mod.rs @@ -6,7 +6,7 @@ mod request_web; use crate::ProgramArgs; use anyhow::anyhow; -use clap::{Args, Subcommand}; +use clap::{Args, CommandFactory, Subcommand, ValueHint}; use is_root::is_root; use std::path::PathBuf; use std::process::exit; @@ -14,7 +14,12 @@ use std::process::exit; #[derive(Debug, Subcommand)] pub(crate) enum ProxyOptions { /// Set group's proxy - Set { group: String, proxy: String }, + Set { + #[clap(value_hint = ValueHint::Other)] + group: String, + #[clap(value_hint = ValueHint::Other)] + proxy: String, + }, /// List all groups List, } @@ -24,14 +29,22 @@ pub(crate) enum ConnOptions { /// List all active connections List, Stop { + #[clap(value_hint = ValueHint::Other)] nth: Option, }, } -#[derive(Debug, Subcommand)] +#[derive(Debug, Clone, Copy, Subcommand)] +pub(crate) enum TunSetOptions { + On, + Off, +} + +#[derive(Debug, Clone, Copy, Subcommand)] pub(crate) enum TunOptions { /// Set TUN - Set { s: String }, + #[command(subcommand)] + Set(TunSetOptions), /// Get TUN status Get, } @@ -45,9 +58,15 @@ pub(crate) struct CertOptions { #[derive(Debug, Subcommand)] pub(crate) enum TempRuleOptions { /// Add a temporary rule to the head of rule list - Add { literal: String }, + Add { + #[clap(value_hint = ValueHint::Other)] + literal: String, + }, /// Delete temporary rules matching this prefix - Delete { literal: String }, + Delete { + #[clap(value_hint = ValueHint::Other)] + literal: String, + }, /// Delete all temporary rules Clear, } @@ -57,9 +76,17 @@ pub(crate) enum InterceptOptions { /// List all captured data List, /// List data ranged from *start* to *end* - Range { start: u32, end: Option }, + Range { + #[clap(value_hint = ValueHint::Other)] + start: u32, + #[clap(value_hint = ValueHint::Other)] + end: Option, + }, /// Get details of the packet - Get { id: u32 }, + Get { + #[clap(value_hint = ValueHint::Other)] + id: u32, + }, } #[derive(Debug, Args)] @@ -85,33 +112,51 @@ pub(crate) struct InitOptions { pub app_data: Option, } +#[derive(Debug, Clone, Copy, Subcommand)] +pub(crate) enum PromptOptions { + Bash, + Zsh, + Fish, +} + #[derive(Debug, Subcommand)] -pub(crate) enum SubCommand { - /// Start Main Program - Start(StartOptions), - /// Create Configurations +pub(crate) enum GenerateOptions { + /// Create configurations Init(InitOptions), - /// Proxy Settings + /// Generate certificates + Cert(CertOptions), + /// Generate auto-completion profiles for shells #[command(subcommand)] - Proxy(ProxyOptions), - /// Connection Settings + Prompt(PromptOptions), +} + +#[derive(Debug, Subcommand)] +pub(crate) enum SubCommand { + /// Start the main program + Start(StartOptions), + /// Reload configurations + Reload, + /// Connection settings #[command(subcommand)] Conn(ConnOptions), - /// Generate Certificates - Cert(CertOptions), - /// Captured HTTP Data + /// Captured HTTP data #[command(subcommand)] Intercept(InterceptOptions), - /// Adjust TUN Status + /// Proxy settings #[command(subcommand)] - Tun(TunOptions), - /// Modify Temporary Rules + Proxy(ProxyOptions), + /// Modify temporary rules #[command(subcommand)] - Rule(TempRuleOptions), - /// Clean Unexpected Shutdown + TempRule(TempRuleOptions), + /// Adjust TUN status + #[command(subcommand)] + Tun(TunOptions), + /// Clean unexpected shutdown Clean, - /// Reload Configuration - Reload, + /// Generate necessary files before the first run + #[command(subcommand)] + Generate(GenerateOptions), + #[cfg(feature = "internal-test")] #[clap(hide = true)] Internal, } @@ -119,7 +164,7 @@ pub(crate) enum SubCommand { pub(crate) async fn controller_main(args: ProgramArgs) -> ! { let default_uds_path = "/var/run/boltconn.sock"; match args.cmd { - SubCommand::Init(init) => { + SubCommand::Generate(GenerateOptions::Init(init)) => { fn create(init: InitOptions) -> anyhow::Result<()> { let (config, data, _) = crate::config::parse_paths(&init.config, &init.app_data, &None)?; @@ -142,7 +187,7 @@ pub(crate) async fn controller_main(args: ProgramArgs) -> ! { } } } - SubCommand::Cert(opt) => { + SubCommand::Generate(GenerateOptions::Cert(opt)) => { if !is_root() { eprintln!("Must be run with root/admin privilege"); exit(-1) @@ -173,6 +218,17 @@ pub(crate) async fn controller_main(args: ProgramArgs) -> ! { } } } + SubCommand::Generate(GenerateOptions::Prompt(shell)) => { + let generator = match shell { + PromptOptions::Bash => clap_complete::Shell::Bash, + PromptOptions::Zsh => clap_complete::Shell::Zsh, + PromptOptions::Fish => clap_complete::Shell::Fish, + }; + let mut command = ProgramArgs::command(); + let bin_name = command.get_name().to_string(); + clap_complete::generate(generator, &mut command, bin_name, &mut std::io::stdout()); + exit(0) + } SubCommand::Clean => { if !is_root() { eprintln!("Must be run with root/admin privilege"); @@ -206,7 +262,14 @@ pub(crate) async fn controller_main(args: ProgramArgs) -> ! { }, SubCommand::Tun(opt) => match opt { TunOptions::Get => requester.get_tun().await, - TunOptions::Set { s } => requester.set_tun(s.as_str()).await, + TunOptions::Set(s) => { + requester + .set_tun(match s { + TunSetOptions::On => true, + TunSetOptions::Off => false, + }) + .await + } }, SubCommand::Intercept(opt) => match opt { InterceptOptions::List => requester.intercept(None).await, @@ -214,16 +277,16 @@ pub(crate) async fn controller_main(args: ProgramArgs) -> ! { InterceptOptions::Get { id } => requester.get_intercept_payload(id).await, }, SubCommand::Reload => requester.reload_config().await, - SubCommand::Rule(opt) => match opt { + SubCommand::TempRule(opt) => match opt { TempRuleOptions::Add { literal } => requester.add_temporary_rule(literal).await, TempRuleOptions::Delete { literal } => requester.delete_temporary_rule(literal).await, TempRuleOptions::Clear => requester.clear_temporary_rule().await, }, - SubCommand::Start(_) - | SubCommand::Init(_) - | SubCommand::Cert(_) - | SubCommand::Clean - | SubCommand::Internal => { + SubCommand::Start(_) | SubCommand::Generate(_) | SubCommand::Clean => { + unreachable!() + } + #[cfg(feature = "internal-test")] + SubCommand::Internal => { unreachable!() } }; diff --git a/boltconn/src/cli/request.rs b/boltconn/src/cli/request.rs index 75d2e4f..e84a65f 100644 --- a/boltconn/src/cli/request.rs +++ b/boltconn/src/cli/request.rs @@ -109,14 +109,8 @@ impl Requester { Ok(()) } - pub async fn set_tun(&self, content: &str) -> Result<()> { - let enabled = boltapi::TunStatusSchema { - enabled: match content.to_lowercase().as_str() { - "on" => true, - "off" => false, - _ => return Err(anyhow::anyhow!("Unknown TUN setting: {}", content)), - }, - }; + pub async fn set_tun(&self, enabled: bool) -> Result<()> { + let enabled = boltapi::TunStatusSchema { enabled }; match &self.inner { Inner::Web(c) => c.set_tun(enabled).await, Inner::Uds(c) => c.set_tun(enabled).await, diff --git a/boltconn/src/main.rs b/boltconn/src/main.rs index 07cf046..9817ff1 100644 --- a/boltconn/src/main.rs +++ b/boltconn/src/main.rs @@ -43,6 +43,7 @@ fn main() -> ExitCode { let args: ProgramArgs = ProgramArgs::parse(); let cmds = match args.cmd { SubCommand::Start(sub) => sub, + #[cfg(feature = "internal-test")] SubCommand::Internal => return internal_code(), _ => rt.block_on(cli::controller_main(args)), }; @@ -113,6 +114,7 @@ fn main() -> ExitCode { } /// This function is a shortcut for testing things conveniently. Only for development use. +#[cfg(feature = "internal-test")] fn internal_code() -> ExitCode { println!("This option is not for end-user."); ExitCode::SUCCESS