v2x-patcher 0.2.0

use std::io::Read;
use std::{process, thread, time};

use clap::Parser;
use sha2::{Digest, Sha256};

const FIRMWARE_URL: &str = "https://files.creative.com/creative/bin/firmware/fwureleases/MF8400/zifs/fc00/MF8400_fc00_21BB4E50.zif";
const FIRMWARE_SHA256: &str = "36ec53321efcb0b9d40c400de07796eccd5ec422e710f8217c30ad34efad1c8c";

const CIFF_MAGIC: &[u8; 4] = b"CIFF";
const FMAIN_DATA_OFFSET: usize = 0x387DC;
const CHK2_MAGIC: &[u8; 4] = b"CHK2";
const WAKEUP_SLEEP: u64 = 10; // secs

// patch: disable BLE CTP bridge
// file_offset = (0x40026a5e - 0x40006520) + 0xEEFC = 0x2F43A
const FMAIN_BLE_BRIDGE_FILE_OFFSET: usize = 0x2F43A;
const FMAIN_BLE_BRIDGE_PATCH_LEN: usize = 74;
const FMAIN_BLE_BRIDGE_ORIGINAL: [u8; 4] = [0x0a, 0x23, 0xa4, 0xf8];

// optional patch: WELCOME -> PATCHED
const FMAIN_WELCOME_OFFSET: usize = 0xB8F00;

// optional patch: replace FMAIN opcode 0x54 handler with multi-cmd mem handler
// file_offset = (ram_addr - 0x40006520) + 0xEEFC
const FMAIN_HANDLER_FILE_OFFSET: usize = 0x2D43E; // (0x40024A62 - 0x40006520) + 0xEEFC
const FMAIN_HANDLER_ORIGINAL: [u8; 4] = [0xc0, 0xef, 0x50, 0x00]; // vmov.i32 q8, #0

#[rustfmt::skip]
const FMAIN_HANDLER: [u8; 96] = [
    // sub-command dispatch
    0x3b, 0x78,               // ldrb    r3, [r7, #0]       ; sub-command
    0x01, 0x2b,               // cmp     r3, #1
    0x0f, 0xd0,               // beq     write              ; +0x26
    0x02, 0x2b,               // cmp     r3, #2
    0x1e, 0xd0,               // beq     exec               ; +0x48

    // sub-cmd 0x00: mem-read
    0xd7, 0xf8, 0x01, 0x00,   // ldr.w   r0, [r7, #1]       ; address
    0x79, 0x79,               // ldrb    r1, [r7, #5]       ; count
    0xad, 0xf8, 0x2a, 0x10,   // strh.w  r1, [sp, #0x2a]    ; response length
    0x2a, 0x46,               // mov     r2, r5             ; dst = response buffer
    0x29, 0xb1,               // cbz     r1, read_done
    0x10, 0xf8, 0x01, 0x3b,   // ldrb    r3, [r0], #1       ; loop: *addr++
    0x02, 0xf8, 0x01, 0x3b,   // strb    r3, [r2], #1       ;       *dst++
    0x01, 0x39,               // subs    r1, #1
    0xf9, 0xd1,               // bne     read_loop
    0x13, 0xe6,               // b.n     0x400246b0         ; exit

    // sub-cmd 0x01: mem-write
    0xd7, 0xf8, 0x01, 0x00,   // ldr.w   r0, [r7, #1]       ; dest address
    0x79, 0x79,               // ldrb    r1, [r7, #5]       ; count
    0xba, 0x1d,               // adds    r2, r7, #6         ; src = payload+6
    0x29, 0xb1,               // cbz     r1, write_done
    0x12, 0xf8, 0x01, 0x3b,   // ldrb    r3, [r2], #1       ; loop: *src++
    0x00, 0xf8, 0x01, 0x3b,   // strb    r3, [r0], #1       ;       *dst++
    0x01, 0x39,               // subs    r1, #1
    0xf9, 0xd1,               // bne     write_loop
    0x00, 0x23,               // movs    r3, #0             ; write_done:
    0x2b, 0x70,               // strb    r3, [r5, #0]       ; response[0] = OK
    0x01, 0x23,               // movs    r3, #1
    0xad, 0xf8, 0x2a, 0x30,   // strh.w  r3, [sp, #0x2a]    ; response length = 1
    0x02, 0xe6,               // b.n     0x400246b0         ; exit

    // sub-cmd 0x02: mem-exec
    0xd7, 0xf8, 0x05, 0x00,   // ldr.w   r0, [r7, #5]       ; r0 argument
    0xd7, 0xf8, 0x01, 0x30,   // ldr.w   r3, [r7, #1]       ; function addr
    0x43, 0xf0, 0x01, 0x03,   // orr.w   r3, r3, #1         ; ensure thumb bit
    0x98, 0x47,               // blx     r3                 ; call function
    0x28, 0x60,               // str     r0, [r5, #0]       ; response = r0
    0x04, 0x23,               // movs    r3, #4
    0xad, 0xf8, 0x2a, 0x30,   // strh.w  r3, [sp, #0x2a]    ; response length = 4
    0xf6, 0xe5,               // b.n     0x400246b0         ; exit
];

