pfs-tftp-claude/crates/pfs-tftp/src/bin/tftp.rs

//! TFTP Client (tftp)
//!
//! A simple TFTP client implementation following RFC 1350.
//!
//! # Usage
//!
//! ```text
//! tftp [OPTIONS] <COMMAND>
//!
//! Commands:
//!   get <HOST> <REMOTE_FILE> [LOCAL_FILE]    Download a file
//!   put <HOST> <LOCAL_FILE> [REMOTE_FILE]    Upload a file
//!
//! Options:
//!   -p, --port <PORT>    Server port (default: 69)
//!   -m, --mode <MODE>    Transfer mode: octet (default) or netascii
//!   -v, --verbose        Enable verbose output
//!   -h, --help           Print help
//! ```
//!
//! # Examples
//!
//! ```text
//! # Download a file
//! tftp get 192.168.1.1 config.txt
//!
//! # Download from a custom port
//! tftp -p 6969 get 192.168.1.1 config.txt
//!
//! # Upload a file
//! tftp put 192.168.1.1 firmware.bin
//!
//! # Upload with netascii mode
//! tftp -m netascii put 192.168.1.1 readme.txt
//! ```

use std::{
    env,
    fs::File,
    io::{self, Write},
    path::{Path, PathBuf},
    process::ExitCode,
};

use pfs_tftp::{Client, Mode};

/// A writer that defers file creation until the first write.
///
/// This ensures the local file is only created after we've confirmed the remote
/// file exists (i.e., after receiving the first DATA packet, not an ERROR).
struct DeferredFileWriter {
    path: PathBuf,
    file: Option<File>,
    bytes_written: u64,
}

impl DeferredFileWriter {
    fn new(path: PathBuf) -> Self {
        Self {
            path,
            file: None,
            bytes_written: 0,
        }
    }

    fn bytes_written(&self) -> u64 {
        self.bytes_written
    }
}

impl Write for DeferredFileWriter {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if self.file.is_none() {
            self.file = Some(File::create(&self.path)?);
        }
        // SAFETY: We just ensured self.file is Some above
        let file = self.file.as_mut().expect("file is Some");
        let n = file.write(buf)?;
        self.bytes_written += n as u64;
        Ok(n)
    }

    fn flush(&mut self) -> io::Result<()> {
        if let Some(ref mut file) = self.file {
            file.flush()
        } else {
            Ok(())
        }
    }
}

/// Print usage information.
fn print_usage(program: &str) {
    eprintln!("TFTP Client (RFC 1350)");
    eprintln!();
    eprintln!("Usage: {program} [OPTIONS] <COMMAND>");
    eprintln!();
    eprintln!("Commands:");
    eprintln!("  get <HOST> <REMOTE_FILE> [LOCAL_FILE]    Download a file");
    eprintln!("  put <HOST> <LOCAL_FILE> [REMOTE_FILE]    Upload a file");
    eprintln!();
    eprintln!("Options:");
    eprintln!("  -p, --port <PORT>    Server port (default: 69)");
    eprintln!("  -m, --mode <MODE>    Transfer mode: octet (default) or netascii");
    eprintln!("  -v, --verbose        Enable verbose output");
    eprintln!("  -h, --help           Print help");
    eprintln!();
    eprintln!("Examples:");
    eprintln!("  {program} get 192.168.1.1 config.txt");
    eprintln!("  {program} -p 6969 get 192.168.1.1 config.txt");
    eprintln!("  {program} -m netascii put 192.168.1.1 readme.txt");
}

/// Command to execute.
enum Command {
    Get {
        host: String,
        remote_file: String,
        local_file: String,
    },
    Put {
        host: String,
        local_file: String,
        remote_file: String,
    },
}

/// Parsed command-line arguments.
struct Args {
    command: Command,
    port: Option<u16>,
    mode: Mode,
    verbose: bool,
}

