Use clap for argumentsHEADmaster

1 files changed, 13 insertions, 3 deletions

diff --git a/src/main.rs b/src/main.rs
index f80b06a..3aeba9c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,8 +1,18 @@
+use clap::Parser;
+
+#[derive(clap::clap_derive::Parser, Debug)]
+struct Args {
+    target: String,
+    flag_ids_path: String,
+}
+
 #[tokio::main]
 async fn main() {
-    let target = std::env::args().nth(1).unwrap();
-    let flag_ids_path = std::env::args().nth(2).unwrap();
-    let flag_ids = std::fs::read_to_string(flag_ids_path).unwrap();
+    let args = Args::parse();
+    let target = args.target;
+    let flag_ids = tokio::fs::read_to_string(&args.flag_ids_path)
+        .await
+        .expect("Failed to read flag IDs file");
 
     exploit(&target, &flag_ids).await;
 }