ssh-stamp-ota 1.0.4

// SPDX-FileCopyrightText: 2026 Julio Beltran Ortega <jubeormk1@gmail.com>
//
// SPDX-License-Identifier: GPL-3.0-or-later

use ssh_stamp_ota::{OtaHeader, tlv};

use clap::{ArgAction, Command};
use sha2::{Digest, Sha256};
use std::{
    io::{Read, Seek, SeekFrom, Write},
    path::{Path, PathBuf},
};

const OTA_PACKER_VERSION: &str = env!("CARGO_PKG_VERSION");

const OK: i32 = 0;
const USAGE: i32 = 1;
const FILE_NOT_FOUND: i32 = 2;
const NOT_A_FILE: i32 = 3;
const OPEN_FAILED: i32 = 4;
const READ_FAILED: i32 = 5;
const CREATE_FAILED: i32 = 6;
const WRITE_FAILED: i32 = 7;
const SEEK_FAILED: i32 = 8;
const CHECKSUM_MISMATCH: i32 = 9;
const FILE_TOO_LARGE: i32 = 10;

fn main() {
    let matches = Command::new("packer")
        .about(format!("SSH-Stamp utility {OTA_PACKER_VERSION} to pack (unpack) OTA update files adding the required metadata."))
        .arg(clap::arg!(<FILE> "The file to process").required(true))
        .arg(
            clap::arg!(-u --unpack "Unpacks a OTA file. Will save to <file> with .ota.npkd extension")
                .action(ArgAction::SetTrue)
                .conflicts_with("pack"),
        )
        .arg(
            clap::arg!(-p --pack "(default) Packs a binary file as an OTA file. Will save to <file>.ota")
                .action(ArgAction::SetTrue)
                .conflicts_with("unpack"),
        )
        .get_matches();
    let Some(file_path) = matches.get_one::<String>("FILE") else {
        eprintln!("Error: No file provided");
        std::process::exit(USAGE);
    };

    let file_path = PathBuf::from(file_path);
    if !file_path.exists() {
        eprintln!("Error: File '{}' does not exist", file_path.display());
        std::process::exit(FILE_NOT_FOUND);
    }
    if !file_path.is_file() {
        eprintln!(
            "Error: File '{}' is not a regular file",
            file_path.display()
        );
        std::process::exit(NOT_A_FILE);
    }

    if matches.get_flag("unpack") {
        std::process::exit(unpack_ota(&file_path));
    }

    std::process::exit(pack_bin(&file_path));
}

fn unpack_ota(file_path: &Path) -> i32 {
    println!("Unpacking BIN from OTA file {}...", file_path.display());
    let Ok(file) = std::fs::File::open(file_path) else {
        eprintln!("Error: Could not open file '{}'", file_path.display());
        return OPEN_FAILED;
    };
    let mut reader = std::io::BufReader::new(file);
    let mut buffer = [0u8; 512];
    let Ok(_) = reader.read(&mut buffer) else {
        eprintln!("Error: Could not read from file '{}'", file_path.display());
        return READ_FAILED;
    };
    let Ok((header, seek_to_bin)) = OtaHeader::deserialize(&buffer) else {
        eprintln!(
            "Error: Could not parse OTA header from file '{}'",
            file_path.display(),
        );
        return READ_FAILED;
    };

    println!("Found OTA header: {header:?}");

    let mut file_path_bin = file_path.to_path_buf();
    file_path_bin.set_extension("ota.npkd");
    println!("Saving unpacked BIN file to: {}", file_path_bin.display());

    let Ok(mut bin_file) = std::fs::File::create(&file_path_bin) else {
        eprintln!(
            "Error: Could not create BIN file '{}'",
            file_path_bin.display(),
        );
        return CREATE_FAILED;
    };

    if let Err(e) = reader.seek(SeekFrom::Start(seek_to_bin as u64)) {
        eprintln!("Error: Could not seek to binary data: {e}");
        return SEEK_FAILED;
    }

    let mut recover_ota_bin_hasher = Sha256::new();

    let mut r: usize;
    while {
        r = reader.read(&mut buffer).unwrap_or(0);
        r
    } > 0
    {
        let Ok(_) = bin_file.write(&buffer[..r]) else {
            eprintln!(
                "Error: Could not write to BIN file '{}'",
                file_path_bin.display(),
            );
            return WRITE_FAILED;
        };
        recover_ota_bin_hasher.update(&buffer[..r]);
    }

    let recovered_firmware_sha256 = recover_ota_bin_hasher.finalize();

    let Some(expected_sha256) = header.sha256_checksum else {
        eprintln!("Error: OTA header has no SHA-256 checksum to verify against");
        return CHECKSUM_MISMATCH;
    };

    if recovered_firmware_sha256.as_slice() != expected_sha256 {
        eprintln!(
            "Error: Recovered firmware SHA-256 does not match expected value!\nExpected: {expected_sha256:x?}\nRecovered: {recovered_firmware_sha256:x?}"
        );
        return CHECKSUM_MISMATCH;
    }
    println!("Recovered firmware SHA-256 matches expected value.");

    OK
}

// TODO: Optimize memory usage by streaming the file instead of reading it all at once
fn pack_bin(file_path: &Path) -> i32 {
    println!("Packing {} as OTA...", file_path.display());

    let metadata = match file_path.metadata() {
        Ok(metadata) => metadata,
        Err(e) => {
            eprintln!(
                "Error: Could not retrieve metadata for file '{}': {e}",
                file_path.display()
            );
            return OPEN_FAILED;
        }
    };
    let Ok(firmware_size) = u32::try_from(metadata.len()) else {
        eprintln!(
            "Error: File '{}' is too large (max 4GB supported)",
            file_path.display()
        );
        return FILE_TOO_LARGE;
    };
    println!("Bin file size: {firmware_size} bytes");

    let mut hasher = Sha256::new();
    let Ok(read) = std::fs::read(file_path) else {
        eprintln!("Error: Could not read file '{}'", file_path.display());
        return READ_FAILED;
    };
    hasher.update(&read);

    let firmware_sha256 = hasher.finalize();
    println!("Firmware SHA-256: {firmware_sha256:x}");

    // We could read an u32 from an argument if we want to support multiple OTA types...
    let ota_type = tlv::OTA_TYPE_VALUE_SSH_STAMP;
    println!("OTA Type Number: {ota_type} (SSH-Stamp)");

    let mut ota_file_path = file_path.to_path_buf();
    ota_file_path.set_extension("ota");

    println!("Saving OTA file to: {}", ota_file_path.display());

    let Ok(mut ota_file) = std::fs::File::create(&ota_file_path) else {
        eprintln!(
            "Error: Could not create OTA file '{}'",
            ota_file_path.display(),
        );
        return CREATE_FAILED;
    };

    // More than enough for the header
    let mut buf = [0u8; 512];

    let header_len =
        OtaHeader::new(ota_type, firmware_sha256.as_slice(), firmware_size).serialize(&mut buf);

    println!("OTA header length: {header_len} bytes");

    let Ok(bytes) = ota_file.write(&buf[..header_len]) else {
        eprintln!(
            "Error: Could not write to OTA file '{}'",
            ota_file_path.display(),
        );
        return WRITE_FAILED;
    };
    println!("Wrote {bytes} bytes of OTA header");

    let Ok(bytes) = ota_file.write(&read) else {
        eprintln!(
            "Error: Could not write firmware data to OTA file '{}'",
            ota_file_path.display(),
        );
        return WRITE_FAILED;
    };
    println!("Wrote {bytes} bytes of firmware data");

    OK
}