/// Parse command-line arguments.
fn parse_args() -> Result<Args, String> {
    let args: Vec<String> = env::args().collect();
    let program = &args[0];

    let mut port: Option<u16> = None;
    let mut mode = Mode::Octet;
    let mut verbose = false;

    let mut i = 1;

    // Parse options
    while i < args.len() && args[i].starts_with('-') {
        match args[i].as_str() {
            "-h" | "--help" => {
                print_usage(program);
                std::process::exit(0);
            }
            "-p" | "--port" => {
                i += 1;
                if i >= args.len() {
                    return Err("--port requires an argument".to_string());
                }
                port = Some(
                    args[i]
                        .parse()
                        .map_err(|_| format!("invalid port: {}", args[i]))?,
                );
            }
            "-m" | "--mode" => {
                i += 1;
                if i >= args.len() {
                    return Err("--mode requires an argument".to_string());
                }
                mode = match args[i].to_lowercase().as_str() {
                    "octet" | "binary" => Mode::Octet,
                    "netascii" | "ascii" => Mode::NetAscii,
                    _ => return Err(format!("invalid mode: {}", args[i])),
                };
            }
            "-v" | "--verbose" => {
                verbose = true;
            }
            opt => {
                return Err(format!("unknown option: {opt}"));
            }
        }
        i += 1;
    }

    // Parse command
    if i >= args.len() {
        return Err("missing command".to_string());
    }

    let command = match args[i].as_str() {
        "get" => {
            i += 1;
            if i + 1 >= args.len() {
                return Err("get requires HOST and REMOTE_FILE".to_string());
            }
            let host = args[i].clone();
            let remote_file = args[i + 1].clone();
            let local_file = if i + 2 < args.len() {
                args[i + 2].clone()
            } else {
                // Use the filename part of the remote file
                Path::new(&remote_file)
                    .file_name()
                    .map_or_else(|| remote_file.clone(), |s| s.to_string_lossy().to_string())
            };
            Command::Get {
                host,
                remote_file,
                local_file,
            }
        }
        "put" => {
            i += 1;
            if i + 1 >= args.len() {
                return Err("put requires HOST and LOCAL_FILE".to_string());
            }
            let host = args[i].clone();
            let local_file = args[i + 1].clone();
            let remote_file = if i + 2 < args.len() {
                args[i + 2].clone()
            } else {
                // Use the filename part of the local file
                Path::new(&local_file)
                    .file_name()
                    .map_or_else(|| local_file.clone(), |s| s.to_string_lossy().to_string())
            };
            Command::Put {
                host,
                local_file,
                remote_file,
            }
        }
        cmd => {
            return Err(format!("unknown command: {cmd}"));
        }
    };

    Ok(Args {
        command,
        port,
        mode,
        verbose,
    })
}

/// Format host for connection.
///
/// If a port override is provided, it is used. Otherwise, if the host already
/// contains a port (has ':'), it is used as-is. Otherwise, the default TFTP
/// port 69 is appended.
fn format_host(host: &str, port_override: Option<u16>) -> String {
    if let Some(port) = port_override {
        // Strip any existing port from host and use override
        let host_part = host.split(':').next().unwrap_or(host);
        format!("{host_part}:{port}")
    } else if host.contains(':') {
        host.to_string()
    } else {
        format!("{host}:69")
    }
}

#[allow(clippy::too_many_lines)]
fn main() -> ExitCode {
    let args = match parse_args() {
        Ok(args) => args,
        Err(e) => {
            eprintln!("Error: {e}");
            eprintln!("Try '--help' for more information.");
            return ExitCode::FAILURE;
        }
    };

    match args.command {
        Command::Get {
            host,
            remote_file,
            local_file,
        } => {
            let addr = format_host(&host, args.port);

            if args.verbose {
                eprintln!("Getting '{remote_file}' from {addr} -> '{local_file}'");
                eprintln!(
                    "Mode: {}",
                    match args.mode {
                        Mode::Octet => "octet",
                        Mode::NetAscii => "netascii",
                    }
                );
            }

            let client = match Client::new(&addr) {
                Ok(c) => c,
                Err(e) => {
                    eprintln!("Error connecting to {addr}: {e}");
                    return ExitCode::FAILURE;
                }
            };

            // Use deferred writer - file is only created after first DATA packet
            let mut writer = DeferredFileWriter::new(PathBuf::from(&local_file));

            if let Err(e) = client.get_to_writer(&remote_file, args.mode, &mut writer) {
                eprintln!("Error: {e}");
                return ExitCode::FAILURE;
            }

            let bytes = writer.bytes_written();
            if args.verbose {
                eprintln!("Received {bytes} bytes");
            }
            println!("Downloaded '{remote_file}' -> '{local_file}' ({bytes} bytes)");
        }

        Command::Put {
            host,
            local_file,
            remote_file,
        } => {
            let addr = format_host(&host, args.port);

            if args.verbose {
                eprintln!("Putting '{local_file}' -> {addr} as '{remote_file}'");
                eprintln!(
                    "Mode: {}",
                    match args.mode {
                        Mode::Octet => "octet",
                        Mode::NetAscii => "netascii",
                    }
                );
            }

            let client = match Client::new(&addr) {
                Ok(c) => c,
                Err(e) => {
                    eprintln!("Error connecting to {addr}: {e}");
                    return ExitCode::FAILURE;
                }
            };

            // Open input file
            let mut file = match File::open(&local_file) {
                Ok(f) => f,
                Err(e) => {
                    eprintln!("Error opening file '{local_file}': {e}");
                    return ExitCode::FAILURE;
                }
            };

            // Upload
            match client.put_from_reader(&remote_file, args.mode, &mut file) {
                Ok(bytes) => {
                    if args.verbose {
                        eprintln!("Sent {bytes} bytes");
                    }
                    println!("Uploaded '{local_file}' -> '{remote_file}' ({bytes} bytes)");
                }
                Err(e) => {
                    eprintln!("Error: {e}");
                    return ExitCode::FAILURE;
                }
            }
        }
    }

    ExitCode::SUCCESS
}