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(group(
ArgGroup::new("cmds")
.required(true)
.args(&["CONFIG", "genkey"]),
))]
pub struct Cli {
/// The path to the configuration file
///
/// Running as a client or a server is automatically determined
/// according to the configuration file.
#[clap(parse(from_os_str), name = "config")]
pub config_path: std::path::PathBuf,
#[clap(parse(from_os_str), name = "CONFIG")]
pub config_path: Option<std::path::PathBuf>,
/// Run as a server
#[clap(long, short)]
#[clap(long, short, group = "mode")]
pub server: bool,
/// Run as a client
#[clap(long, short)]
#[clap(long, short, group = "mode")]
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;
pub use cli::Cli;
use cli::KeypairType;
pub use config::Config;
pub use constants::UDP_BUFFER_SIZE;
@ -24,8 +25,37 @@ mod server;
#[cfg(feature = "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<()> {
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);
@ -155,9 +185,10 @@ mod tests {
};
let args = Cli {
config_path: std::path::PathBuf::new(),
config_path: Some(std::path::PathBuf::new()),
server: t.arg_s,
client: t.arg_c,
..Default::default()
};
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,
};
let params: NoiseParams = config.pattern.clone().parse()?;
let params: NoiseParams = config.pattern.parse()?;
Ok(NoiseTransport {
config,

View File

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