geist_supervisor 0.1.28

use anyhow::{Context, Result};
use ed25519_dalek::{Signature, Verifier, VerifyingKey};
use sha2::{Digest, Sha256};
use std::path::Path;

const OTA_PUBLIC_KEY: &[u8; 32] = include_bytes!("../../keys/ota_signing_key.pub.raw");

/// Whether to hard-fail on missing signatures.
/// Set to false during transition (old unsigned releases still exist).
/// Flip to true once all releases are signed.
pub const REQUIRE_SIGNATURE: bool = false;

/// Verify a bundle's Ed25519 signature using the embedded public key.
pub fn verify_bundle(bundle_path: &Path, sig_path: &Path) -> Result<()> {
    verify_bundle_with_key(bundle_path, sig_path, OTA_PUBLIC_KEY)
}

/// Verify a bundle's Ed25519 signature using a provided public key.
/// The signature is over the SHA-256 digest of the bundle file.
pub fn verify_bundle_with_key(
    bundle_path: &Path,
    sig_path: &Path,
    public_key_bytes: &[u8; 32],
) -> Result<()> {
    tracing::info!("Verifying bundle signature: {}", bundle_path.display());

    let verifying_key =
        VerifyingKey::from_bytes(public_key_bytes).context("Invalid embedded public key")?;

    // SHA-256 of the bundle (matches what release.sh signs)
    let bundle_bytes = std::fs::read(bundle_path).context("Failed to read bundle")?;
    let digest = Sha256::digest(&bundle_bytes);

    // Load 64-byte Ed25519 signature
    let sig_bytes = std::fs::read(sig_path).context("Failed to read signature")?;
    let signature = Signature::from_slice(&sig_bytes)
        .context("Invalid signature format (expected 64 bytes)")?;

    verifying_key
        .verify(digest.as_slice(), &signature)
        .map_err(|e| anyhow::anyhow!("Signature verification FAILED: {}", e))?;

    tracing::info!("Signature verified OK");
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use ed25519_dalek::SigningKey;
    use sha2::{Digest, Sha256};
    use std::fs;
    use tempfile::tempdir;

    fn test_keypair() -> (SigningKey, VerifyingKey) {
        let signing_key = SigningKey::from_bytes(&[1u8; 32]);
        let verifying_key = signing_key.verifying_key();
        (signing_key, verifying_key)
    }

    fn sign_file(signing_key: &SigningKey, file_path: &Path) -> Vec<u8> {
        use ed25519_dalek::Signer;
        let content = fs::read(file_path).unwrap();
        let digest = Sha256::digest(&content);
        let sig = signing_key.sign(digest.as_slice());
        sig.to_bytes().to_vec()
    }

    #[test]
    fn test_valid_signature() {
        let dir = tempdir().unwrap();
        let bundle_path = dir.path().join("bundle.tar.gz");
        let sig_path = dir.path().join("bundle.tar.gz.sig");

        fs::write(&bundle_path, b"test bundle content").unwrap();

        let (signing_key, verifying_key) = test_keypair();
        let sig = sign_file(&signing_key, &bundle_path);
        fs::write(&sig_path, &sig).unwrap();

        let pubkey_bytes: &[u8; 32] = verifying_key.as_bytes();
        let result = verify_bundle_with_key(&bundle_path, &sig_path, pubkey_bytes);
        assert!(
            result.is_ok(),
            "Valid signature should verify: {:?}",
            result
        );
    }

    #[test]
    fn test_corrupted_signature_fails() {
        let dir = tempdir().unwrap();
        let bundle_path = dir.path().join("bundle.tar.gz");
        let sig_path = dir.path().join("bundle.tar.gz.sig");

        fs::write(&bundle_path, b"test bundle content").unwrap();

        let (signing_key, verifying_key) = test_keypair();
        let mut sig = sign_file(&signing_key, &bundle_path);
        sig[0] ^= 0xff; // corrupt one byte
        fs::write(&sig_path, &sig).unwrap();

        let pubkey_bytes: &[u8; 32] = verifying_key.as_bytes();
        let result = verify_bundle_with_key(&bundle_path, &sig_path, pubkey_bytes);
        assert!(result.is_err(), "Corrupted signature should fail");
    }

    #[test]
    fn test_wrong_key_fails() {
        let dir = tempdir().unwrap();
        let bundle_path = dir.path().join("bundle.tar.gz");
        let sig_path = dir.path().join("bundle.tar.gz.sig");

        fs::write(&bundle_path, b"test bundle content").unwrap();

        let (signing_key, _) = test_keypair();
        let sig = sign_file(&signing_key, &bundle_path);
        fs::write(&sig_path, &sig).unwrap();

        // Different key
        let wrong_key = SigningKey::from_bytes(&[2u8; 32]);
        let wrong_verifying_key = wrong_key.verifying_key();
        let wrong_pubkey: &[u8; 32] = wrong_verifying_key.as_bytes();
        let result = verify_bundle_with_key(&bundle_path, &sig_path, wrong_pubkey);
        assert!(result.is_err(), "Wrong key should fail verification");
    }

    #[test]
    fn test_tampered_bundle_fails() {
        let dir = tempdir().unwrap();
        let bundle_path = dir.path().join("bundle.tar.gz");
        let sig_path = dir.path().join("bundle.tar.gz.sig");

        fs::write(&bundle_path, b"original content").unwrap();

        let (signing_key, verifying_key) = test_keypair();
        let sig = sign_file(&signing_key, &bundle_path);
        fs::write(&sig_path, &sig).unwrap();

        // Tamper with the bundle after signing
        fs::write(&bundle_path, b"tampered content").unwrap();

        let pubkey_bytes: &[u8; 32] = verifying_key.as_bytes();
        let result = verify_bundle_with_key(&bundle_path, &sig_path, pubkey_bytes);
        assert!(result.is_err(), "Tampered bundle should fail verification");
    }

    #[test]
    fn test_missing_sig_file() {
        let dir = tempdir().unwrap();
        let bundle_path = dir.path().join("bundle.tar.gz");
        let sig_path = dir.path().join("nonexistent.sig");

        fs::write(&bundle_path, b"test content").unwrap();

        let (_, verifying_key) = test_keypair();
        let pubkey_bytes: &[u8; 32] = verifying_key.as_bytes();
        let result = verify_bundle_with_key(&bundle_path, &sig_path, pubkey_bytes);
        assert!(result.is_err(), "Missing sig file should fail");
    }

    #[test]
    fn test_truncated_signature() {
        let dir = tempdir().unwrap();
        let bundle_path = dir.path().join("bundle.tar.gz");
        let sig_path = dir.path().join("bundle.tar.gz.sig");

        fs::write(&bundle_path, b"test content").unwrap();
        fs::write(&sig_path, b"too short").unwrap();

        let (_, verifying_key) = test_keypair();
        let pubkey_bytes: &[u8; 32] = verifying_key.as_bytes();
        let result = verify_bundle_with_key(&bundle_path, &sig_path, pubkey_bytes);
        assert!(result.is_err(), "Truncated signature should fail");
    }

    #[test]
    fn test_empty_bundle() {
        let dir = tempdir().unwrap();
        let bundle_path = dir.path().join("bundle.tar.gz");
        let sig_path = dir.path().join("bundle.tar.gz.sig");

        fs::write(&bundle_path, b"").unwrap();

        let (signing_key, verifying_key) = test_keypair();
        let sig = sign_file(&signing_key, &bundle_path);
        fs::write(&sig_path, &sig).unwrap();

        let pubkey_bytes: &[u8; 32] = verifying_key.as_bytes();
        let result = verify_bundle_with_key(&bundle_path, &sig_path, pubkey_bytes);
        assert!(
            result.is_ok(),
            "Empty bundle with valid signature should verify"
        );
    }
}