1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
|
/*
* Takes the first n lines of input, passes to stdout.
*/
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(())
}
|
