implement aft-head

1 files changed, 66 insertions, 3 deletions

diff --git a/src/bin/aft-head.rs b/src/bin/aft-head.rs
index b8c848c..3fecfa2 100644
--- a/src/bin/aft-head.rs
+++ b/src/bin/aft-head.rs
@@ -1,8 +1,71 @@
-
 /*
  * Takes the first n lines of input, passes to stdout.
  */
 
-fn main() {
-    println!("Hello, world2!");
+use std::error::Error;
+use std::io::{self, Read, Write, BufRead};
+use structopt::StructOpt;
+use std::path::PathBuf;
+
+#[derive(StructOpt)]
+#[structopt(rename_all = "kebab-case", about = "take first N rows of input (AFT-aware)" )]
+struct Opt {
+    #[structopt(long = "--lines", short = "-n", default_value = "10")]
+    lines: u64,
+
+    #[structopt(parse(from_os_str))]
+    input: Option<PathBuf>,
+}
+
+fn main() -> Result<(), Box<dyn Error>> {
+
+    let opt = Opt::from_args();
+
+    // treat "-" as "use stdin"
+    let input = match opt.input {
+        Some(s) if s.to_string_lossy() == "-" => None,
+        other => other,
+    };
+    let stdin = io::stdin();
+    let stdin = stdin.lock();
+    let mut input: Box<dyn BufRead> = match input {
+        None => {
+            Box::new(io::BufReader::new(stdin))
+        },
+        Some(path) => {
+            let f = std::fs::File::open(path)?;
+            Box::new(io::BufReader::new(f))
+        },
+    };
+    let mut stdout = io::stdout();
+
+    // read first byte of input to check if we are getting AFT; if not just pass through
+    let mut first_byte = [0; 1];
+    let got = input.read(&mut first_byte)?;
+    if got != 1 {
+        panic!("couldn't read a byte from input");
+    };
+
+    let is_aft = first_byte[0] == 0x01;
+    if !is_aft {
+        // simple implementation of lines
+        stdout.write(&first_byte)?;
+        for l in input.lines().take(opt.lines as usize) {
+            writeln!(stdout, "{}", l?)?;
+        }
+        return Ok(())
+    };
+
+    let mut header_bytes = vec![];
+    // TODO: sanity limit on length
+    input.read_until(0x02, &mut header_bytes)?;
+    if header_bytes[header_bytes.len()-1] != 0x02 {
+        panic!("expected an AFT header");
+    };
+    stdout.write(&first_byte)?;
+    stdout.write(&header_bytes)?;
+    for l in input.lines().take(opt.lines as usize) {
+        writeln!(stdout, "{}", l?)?;
+    }
+    Ok(())
 }