use sha2::{Sha256, Digest};

fn hash_chunk(data: &[u8]) -> Vec<u8> {
    Sha256::digest(data).to_vec()
}

fn combine_hashes(left: &[u8], right: &[u8]) -> Vec<u8> {
    let mut hasher = Sha256::new();
    hasher.update(left);
    hasher.update(right);
    hasher.finalize().to_vec()
}

streamlit run ev_firmware_sim.py

use anyhow::{Context, Result};
use ed25519_dalek::{Signer, SigningKey, Verifier, VerifyingKey, Signature};
use rand::rngs::OsRng;
use sha2::{Digest, Sha256};
use std::fs::File;
use std::io::{Read, Write};
use std::path::Path;

const CHUNK_SIZE: usize = 4 * 1024 * 1024; // 4 MB

type Hash32 = [u8; 32];

fn sha256_bytes(data: &[u8]) -> Hash32 {
    let mut h = Sha256::new();
    h.update(data);
    h.finalize().into()
}

fn hash_pair(left: &Hash32, right: &Hash32) -> Hash32 {
    let mut h = Sha256::new();
    h.update(left);
    h.update(right);
    h.finalize().into()
}

fn leaf_hashes_from_file(path: &str, chunk_size: usize) -> Result<Vec<Hash32>> {
    let mut f = File::open(path).with_context(|| format!("open {}", path))?;
    let mut buf = vec![0u8; chunk_size];
    let mut hashes = Vec::new();
    loop {
        let n = f.read(&mut buf)?;
        if n == 0 { break; }
        hashes.push(sha256_bytes(&buf[..n]));
    }
    Ok(hashes)
}

fn build_merkle_root(mut level: Vec<Hash32>) -> Hash32 {
    if level.is_empty() { return sha256_bytes(b""); }
    while level.len() > 1 {
        if level.len() % 2 == 1 {
            let last = *level.last().unwrap();
            level.push(last);
        }
        let mut next = Vec::with_capacity(level.len()/2);
        for i in (0..level.len()).step_by(2) {
            next.push(hash_pair(&level[i], &level[i+1]));
        }
        level = next;
    }
    level[0]
}

fn gen_proof(mut level: Vec<Hash32>, mut idx: usize) -> Vec<Hash32> {
    let mut proof = Vec::new();
    while level.len() > 1 {
        if level.len() % 2 == 1 {
            let last = *level.last().unwrap();
            level.push(last);
        }
        let mut next = Vec::with_capacity(level.len()/2);
        for i in (0..level.len()).step_by(2) {
            let left = level[i];
            let right = level[i+1];
            if i == idx || i+1 == idx {
                proof.push(if i == idx { right } else { left });
                idx = next.len();
            }
            next.push(hash_pair(&left, &right));
        }
        level = next;
    }
    proof
}

fn verify_proof(leaf: &Hash32, mut idx: usize, proof: &[Hash32], expected_root: &Hash32) -> bool {
    let mut cur = *leaf;
    for sib in proof {
        cur = if idx % 2 == 0 { hash_pair(&cur, sib) } else { hash_pair(sib, &cur) };
        idx /= 2;
    }
    &cur == expected_root
}

fn ensure_demo_file(path: &str, size_mb: usize) -> Result<()> {
    if Path::new(path).exists() { return Ok(()); }
    println!("Creating demo firmware: {} ({} MB)", path, size_mb);
    let mut f = File::create(path)?;
    let block = 1024; // 1 KB
    for i in 0..(size_mb * 1024) {
        let mut buf = vec![0u8; block];
        for (j, b) in buf.iter_mut().enumerate() { *b = ((i + j) % 251) as u8; }
        f.write_all(&buf)?;
    }
    Ok(())
}

fn main() -> Result<()> {
    // 1) Demo artifact
    let demo = "demo_firmware.bin";
    ensure_demo_file(demo, 16)?; // 16 MB

    // 2) Leaves and root
    let leaves = leaf_hashes_from_file(demo, CHUNK_SIZE)?;
    println!("Leaf count: {}", leaves.len());
    let root = build_merkle_root(leaves.clone());
    println!("Merkle root: {}", hex::encode(root));

    // 3) Root signing (toy keys; in prod these come from HSM/TPM)
    let mut csprng = OsRng;
    let sk = SigningKey::generate(&mut csprng);
    let vk: VerifyingKey = sk.verifying_key();

    let metadata = b"v=1.2.3;ts=2025-08-18"; // include anti-rollback data
    let mut to_sign = Vec::new();
    to_sign.extend_from_slice(&root);
    to_sign.extend_from_slice(metadata);
    let sig: Signature = sk.sign(&to_sign);
    vk.verify(&to_sign, &sig).expect("signature must verify");
    println!("Signed root with metadata: {} bytes", to_sign.len());

    // 4) Simulate a changed chunk and proof verification against old root
    let changed_idx = 0usize.min(leaves.len()-1);
    let proof = gen_proof(leaves.clone(), changed_idx);
    let ok = verify_proof(&leaves[changed_idx], changed_idx, &proof, &root);
    println!("Proof OK against current root? {}", ok);
    assert!(ok);

    // Pretend chunk changed: new leaf hash fails against old root
    let mut mutated = leaves[changed_idx];
    mutated[0] ^= 0xFF;
    let ok2 = verify_proof(&mutated, changed_idx, &proof, &root);
    println!("Proof OK after mutation (should be false)? {}", ok2);
    assert!(!ok2);

    Ok(())
}