Add "output" argument to client

1 files changed, 9 insertions, 3 deletions

diff --git a/src/bin/roughenough-client.rs b/src/bin/roughenough-client.rs
index f7afea5..e1122b4 100644
--- a/src/bin/roughenough-client.rs
+++ b/src/bin/roughenough-client.rs
@@ -230,6 +230,12 @@ fn main() {
       .long("stress")
       .help("Stress-tests the server by sending the same request as fast as possible. Please only use this on your own server")
     )
+    .arg(Arg::with_name("output")
+      .short("o")
+      .long("output")
+      .takes_value(true)
+      .help("Writes all requsts to the specified file, in addition to sending them to the server. Useful for generating fuzer inputs")
+    )
     .get_matches();
 
     let host = matches.value_of("host").unwrap();
@@ -240,6 +246,7 @@ fn main() {
     let pub_key = matches
         .value_of("public-key")
         .map(|pkey| hex::decode(pkey).expect("Error parsing public key!"));
+    let out = matches.value_of("output");
 
     println!("Requesting time from: {:?}:{:?}", host, port);
 
@@ -266,14 +273,13 @@ fn main() {
     }
 
     let mut requests = Vec::with_capacity(num_requests);
-
-    let mut file = File::create("requests.bin").unwrap();
+    let mut file = out.map(|o | File::create(o).expect("Failed to create file!"));
 
     for _ in 0..num_requests {
         let nonce = create_nonce();
         let mut socket = UdpSocket::bind("0.0.0.0:0").expect("Couldn't open UDP socket");
         let request = make_request(&nonce);
-        file.write_all(&request);
+        file.as_mut().map(|f| f.write_all(&request).expect("Failed to write to file!"));
 
         requests.push((nonce, request, socket));
     }