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;
#[derive(FromArgs, Debug)]
struct Arguments {
#[argh(option, short = 'p', default = "String::from(\"/dev/ttyS0\")")]
port: String,
#[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 = 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;
while !session_done || !receiver.drain_outgoing().is_empty() {
let mut progressed = false;
if !receiver.drain_outgoing().is_empty() {
match port.write(receiver.drain_outgoing()) {
Ok(0) => {}
Ok(n) => {
receiver.advance_outgoing(n);
progressed = true;
}
Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {}
Err(e) => bail!("wire write failed: {e}"),
}
}
if !receiver.drain_file().is_empty() {
if let Some((file, pb)) = open_files.get_mut(¤t_filename) {
match file.write(receiver.drain_file()) {
Ok(0) => {}
Ok(n) => {
receiver
.advance_file(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(receiver.drain_file()) {
Ok(0) => {}
Ok(n) => {
receiver
.advance_file(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}"),
}
}
}
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}"),
}
if receiver.drain_outgoing().is_empty()
&& receiver.drain_file().is_empty()
&& input_offset < input_buf.len()
{
let consumed = receiver
.feed_incoming(&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;
}
}
}
while let Some(event) = receiver.poll_event() {
match event {
zmodem2::ReceiverEvent::FileStart => {
let name = receiver.file_name();
if name != current_filename.as_slice() {
let received_filename = filename_from_u8(name)?;
let path = dest_dir.join(received_filename);
let file = File::create(&path)
.with_context(|| format!("could not create '{}'", path.display()))?;
let pb = if show_progress {
let pb = ProgressBar::new(u64::from(receiver.file_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(name.to_vec(), (file, pb));
current_filename = name.to_vec();
}
}
zmodem2::ReceiverEvent::FileComplete => {
if let Some((_file, pb)) = open_files.remove(¤t_filename) {
if show_progress {
pb.finish_and_clear();
} else {
pb.finish();
}
}
current_filename.clear();
}
zmodem2::ReceiverEvent::SessionComplete => {
session_done = true;
}
zmodem2::ReceiverEvent::Aborted => bail!("ZMODEM session aborted"),
}
}
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(())
}