crablock-core 0.1.2

use std::process::Command;

use base64::Engine;
use chrono::Utc;
use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey};

use crate::crypto::compute_sha256;
use crate::error::{CrablockError, Result};
use crate::format::Package;

pub const SIGNATURE_ALGORITHM_ED25519: &str = "ed25519";

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SignatureVerificationStatus {
    NotSigned,
    Verified,
}

pub fn is_signature_required(package: &Package, require_signature_flag: bool) -> bool {
    // The CLI flag and the manifest policy both mean the same thing: unsigned packages should fail.
    require_signature_flag || package.manifest.require_signature.unwrap_or(false)
}

pub fn read_signing_key_from_source(source: &str) -> Result<SigningKey> {
    let raw = read_key_source(source)?;

    match raw.len() {
        32 => {
            let mut key = [0u8; 32];
            key.copy_from_slice(&raw);
            Ok(SigningKey::from_bytes(&key))
        }
        64 => {
            let mut key = [0u8; 64];
            key.copy_from_slice(&raw);
            SigningKey::from_keypair_bytes(&key).map_err(|e| {
                CrablockError::InvalidKey(format!("Invalid Ed25519 keypair bytes: {e}"))
            })
        }
        len => Err(CrablockError::InvalidKey(format!(
            "Signing key must decode to 32-byte seed or 64-byte keypair, got {len} bytes"
        ))),
    }
}

pub fn read_verifying_key_from_source(source: &str) -> Result<VerifyingKey> {
    let raw = read_key_source(source)?;

    match raw.len() {
        32 => {
            let mut key = [0u8; 32];
            key.copy_from_slice(&raw);
            VerifyingKey::from_bytes(&key)
                .map_err(|e| CrablockError::InvalidKey(format!("Invalid Ed25519 pubkey: {e}")))
        }
        64 => {
            let mut key = [0u8; 64];
            key.copy_from_slice(&raw);
            let signing_key = SigningKey::from_keypair_bytes(&key).map_err(|e| {
                CrablockError::InvalidKey(format!("Invalid Ed25519 keypair bytes: {e}"))
            })?;
            Ok(signing_key.verifying_key())
        }
        len => Err(CrablockError::InvalidKey(format!(
            "Public key must decode to 32-byte Ed25519 pubkey or 64-byte keypair, got {len} bytes"
        ))),
    }
}

pub fn public_key_fingerprint(verifying_key: &VerifyingKey) -> String {
    compute_sha256(verifying_key.as_bytes())
}

pub fn sign_package(
    package: &mut Package,
    signing_key: &SigningKey,
    pubkey_id_override: Option<&str>,
) -> Result<()> {
    // Signature metadata is copied to both the package trailer and the manifest.
    // That keeps old and new readers aligned during migration.
    let verifying_key = signing_key.verifying_key();
    let fingerprint = pubkey_id_override
        .map(str::to_string)
        .unwrap_or_else(|| public_key_fingerprint(&verifying_key));
    let signed_at = Utc::now().to_rfc3339();

    package.signature_algorithm = Some(SIGNATURE_ALGORITHM_ED25519.to_string());
    package.signing_pubkey_fingerprint = Some(fingerprint.clone());
    package.manifest.signature_algorithm = Some(SIGNATURE_ALGORITHM_ED25519.to_string());
    package.manifest.signing_pubkey_fingerprint = Some(fingerprint);
    package.manifest.signature_created_at = Some(signed_at);

    let canonical_bytes = package.canonical_signing_bytes()?;
    let signature = signing_key.sign(&canonical_bytes);
    package.signature = Some(signature.to_bytes().to_vec());

    Ok(())
}

pub fn verify_package_signature(
    package: &Package,
    verifying_key: &VerifyingKey,
    require_signature: bool,
) -> Result<SignatureVerificationStatus> {
    let signature_bytes = match package.signature.as_ref() {
        Some(signature) => signature,
        None => {
            if require_signature {
                return Err(CrablockError::SignatureMissing);
            }
            return Ok(SignatureVerificationStatus::NotSigned);
        }
    };

    let algorithm = package
        .signature_algorithm
        .as_deref()
        .or(package.manifest.signature_algorithm.as_deref())
        .unwrap_or(SIGNATURE_ALGORITHM_ED25519);
    if algorithm != SIGNATURE_ALGORITHM_ED25519 {
        return Err(CrablockError::UnsupportedSignatureAlgorithm(
            algorithm.to_string(),
        ));
    }

    if signature_bytes.len() != 64 {
        return Err(CrablockError::SignatureInvalid);
    }
    let signature =
        Signature::from_slice(signature_bytes).map_err(|_| CrablockError::SignatureInvalid)?;

    // We verify the exact bytes that were signed at package creation time.
    let canonical_bytes = package.canonical_signing_bytes()?;
    verifying_key
        .verify(&canonical_bytes, &signature)
        .map_err(|_| CrablockError::SignatureInvalid)?;

    let expected_fingerprint = package
        .signing_pubkey_fingerprint
        .as_deref()
        .or(package.manifest.signing_pubkey_fingerprint.as_deref());
    if let Some(expected) = expected_fingerprint {
        let actual = public_key_fingerprint(verifying_key);
        if expected != actual {
            return Err(CrablockError::SignatureInvalid);
        }
    }

    Ok(SignatureVerificationStatus::Verified)
}