/// Firmware patcher for the Creative Sound Blaster Katana V2X.
///
/// Downloads the official firmware, applies security patches, verifies the
/// device model, and uploads the patched firmware.
#[derive(Parser)]
#[command(name = "v2x-patcher")]
struct Cli {
    /// Do not apply the Bluetooth CTP disabling patch
    #[arg(long)]
    no_bt_patch: bool,

    /// Apply mem-read/write/exec debug handler and WELCOME -> PATCHED patches.
    #[arg(long)]
    mem_patches: bool,

    /// Serial port path (e.g. /dev/ttyACM0). Auto-detected if not specified.
    #[arg(short, long)]
    port: Option<String>,

    /// Write patched firmware to file instead of uploading to device.
    #[arg(short, long)]
    output: Option<String>,

    /// Do not prompt to exit at the end.
    #[arg(short, long)]
    yes: bool,
}

fn run(cli: &Cli) -> Result<(), Box<dyn std::error::Error>> {
    println!("Downloading firmware...");
    let firmware = dl_firmware()?;
    println!("  Downloaded {} bytes", firmware.len());

    println!("Verifying SHA-256...");
    let hash = Sha256::digest(&firmware);
    let hash_hex = hex::encode(hash);
    if hash_hex != FIRMWARE_SHA256 {
        return Err(format!(
            "SHA-256 mismatch!\n  expected: {FIRMWARE_SHA256}\n  got:      {hash_hex}"
        )
        .into());
    }

    println!("  SHA-256 OK: {hash_hex}");

    println!("Verifying patch sites...");
    verify_patch_sites(&firmware)?;
    println!("  All patch sites verified clean");

    let mut data = firmware;
    apply_patches(&mut data, !cli.no_bt_patch, cli.mem_patches);
    update_chk2(&mut data);

    if let Some(ref path) = cli.output {
        std::fs::write(path, &data)?;
        println!("Wrote {} bytes to {path}", data.len());
    } else {
        println!("Connecting to device...");
        let mut dev = match &cli.port {
            Some(path) => v2x::KatanaV2X::open(path)?,
            None => v2x::KatanaV2X::open_auto()?,
        };
        println!(
            "  Connected: {}",
            dev.port_name().unwrap_or_else(|| "unknown".into())
        );

        println!("Verifying device model...");
        let model = dev.device_model()?;
        if model != "SB Katana V2X" {
            return Err(format!(
                "Unexpected device model: \"{model}\", expected \"SB Katana V2X\""
            )
            .into());
        }
        println!("  Model: {model}");

        println!("Waiting {WAKEUP_SLEEP} seconds for the device to come out of sleep mode...");
        thread::sleep(time::Duration::from_secs(WAKEUP_SLEEP));

        println!("Uploading firmware ({} bytes)...", data.len());
        dev.upload_firmware(&data, |sent, total| {
            let pct = sent * 100 / total;
            eprint!("\r  Progress: {pct:3}% ({sent}/{total} bytes)");
        })?;
        eprintln!();
        println!("Firmware upload complete. Device will reboot.");
    }

    Ok(())
}

fn verify_patch_sites(data: &[u8]) -> Result<(), Box<dyn std::error::Error>> {
    if &data[..4] != CIFF_MAGIC {
        return Err(format!("Not a CIFF file: {:02x?}", &data[..4]).into());
    }

    // FMAIN data check
    let fmain_check = &data[FMAIN_DATA_OFFSET..FMAIN_DATA_OFFSET + 4];
    if fmain_check != [0x44, 0xf0, 0x9f, 0xe5] {
        return Err(format!(
            "FMAIN data mismatch at 0x{:X}: {}",
            FMAIN_DATA_OFFSET,
            hex::encode(fmain_check)
        )
        .into());
    }

    // verify BLE bridge site
    let off = FMAIN_DATA_OFFSET + FMAIN_BLE_BRIDGE_FILE_OFFSET;
    let actual = &data[off..off + 4];
    if actual != FMAIN_BLE_BRIDGE_ORIGINAL {
        return Err(format!(
            "BLE bridge site has {}, expected {}",
            hex::encode(actual),
            hex::encode(FMAIN_BLE_BRIDGE_ORIGINAL)
        )
        .into());
    }

    // verify CHK2 trailer
    let chk2_offset = data.len() - 40;
    if &data[chk2_offset..chk2_offset + 4] != CHK2_MAGIC {
        return Err("CHK2 trailer not found".into());
    }

    // verify mem patches as well even if they're not applied
    // verify handler site
    let off = FMAIN_DATA_OFFSET + FMAIN_HANDLER_FILE_OFFSET;
    let actual = &data[off..off + 4];
    if actual != FMAIN_HANDLER_ORIGINAL {
        return Err(format!(
            "FMAIN handler site has {}, expected {}",
            hex::encode(actual),
            hex::encode(FMAIN_HANDLER_ORIGINAL)
        )
            .into());
    }

    // verify WELCOME string
    let off = FMAIN_DATA_OFFSET + FMAIN_WELCOME_OFFSET;
    if &data[off..off + 7] != b"WELCOME" {
        return Err("FMAIN WELCOME string not found".into());
    }

    Ok(())
}

