feat: add `--genkey`

This commit is contained in:
Yujia Qiao 2021-12-24 20:42:01 +08:00 committed by Yujia Qiao
parent f472814013
commit 46ec6400e4
4 changed files with 61 additions and 11 deletions

View File

@ -1,20 +1,37 @@
use clap::{AppSettings, Parser}; use clap::{AppSettings, ArgGroup, Parser};
#[derive(Parser, Debug)] #[derive(clap::ArgEnum, Clone, Debug, Copy)]
pub enum KeypairType {
X25519,
X448,
}
#[derive(Parser, Debug, Default)]
#[clap(about, version, setting(AppSettings::DeriveDisplayOrder))] #[clap(about, version, setting(AppSettings::DeriveDisplayOrder))]
#[clap(group(
ArgGroup::new("cmds")
.required(true)
.args(&["CONFIG", "genkey"]),
))]
pub struct Cli { pub struct Cli {
/// The path to the configuration file /// The path to the configuration file
/// ///
/// Running as a client or a server is automatically determined /// Running as a client or a server is automatically determined
/// according to the configuration file. /// according to the configuration file.
#[clap(parse(from_os_str), name = "config")] #[clap(parse(from_os_str), name = "CONFIG")]
pub config_path: std::path::PathBuf, pub config_path: Option<std::path::PathBuf>,
/// Run as a server /// Run as a server
#[clap(long, short)] #[clap(long, short, group = "mode")]
pub server: bool, pub server: bool,
/// Run as a client /// Run as a client
#[clap(long, short)] #[clap(long, short, group = "mode")]
pub client: bool, pub client: bool,
/// Generate a keypair for the use of the noise protocol
///
/// The DH function to use is x25519
#[clap(long, arg_enum, value_name = "CURVE")]
pub genkey: Option<Option<KeypairType>>,
} }

View File

@ -7,6 +7,7 @@ mod protocol;
mod transport; mod transport;
pub use cli::Cli; pub use cli::Cli;
use cli::KeypairType;
pub use config::Config; pub use config::Config;
pub use constants::UDP_BUFFER_SIZE; pub use constants::UDP_BUFFER_SIZE;
@ -24,8 +25,37 @@ mod server;
#[cfg(feature = "server")] #[cfg(feature = "server")]
use server::run_server; use server::run_server;
const DEFAULT_CURVE: KeypairType = KeypairType::X25519;
fn get_str_from_keypair_type(curve: KeypairType) -> &'static str {
match curve {
KeypairType::X25519 => "25519",
KeypairType::X448 => "448",
}
}
fn genkey(curve: Option<KeypairType>) -> Result<()> {
let curve = curve.unwrap_or(DEFAULT_CURVE);
let builder = snowstorm::Builder::new(
format!(
"Noise_KK_{}_ChaChaPoly_BLAKE2s",
get_str_from_keypair_type(curve)
)
.parse()?,
);
let keypair = builder.generate_keypair()?;
println!("Private Key:\n{}\n", base64::encode(keypair.private));
println!("Public Key:\n{}", base64::encode(keypair.public));
Ok(())
}
pub async fn run(args: &Cli, shutdown_rx: broadcast::Receiver<bool>) -> Result<()> { pub async fn run(args: &Cli, shutdown_rx: broadcast::Receiver<bool>) -> Result<()> {
let config = Config::from_file(&args.config_path).await?; if args.genkey.is_some() {
return genkey(args.genkey.unwrap());
}
let config = Config::from_file(args.config_path.as_ref().unwrap()).await?;
debug!("{:?}", config); debug!("{:?}", config);
@ -155,9 +185,10 @@ mod tests {
}; };
let args = Cli { let args = Cli {
config_path: std::path::PathBuf::new(), config_path: Some(std::path::PathBuf::new()),
server: t.arg_s, server: t.arg_s,
client: t.arg_c, client: t.arg_c,
..Default::default()
}; };
assert_eq!(determine_run_mode(&config, &args), t.run_mode); assert_eq!(determine_run_mode(&config, &args), t.run_mode);

View File

@ -54,7 +54,7 @@ impl Transport for NoiseTransport {
None => builder.generate_keypair()?.private, None => builder.generate_keypair()?.private,
}; };
let params: NoiseParams = config.pattern.clone().parse()?; let params: NoiseParams = config.pattern.parse()?;
Ok(NoiseTransport { Ok(NoiseTransport {
config, config,

View File

@ -15,9 +15,10 @@ pub async fn run_rathole_server(
shutdown_rx: broadcast::Receiver<bool>, shutdown_rx: broadcast::Receiver<bool>,
) -> Result<()> { ) -> Result<()> {
let cli = rathole::Cli { let cli = rathole::Cli {
config_path: PathBuf::from(config_path), config_path: Some(PathBuf::from(config_path)),
server: true, server: true,
client: false, client: false,
..Default::default()
}; };
rathole::run(&cli, shutdown_rx).await rathole::run(&cli, shutdown_rx).await
} }
@ -27,9 +28,10 @@ pub async fn run_rathole_client(
shutdown_rx: broadcast::Receiver<bool>, shutdown_rx: broadcast::Receiver<bool>,
) -> Result<()> { ) -> Result<()> {
let cli = rathole::Cli { let cli = rathole::Cli {
config_path: PathBuf::from(config_path), config_path: Some(PathBuf::from(config_path)),
server: false, server: false,
client: true, client: true,
..Default::default()
}; };
rathole::run(&cli, shutdown_rx).await rathole::run(&cli, shutdown_rx).await
} }