fn read_key_source(source: &str) -> Result<Vec<u8>> {
    // Signing keys and public keys support the same source prefixes as encryption keys.
    if let Some(var) = source.strip_prefix("env:") {
        let value = std::env::var(var)
            .map_err(|_| CrablockError::KeySource(format!("Environment variable {var} not set")))?;
        return decode_key_text(&value);
    }

    if let Some(path) = source.strip_prefix("file:") {
        let bytes = std::fs::read(path).map_err(|e| {
            CrablockError::KeySource(format!("Failed to read key file {path}: {e}"))
        })?;
        return decode_key_bytes(bytes);
    }

    if let Some(cmd) = source.strip_prefix("cmd:") {
        let output = Command::new("sh")
            .arg("-c")
            .arg(cmd)
            .output()
            .map_err(|e| CrablockError::KeySource(format!("Failed to execute key command: {e}")))?;

        if !output.status.success() {
            return Err(CrablockError::KeySource(format!(
                "Key command failed with exit code: {:?}",
                output.status.code()
            )));
        }
        return decode_key_bytes(output.stdout);
    }

    decode_key_text(source)
}

fn decode_key_bytes(bytes: Vec<u8>) -> Result<Vec<u8>> {
    match String::from_utf8(bytes.clone()) {
        Ok(value) => decode_key_text(&value),
        Err(_) => Ok(bytes),
    }
}

fn decode_key_text(value: &str) -> Result<Vec<u8>> {
    // Accept explicit prefixes first, then try common raw forms.
    let trimmed = value.trim();
    if trimmed.is_empty() {
        return Err(CrablockError::InvalidKey(
            "Key material cannot be empty".to_string(),
        ));
    }

    if let Some(raw) = trimmed.strip_prefix("hex:") {
        return hex::decode(raw)
            .map_err(|e| CrablockError::InvalidKey(format!("Invalid hex: {e}")));
    }

    if let Some(raw) = trimmed.strip_prefix("base64:") {
        return base64::engine::general_purpose::STANDARD
            .decode(raw)
            .map_err(|e| CrablockError::InvalidKey(format!("Invalid base64: {e}")));
    }

    if trimmed.len().is_multiple_of(2) && trimmed.chars().all(|ch| ch.is_ascii_hexdigit()) {
        if let Ok(decoded) = hex::decode(trimmed) {
            return Ok(decoded);
        }
    }

    if let Ok(decoded) = base64::engine::general_purpose::STANDARD.decode(trimmed) {
        return Ok(decoded);
    }

    Ok(trimmed.as_bytes().to_vec())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::crypto::EncryptionAlgorithm;
    use crate::manifest::Manifest;

    #[test]
    fn test_sign_and_verify_roundtrip() {
        let signing_key = SigningKey::from_bytes(&[42u8; 32]);
        let mut package = Package::new(
            Manifest::new(
                "test".to_string(),
                4,
                EncryptionAlgorithm::Aes256Gcm,
                &[0u8; 12],
                "artifact_hash",
                "payload_hash",
            ),
            vec![1, 2, 3, 4],
            None,
        );

        sign_package(&mut package, &signing_key, None).unwrap();
        let status =
            verify_package_signature(&package, &signing_key.verifying_key(), true).unwrap();

        assert_eq!(status, SignatureVerificationStatus::Verified);
    }

    #[test]
    fn test_tampered_payload_invalidates_signature() {
        let signing_key = SigningKey::from_bytes(&[11u8; 32]);
        let mut package = Package::new(
            Manifest::new(
                "test".to_string(),
                4,
                EncryptionAlgorithm::Aes256Gcm,
                &[0u8; 12],
                "artifact_hash",
                "payload_hash",
            ),
            vec![1, 2, 3, 4],
            None,
        );

        sign_package(&mut package, &signing_key, None).unwrap();
        package.payload[0] ^= 0xFF;

        let result = verify_package_signature(&package, &signing_key.verifying_key(), true);
        assert!(matches!(result, Err(CrablockError::SignatureInvalid)));
    }

    #[test]
    fn test_unsigned_package_require_signature_behavior() {
        let signing_key = SigningKey::from_bytes(&[7u8; 32]);
        let package = Package::new(
            Manifest::new(
                "test".to_string(),
                4,
                EncryptionAlgorithm::Aes256Gcm,
                &[0u8; 12],
                "artifact_hash",
                "payload_hash",
            ),
            vec![1, 2, 3, 4],
            None,
        );

        let optional =
            verify_package_signature(&package, &signing_key.verifying_key(), false).unwrap();
        assert_eq!(optional, SignatureVerificationStatus::NotSigned);

        let required = verify_package_signature(&package, &signing_key.verifying_key(), true);
        assert!(matches!(required, Err(CrablockError::SignatureMissing)));
    }

    #[test]
    fn test_read_signing_key_from_base64_prefix() {
        let source = "base64:KioqKioqKioqKioqKioqKioqKioqKioqKioqKioqKio=";
        let key = read_signing_key_from_source(source).unwrap();
        assert_eq!(key.to_bytes(), [42u8; 32]);
    }
}