use std::env;
use std::fs::File;
use std::io::{self, Write};
use std::net::{SocketAddr, ToSocketAddrs, UdpSocket};
use std::path::{Path, PathBuf};
use std::time::Duration;
extern crate getopts;
use getopts::Options;
extern crate rtftp;
enum Operation {
RRQ,
WRQ,
}
struct Configuration {
operation: Operation,
mode: rtftp::Mode,
filename: PathBuf,
remote: SocketAddr,
blksize: usize,
}
struct Tftpc {
tftp: rtftp::Tftp,
conf: Configuration,
}
fn update_progress(current: u64, total: u64, last: u64) -> u64 {
if total == 0 {
return 0;
}
let onepercent = total / 100;
if current < total && current < last + onepercent {
return last;
}
let percent = 100 * current / total;
print!("\r {}% ", percent);
io::stdout().flush().expect("flushing stdout failed");
if current >= total {
print!("\r");
}
current
}
impl Tftpc {
pub fn new(conf: Configuration) -> Tftpc {
let mut tftp = rtftp::Tftp::new();
tftp.set_mode(conf.mode);
Tftpc {
tftp,
conf,
}
}
fn wait_for_option_ack(&mut self, sock: &UdpSocket) -> Option<SocketAddr> {
let mut buf = [0; 512];
sock.peek_from(&mut buf).ok()?;
let opcode = u16::from_be_bytes([buf[0], buf[1]]);
if opcode != rtftp::Opcode::OACK as u16 {
return None;
}
let (len, remote) = sock.recv_from(&mut buf).ok()?;
let mut options = self.tftp.parse_options(&buf[2..len]);
self.tftp.init_tftp_options(sock, &mut options).ok()?;
Some(remote)
}
fn wait_for_response(&self, sock: &UdpSocket, expected_opcode: rtftp::Opcode, expected_block: u16, expected_remote: Option<SocketAddr>) -> Result<Option<SocketAddr>, std::io::Error> {
let mut buf = [0; 4];
let (len, remote) = sock.peek_from(&mut buf)?;
if let Some(rem) = expected_remote {
if rem != remote {
return Ok(None);
}
}
let opcode = u16::from_be_bytes([buf[0], buf[1]]);
let block_nr = u16::from_be_bytes([buf[2], buf[3]]);
if opcode == rtftp::Opcode::ERROR as u16 {
let mut buf = [0; 512];
let len = sock.recv(&mut buf)?;
return Err(self.tftp.parse_error(&buf[..len]));
}
if len != 4 || opcode != expected_opcode as u16 || block_nr != expected_block {
return Ok(None);
}
Ok(Some(remote))
}
fn append_option_req(&self, buf: &mut Vec<u8>, fsize: u64) {
self.tftp.append_option(buf, "blksize", &format!("{}", self.conf.blksize));
self.tftp.append_option(buf, "timeout", &format!("{}", 3));
self.tftp.append_option(buf, "tsize", &format!("{}", fsize));
}
fn init_req(&self, opcode: rtftp::Opcode, filename: &str, size: u64) -> Vec<u8> {
let mut buf = Vec::with_capacity(512);
buf.extend((opcode as u16).to_be_bytes().iter());
let mode_str = match self.conf.mode {
rtftp::Mode::OCTET => "octet",
rtftp::Mode::NETASCII => "netascii",
};
self.tftp.append_option(&mut buf, filename, mode_str);
self.append_option_req(&mut buf, size);
buf
}
fn handle_wrq(&mut self, sock: &UdpSocket) -> Result<String, io::Error> {
let mut file = File::open(self.conf.filename.as_path())?;
let err_invalidpath = || io::Error::new(io::ErrorKind::InvalidInput, "Invalid path/filename");
let filename = self.conf.filename.file_name().ok_or_else(err_invalidpath)?
.to_str().ok_or_else(err_invalidpath)?;
let metadata = file.metadata().map_err(|_| err_invalidpath())?;
if !metadata.is_file() {
return Err(err_invalidpath());
}
let tsize = self.tftp.transfersize(&mut file)?;
let buf = self.init_req(rtftp::Opcode::WRQ, filename, tsize);
let mut remote = None;
for _ in 1..3 {
sock.send_to(&buf, self.conf.remote)?;
remote = self.wait_for_option_ack(sock);
if remote.is_none() {
remote = self.wait_for_response(sock, rtftp::Opcode::ACK, 0, None)?;
}
if remote.is_some() {
break;
}
}
match remote {
Some(r) => sock.connect(r).expect("connecting to remote failed"),
None => return Err(io::Error::new(io::ErrorKind::TimedOut, "No response from server")),
}
match self.tftp.send_file(sock, &mut file) {
Ok(_) => Ok(format!("Sent {} to {}.", self.conf.filename.display(), self.conf.remote)),
Err(err) => {
let error = format!("Sending {} to {} failed ({}).", self.conf.filename.display(), self.conf.remote, err);
self.tftp.send_error(sock, 0, "Sending error")?;
Err(io::Error::new(err.kind(), error))
}
}
}
fn handle_rrq(&mut self, sock: &UdpSocket) -> Result<String, io::Error> {
let err_invalidpath = || io::Error::new(io::ErrorKind::InvalidInput, "Invalid path/filename");
let filename = self.conf.filename.file_name().ok_or_else(err_invalidpath)?;
let outpath = env::current_dir().expect("Can't get current directory").join(filename);
let mut file = File::create(outpath)?;
let filename = self.conf.filename.to_str().ok_or_else(err_invalidpath)?;
let buf = self.init_req(rtftp::Opcode::RRQ, filename, 0);
let mut remote = None;
for _ in 1..3 {
sock.send_to(&buf, self.conf.remote)?;
let oack_remote = self.wait_for_option_ack(sock);
if let Some(r) = oack_remote {
self.tftp.send_ack_to(sock, r, 0)?;
}
remote = self.wait_for_response(sock, rtftp::Opcode::DATA, 1, oack_remote)?;
if remote.is_some() {
break;
}
}
match remote {
Some(r) => sock.connect(r).expect("connecting to remote failed"),
None => return Err(io::Error::new(io::ErrorKind::TimedOut, "No response from server")),
}
match self.tftp.recv_file(sock, &mut file) {
Ok(_) => Ok(format!("Received {} from {}.", self.conf.filename.display(), self.conf.remote)),
Err(err) => {
let error = format!("Receiving {} from {} failed ({}).", self.conf.filename.display(), self.conf.remote, err);
self.tftp.send_error(sock, 0, "Receiving error")?;
Err(std::io::Error::new(err.kind(), error))
}
}
}
pub fn start(&mut self) {
self.tftp.set_progress_callback(update_progress);
let socket = UdpSocket::bind("[::]:0").expect("binding failed");
socket.set_read_timeout(Some(Duration::from_secs(5))).expect("setting socket timeout failed");
let err = match self.conf.operation {
Operation::RRQ => self.handle_rrq(&socket),
Operation::WRQ => self.handle_wrq(&socket),
};
match err {
Ok(msg) => println!("{}", msg),
Err(err) => println!("Error: {}", err),
}
}
}
fn usage(opts: &Options, program: &str, error: Option<String>) {
if let Some(err) = error {
println!("{}\n", err);
}
let version = rtftp::VERSION.unwrap_or("");
println!("{}", opts.usage(format!("RusTFTP {}\n\n{} [options] <remote>[:port]", version, program).as_str()));
}
fn parse_commandline(args: &[String]) -> Option<Configuration> {
let program = args[0].clone();
let mut operation = None;
let mut mode = rtftp::Mode::OCTET;
let mut filename = None;
let mut blksize = 1428;
let mut opts = Options::new();
opts.optflag("h", "help", "display usage information");
opts.optopt("g", "get", "download file from remote server", "FILE");
opts.optopt("p", "put", "upload file to remote server", "FILE");
opts.optopt("b", "blksize", format!("negotiate a different block size (default: {})", blksize).as_ref(), "SIZE");
opts.optflag("n", "netascii","use netascii mode (instead of octet)");
let getopts_fail = |err: getopts::Fail| { usage(&opts, &program, Some(err.to_string())) };
let conv_error = |err: std::num::ParseIntError| { usage(&opts, &program, Some(err.to_string())) };
let matches = opts.parse(&args[1..]).map_err(getopts_fail).ok()?;
if matches.opt_present("h") || matches.free.len() != 1 {
usage(&opts, &program, None);
return None;
}
if let Some(f) = matches.opt_str("g") {
operation = Some(Operation::RRQ);
filename = Some(Path::new(&f).to_path_buf());
}
if let Some(f) = matches.opt_str("p") {
operation = Some(Operation::WRQ);
filename = Some(Path::new(&f).to_path_buf());
}
if operation.is_none() || (matches.opt_present("g") && matches.opt_present("p")) {
usage(&opts, &program, Some("Exactly one of g (get) and p (put) required".to_string()));
return None;
}
if matches.opt_present("n") {
mode = rtftp::Mode::NETASCII;
}
let remote_in = matches.free[0].as_str();
let remote = match remote_in.to_socket_addrs() {
Ok(i) => i,
Err(_) => match (remote_in, 69).to_socket_addrs() {
Ok(j) => j,
Err(_) => {
usage(&opts, &program, Some("Failed to parse and lookup specified remote".to_string()));
return None;
}
},
}.next();
blksize = matches.opt_get_default::<usize>("b", blksize).map_err(conv_error).ok()?;
Some(Configuration {
operation: operation.unwrap(),
mode,
filename: filename.unwrap(),
remote: remote.unwrap(),
blksize,
})
}
fn main() {
let args: Vec<String> = env::args().collect();
let conf = match parse_commandline(&args) {
Some(c) => c,
None => return,
};
Tftpc::new(conf).start();
}