feat(main): shell completions, clap

Author: janTatesa

Cargo.lock

@@ -13,6 +13,8 @@ rust-version = "1.69.0"
 [dependencies]
 ureq = { version = "2.9.6", default-features = false, optional = true }
 oorandom = { version = "11.1.3", optional = true }
+clap = { version = "4.5.4", features = ["cargo", "help"] }
+clap_complete = "4.5.2"
 
 [dev-dependencies]
 criterion = "0.5.1"

Cargo.toml

@@ -1,23 +1,74 @@
 #![deny(unsafe_code)]
 
+use std::io;
+
+use clap::{
+	builder::{styling::AnsiColor, Styles},
+	crate_authors, crate_description, crate_name, crate_version, value_parser, Arg, ArgAction,
+	Command,
+};
+use clap_complete::{generate, Generator, Shell};
+fn build_cli() -> Command {
+	Command::new(crate_name!())
+		.arg_required_else_help(true)
+		.author(crate_authors!())
+		.about(crate_description!())
+		.version(crate_version!())
+		.long_about(HELP)
+		.subcommand(Command::new("uaf").about("Run the use-after-free bug"))
+		.subcommand(
+			Command::new("bo").about("Run the buffer overflow exploit. Optionally take a shower"),
+		)
+		.subcommand(Command::new("transition").about("Safely transmute a Boy to a Girl"))
+		.subcommand(Command::new("segfault").about("Segfault yourself"))
+		.subcommand(
+			Command::new("completions")
+				.about("Return shell completions")
+				.arg(
+					Arg::new("shell")
+						.action(ArgAction::Set)
+						.value_parser(value_parser!(Shell))
+						.required(true),
+				),
+		)
+		.styles(STYLE)
+		.help_template(TEMPLATE)
+}
 fn main() {
-	let mut args = std::env::args();
-	let _program = args.next();
-	let Some(subcommand) = args.next() else {
-		println!("{HELP}");
-		return;
-	};
-
-	match subcommand.as_str() {
+	let mut command = build_cli();
+	let matches = build_cli().clone().get_matches();
+	let subcommand = matches.subcommand().unwrap();
+	match subcommand.0 {
 		"uaf" => cve_rs::use_after_free(),
 		"segfault" => cve_rs::segfault(),
 		"bo" => cve_rs::buffer_overflow().unwrap(),
 		"transition" => transmute_demo().unwrap(),
-		"help" | "--help" | "h" | "-h" | "?" | "-?" => println!("{HELP}"),
-		other => println!("Error: Unknown command `{other}`.\n{HELP}"),
+		"completions" => print_completions(
+			subcommand.1.get_one::<Shell>("shell").copied().unwrap(),
+			&mut command,
+		),
+		_ => unreachable!(),
 	}
 }
 
+fn print_completions<G: Generator>(gen: G, cmd: &mut Command) {
+	generate(gen, cmd, cmd.get_name().to_string(), &mut io::stdout());
+}
+
+const STYLE: Styles = Styles::styled()
+	.header(AnsiColor::Yellow.on_default())
+	.usage(AnsiColor::Green.on_default())
+	.literal(AnsiColor::Green.on_default())
+	.placeholder(AnsiColor::Green.on_default());
+
+const TEMPLATE: &str = "\
+{before-help}{name} {version}
+{author-with-newline}{about-with-newline}
+{usage-heading} {usage}
+
+{all-args}{after-help}
+";
+
 const HELP: &str = r"
 cve-rs: Blazingly fast memory vulnerabilities, written in 100% safe rust.
 
@@ -27,13 +78,6 @@ cve-rs exploits a soundness hole in lifetimes that lets us cast any lifetime to
 See: https://github.com/rust-lang/rust/issues/25860
 
 This program is open-source! View the source for all these exploits here: https://github.com/Speykious/cve-rs
-
-Commands:
-    help        Show this help message.
-    uaf         Run the use-after-free bug.
-    bo          Run the buffer overflow exploit. Optionally take a shower.
-    transition  Safely transmute a Boy to a Girl.
-    segfault    Segfault yourself.
 ";
 
 #[repr(C)]