extern crate zmodem2;
use anyhow::{bail, Context};
use argh::FromArgs;
use indicatif::{ProgressBar, ProgressStyle};
use lrzsz2::ReadWrite;
use std::collections::HashMap;
use std::ffi::{OsStr, OsString};
use std::fs::{File, OpenOptions};
use std::io::{IsTerminal, Read, Write};
#[cfg(unix)]
use std::os::unix::ffi::OsStrExt;
use std::path::{Component, Path, PathBuf};
use std::thread::sleep;
use std::time::Duration;
use zmodem2::{Action, Event};
#[derive(FromArgs, Debug)]
struct Arguments {
#[argh(option, short = 'p', default = "String::from(\"/dev/ttyS0\")")]
port: String,
#[argh(switch)]
overwrite: bool,
#[argh(switch)]
no_progress: bool,
#[argh(positional)]
path: String,
}
#[cfg(unix)]
fn filename_from_u8(bytes: &[u8]) -> anyhow::Result<OsString> {
let path = Path::new(OsStr::from_bytes(bytes));
let name = path
.components()
.next_back()
.and_then(|c| match c {
Component::Normal(name) => Some(name),
_ => None,
})
.with_context(|| format!("malformed filename: '{}'", String::from_utf8_lossy(bytes)))?;
Ok(name.to_os_string())
}
#[cfg(not(unix))]
fn filename_from_u8(bytes: &[u8]) -> anyhow::Result<OsString> {
let filename_str = std::str::from_utf8(bytes).context("invalid UTF-8 in filename")?;
let path = Path::new(filename_str);
let name = path
.components()
.next_back()
.and_then(|c| match c {
Component::Normal(name) => Some(name),
_ => None,
})
.with_context(|| format!("malformed filename: '{}'", filename_str))?;
Ok(name.to_os_string())
}
fn main() -> anyhow::Result<()> {
let args: Arguments = argh::from_env();
let port_file = OpenOptions::new()
.read(true)
.write(true)
.open(&args.port)
.with_context(|| format!("'{}'", args.port))?;
lrzsz2::configure_raw_mode(&port_file).ok();
lrzsz2::configure_nonblocking(&port_file).ok();
let mut port: Box<dyn ReadWrite> = Box::new(port_file);
let show_progress = !args.no_progress && std::io::stderr().is_terminal();
let mut receiver =
zmodem2::Receiver::new().map_err(|e| anyhow::anyhow!("ZMODEM error: {:#}", e))?;
let dest_dir = PathBuf::from(&args.path);
if !dest_dir.is_dir() {
bail!("'{}' is not a directory", dest_dir.display());
}
let mut open_files: HashMap<Vec<u8>, (File, ProgressBar)> = HashMap::new();
let mut current_filename = vec![];
let mut sink = std::io::sink();
let mut wire_buf = [0u8; 4096];
let mut input_buf: Vec<u8> = Vec::new();
let mut input_offset: usize = 0;
let mut session_done = false;
loop {
let mut progressed = false;
let mut idle = false;
match port.read(&mut wire_buf) {
Ok(0) => {}
Ok(n) => {
input_buf.extend_from_slice(&wire_buf[..n]);
progressed = true;
}
Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {}
Err(e) => bail!("wire read failed: {e}"),
}
match receiver.poll() {
Action::WriteWire(bytes) => match port.write(bytes) {
Ok(0) => {}
Ok(n) => {
receiver.wire_written(n);
progressed = true;
}
Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {}
Err(e) => bail!("wire write failed: {e}"),
},
Action::WriteFile(bytes) => {
if let Some((file, pb)) = open_files.get_mut(¤t_filename) {
match file.write(bytes) {
Ok(0) => {}
Ok(n) => {
receiver
.file_written(n)
.map_err(|e| anyhow::anyhow!("ZMODEM error: {:#}", e))?;
pb.inc(n as u64);
progressed = true;
}
Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {}
Err(e) => bail!("file write failed: {e}"),
}
} else {
match sink.write(bytes) {
Ok(0) => {}
Ok(n) => {
receiver
.file_written(n)
.map_err(|e| anyhow::anyhow!("ZMODEM error: {:#}", e))?;
progressed = true;
}
Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {}
Err(e) => bail!("file write failed: {e}"),
}
}
}
Action::Event(event) => {
progressed = true;
match event {
Event::FileStarted(info) => {
if info.name != current_filename.as_slice() {
let received_filename = filename_from_u8(info.name)?;
let path = dest_dir.join(received_filename);
let mut options = OpenOptions::new();
options.write(true);
if args.overwrite {
options.create(true).truncate(true);
} else {
options.create_new(true);
}
let file = options.open(&path).with_context(|| {
format!("could not create '{}'", path.display())
})?;
let size = info.size.map(|p| p.get()).unwrap_or(0);
let pb = if show_progress {
let pb = ProgressBar::new(u64::from(size));
pb.set_style(
ProgressStyle::default_bar()
.template(
"{spinner} [{elapsed_precise}] {msg} [{bar:40}] {bytes}/{total_bytes} ({eta})",
)?
.progress_chars("=>-"),
);
pb.set_message(format!("Receiving {}", path.display()));
pb
} else {
ProgressBar::hidden()
};
open_files.insert(info.name.to_vec(), (file, pb));
current_filename = info.name.to_vec();
}
}
Event::FileCompleted => {
if let Some((_file, pb)) = open_files.remove(¤t_filename) {
if show_progress {
pb.finish_and_clear();
} else {
pb.finish();
}
}
current_filename.clear();
}
Event::SessionCompleted => {
session_done = true;
}
Event::Aborted => bail!("ZMODEM session aborted"),
_ => {}
}
}
Action::ReadFile { .. } => unreachable!("receiver never reads files"),
Action::Idle => {
idle = true;
if input_offset < input_buf.len() {
let consumed = receiver
.submit_wire(&input_buf[input_offset..])
.map_err(|e| anyhow::anyhow!("ZMODEM error: {:#}", e))?;
if consumed > 0 {
input_offset += consumed;
progressed = true;
if input_offset == input_buf.len() {
input_buf.clear();
input_offset = 0;
} else if input_offset > 4096 {
input_buf.drain(..input_offset);
input_offset = 0;
}
}
}
}
_ => {}
}
if session_done && idle {
break;
}
if !progressed {
sleep(Duration::from_millis(1));
}
}
for (_filename, (_file, pb)) in open_files.iter() {
if !pb.is_finished() {
if show_progress {
pb.finish_and_clear();
} else {
pb.finish();
}
}
}
Ok(())
}