fn dl_firmware() -> Result<Vec<u8>, Box<dyn std::error::Error>> {
    let body = ureq::get(FIRMWARE_URL).call()?.into_body();
    let mut firmware = Vec::new();
    body.into_reader().read_to_end(&mut firmware)?;

    Ok(firmware)
}

fn apply_patches(data: &mut [u8], bt_patch: bool, mem_patches: bool) {
    if bt_patch {
        // disable BLE CTP bridge
        let off = FMAIN_DATA_OFFSET + FMAIN_BLE_BRIDGE_FILE_OFFSET;
        let nops: Vec<u8> = vec![0x00, 0xbf]
            .into_iter()
            .cycle()
            .take(FMAIN_BLE_BRIDGE_PATCH_LEN)
            .collect();
        data[off..off + FMAIN_BLE_BRIDGE_PATCH_LEN].copy_from_slice(&nops);
        println!(
            "  [ble-bridge]    CIFF 0x{off:05X}: {FMAIN_BLE_BRIDGE_PATCH_LEN} bytes NOPed (BLE CTP disabled)"
        );
    }

    if mem_patches {
        // replace opcode 0x54 handler
        let off = FMAIN_DATA_OFFSET + FMAIN_HANDLER_FILE_OFFSET;
        data[off..off + FMAIN_HANDLER.len()].copy_from_slice(&FMAIN_HANDLER);
        println!(
            "  [fmain-handler] CIFF 0x{off:05X}: {} bytes (mem-read/write/exec)",
            FMAIN_HANDLER.len()
        );

        // WELCOME -> PATCHED
        let off = FMAIN_DATA_OFFSET + FMAIN_WELCOME_OFFSET;
        data[off..off + 7].copy_from_slice(b"PATCHED");
        println!("  [welcome]       CIFF 0x{off:05X}: WELCOME -> PATCHED");
    }
}

fn get_chk2_offset(data: &[u8]) -> usize {
    data.len() - 40
}

fn get_chk2(data: &[u8]) -> &[u8] {
    let chk2_offset = get_chk2_offset(data);
    &data[chk2_offset + 8..chk2_offset + 40]
}

fn calc_chk2(data: &[u8]) -> sha2::digest::Output<Sha256> {
    let chk2_offset = get_chk2_offset(data);
    Sha256::digest(&data[8..chk2_offset])
}

fn update_chk2(data: &mut [u8]) {
    let chk2_offset = get_chk2_offset(data);

    // CHK2 layout: [CHK2 magic (4)] [unknown (4)] [SHA-256 hash (32)]
    // hash covers CIFF[0x08..chk2_offset]
    let hash = calc_chk2(data);
    let old_hash = hex::encode(get_chk2(data));
    data[chk2_offset + 8..chk2_offset + 40].copy_from_slice(&hash);

    println!("  [CHK2] SHA-256 over CIFF[0x08:0x{chk2_offset:X}]");
    println!("    was: {old_hash}");
    println!("    now: {}", hex::encode(hash));
}

fn main() {
    env_logger::init();
    let cli = Cli::parse();

    let mut code = 0;
    if let Err(e) = run(&cli) {
        code = 1;
        eprintln!("Error: {e}");
    }

    if !cli.yes {
        eprint!("Press enter to exit...");
        let _ = std::io::stdin().read_line(&mut String::new());
    }

    process::exit(code);
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::LazyLock;

    static FW: LazyLock<Vec<u8>> = LazyLock::new(|| {
        dl_firmware().unwrap()
    });

    #[test]
    fn test_chk2_get_calc_identical() {
        let old_hash = get_chk2(&FW);
        let calc_hash = calc_chk2(&FW);

        assert_eq!(old_hash, calc_hash.as_slice());
    }

    #[test]
    fn test_update_chk2_produce_identical_fw() {
        let mut data = FW.clone();
        update_chk2(&mut data);

        assert_eq!(FW.as_slice(), data);
    }

    #[test]
    fn test_all_patches_correct_chk2() {
        let mut data = FW.clone();
        apply_patches(&mut data, true, true);
        update_chk2(&mut data);
        let hash = get_chk2(&data);

        assert_eq!(
            hex::encode(hash),
            "38f2892d4b5f1af1e5b03d42437ffbb3d1b6bef2d19611ae6cc7e99342f7ec97"
        );
    }

    #[test]
    fn test_no_patches_produce_identical_fw() {
        let mut data = FW.clone();
        apply_patches(&mut data, false, false);
        update_chk2(&mut data);
        let hash = Sha256::digest(&data);

        assert_eq!(
            hex::encode(hash),
            FIRMWARE_SHA256
        );
    }
}