crablock-core 0.1.0

use std::path::Path;

use crate::crypto::compute_sha256;
use crate::error::{CrablockError, Result};
use crate::manifest::Manifest;

pub const MAGIC_BYTES: &[u8] = b"ENCPKG1";
pub const CURRENT_VERSION: u16 = 2;
pub const LEGACY_VERSION: u16 = 1;
// Header = magic bytes + version + manifest length.
pub const HEADER_SIZE: usize = MAGIC_BYTES.len() + 2 + 4; // magic + version + manifest_len
pub const PACKAGE_EXTENSION: &str = "crablock";

pub fn validate_package_path(path: &Path) -> Result<()> {
    if path.extension().and_then(|ext| ext.to_str()) == Some(PACKAGE_EXTENSION) {
        return Ok(());
    }

    Err(CrablockError::InvalidPackageExtension {
        path: path.display().to_string(),
        expected: format!(".{PACKAGE_EXTENSION}"),
    })
}

#[derive(Debug, Clone)]
pub struct PackageHeader {
    pub magic: Vec<u8>,
    pub version: u16,
    pub manifest_len: u32,
}

impl PackageHeader {
    pub fn new(manifest_len: u32) -> Self {
        Self {
            magic: MAGIC_BYTES.to_vec(),
            version: CURRENT_VERSION,
            manifest_len,
        }
    }

    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        if bytes.len() < HEADER_SIZE {
            return Err(CrablockError::InvalidFormat("Header too short".to_string()));
        }

        let magic = bytes[0..MAGIC_BYTES.len()].to_vec();
        if magic != MAGIC_BYTES {
            return Err(CrablockError::InvalidMagic);
        }

        let version = u16::from_be_bytes([bytes[MAGIC_BYTES.len()], bytes[MAGIC_BYTES.len() + 1]]);

        // We accept both the current version and the older v1 format.
        // This lets new runtimes keep reading packages created before embedded env support existed.
        if version != CURRENT_VERSION && version != LEGACY_VERSION {
            return Err(CrablockError::VersionMismatch {
                expected: CURRENT_VERSION,
                found: version,
            });
        }

        let manifest_len = u32::from_be_bytes([
            bytes[MAGIC_BYTES.len() + 2],
            bytes[MAGIC_BYTES.len() + 3],
            bytes[MAGIC_BYTES.len() + 4],
            bytes[MAGIC_BYTES.len() + 5],
        ]);

        Ok(Self {
            magic,
            version,
            manifest_len,
        })
    }

    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(HEADER_SIZE);
        bytes.extend_from_slice(&self.magic);
        bytes.extend_from_slice(&self.version.to_be_bytes());
        bytes.extend_from_slice(&self.manifest_len.to_be_bytes());
        bytes
    }
}

#[derive(Debug, Clone)]
pub struct Package {
    pub header: PackageHeader,
    pub manifest: Manifest,
    pub payload: Vec<u8>,
    // This second encrypted payload is used only for embedded `.env` content.
    pub embedded_env_payload: Option<Vec<u8>>,
    pub signature: Option<Vec<u8>>,
    pub signature_algorithm: Option<String>,
    pub signing_pubkey_fingerprint: Option<String>,
}

impl Package {
    pub fn new(manifest: Manifest, payload: Vec<u8>, signature: Option<Vec<u8>>) -> Self {
        let manifest_bytes = manifest.to_cbor().unwrap_or_default();
        let header = PackageHeader::new(manifest_bytes.len() as u32);
        let signature_algorithm = manifest.signature_algorithm.clone();
        let signing_pubkey_fingerprint = manifest.signing_pubkey_fingerprint.clone();

        Self {
            header,
            manifest,
            payload,
            embedded_env_payload: None,
            signature,
            signature_algorithm,
            signing_pubkey_fingerprint,
        }
    }

    pub fn to_bytes(&self) -> Result<Vec<u8>> {
        let manifest_bytes = self.manifest.to_cbor()?;
        let header = PackageHeader {
            magic: MAGIC_BYTES.to_vec(),
            version: self.header.version,
            manifest_len: manifest_bytes.len() as u32,
        };

        let mut bytes = Vec::new();

        // Header
        bytes.extend_from_slice(&header.to_bytes());

        // Manifest (CBOR)
        bytes.extend_from_slice(&manifest_bytes);

        // Payload length (u64)
        bytes.extend_from_slice(&(self.payload.len() as u64).to_be_bytes());

        // Payload
        bytes.extend_from_slice(&self.payload);

        if header.version >= CURRENT_VERSION {
            // V2 packages append the optional encrypted env payload after the main payload.
            // V1 packages skip this section completely.
            let embedded_env_payload = self.embedded_env_payload.as_deref().unwrap_or(&[]);
            bytes.extend_from_slice(&(embedded_env_payload.len() as u64).to_be_bytes());
            bytes.extend_from_slice(embedded_env_payload);
        }

        // Signature (optional)
        if let Some(sig) = &self.signature {
            if sig.len() > u16::MAX as usize {
                return Err(CrablockError::InvalidFormat(format!(
                    "Signature too large: {} bytes",
                    sig.len()
                )));
            }
            bytes.extend_from_slice(&(sig.len() as u16).to_be_bytes());
            bytes.extend_from_slice(sig);
        } else {
            bytes.extend_from_slice(&0u16.to_be_bytes());
        }

        // Signature metadata (optional, appended for backward compatibility)
        Self::push_sized_string(&mut bytes, self.signature_algorithm.as_ref())?;
        Self::push_sized_string(&mut bytes, self.signing_pubkey_fingerprint.as_ref())?;

        Ok(bytes)
    }

    pub fn canonical_signing_bytes(&self) -> Result<Vec<u8>> {
        // Signatures should not sign themselves.
        // We clone the package, remove the signature, and serialize the rest.
        let mut canonical = self.clone();
        canonical.signature = None;
        canonical.to_bytes()
    }

    fn push_sized_string(out: &mut Vec<u8>, value: Option<&String>) -> Result<()> {
        if let Some(value) = value {
            let raw = value.as_bytes();
            if raw.len() > u8::MAX as usize {
                return Err(CrablockError::InvalidFormat(format!(
                    "Metadata field too long: {} bytes",
                    raw.len()
                )));
            }
            out.push(raw.len() as u8);
            out.extend_from_slice(raw);
        } else {
            out.push(0u8);
        }
        Ok(())
    }

    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        if bytes.len() < HEADER_SIZE {
            return Err(CrablockError::InvalidFormat(
                "Package too short".to_string(),
            ));
        }

        // Parse header
        let header = PackageHeader::from_bytes(&bytes[0..HEADER_SIZE])?;

        // Parse manifest
        let manifest_start = HEADER_SIZE;
        let manifest_end = manifest_start + header.manifest_len as usize;
        if manifest_end > bytes.len() {
            return Err(CrablockError::InvalidFormat(
                "Manifest extends beyond package".to_string(),
            ));
        }
        let manifest = Manifest::from_cbor(&bytes[manifest_start..manifest_end])?;

        // Parse payload length
        let payload_len_start = manifest_end;
        if payload_len_start + 8 > bytes.len() {
            return Err(CrablockError::InvalidFormat(
                "Missing payload length".to_string(),
            ));
        }
        let payload_len = u64::from_be_bytes([
            bytes[payload_len_start],
            bytes[payload_len_start + 1],
            bytes[payload_len_start + 2],
            bytes[payload_len_start + 3],
            bytes[payload_len_start + 4],
            bytes[payload_len_start + 5],
            bytes[payload_len_start + 6],
            bytes[payload_len_start + 7],
        ]) as usize;

        // Parse payload
        let payload_start = payload_len_start + 8;
        let payload_end = payload_start + payload_len;
        if payload_end > bytes.len() {
            return Err(CrablockError::InvalidFormat(
                "Payload extends beyond package".to_string(),
            ));
        }
        let payload = bytes[payload_start..payload_end].to_vec();

        let (embedded_env_payload, signature_offset) = if header.version >= CURRENT_VERSION {
            // In v2 we read the extra env section before the signature trailer.
            if payload_end + 8 > bytes.len() {
                return Err(CrablockError::InvalidFormat(
                    "Missing embedded env payload length".to_string(),
                ));
            }

            let env_payload_len = u64::from_be_bytes([
                bytes[payload_end],
                bytes[payload_end + 1],
                bytes[payload_end + 2],
                bytes[payload_end + 3],
                bytes[payload_end + 4],
                bytes[payload_end + 5],
                bytes[payload_end + 6],
                bytes[payload_end + 7],
            ]) as usize;
            let env_payload_start = payload_end + 8;
            let env_payload_end = env_payload_start + env_payload_len;
            if env_payload_end > bytes.len() {
                return Err(CrablockError::InvalidFormat(
                    "Embedded env payload extends beyond package".to_string(),
                ));
            }

            let embedded_env_payload = if env_payload_len > 0 {
                Some(bytes[env_payload_start..env_payload_end].to_vec())
            } else {
                None
            };

            (embedded_env_payload, env_payload_end)
        } else {
            (None, payload_end)
        };

        // Parse signature (optional)
        let mut signature = None;
        let mut signature_algorithm = None;
        let mut signing_pubkey_fingerprint = None;
        if signature_offset + 2 <= bytes.len() {
            let sig_len =
                u16::from_be_bytes([bytes[signature_offset], bytes[signature_offset + 1]]) as usize;
            let sig_start = signature_offset + 2;
            let sig_end = sig_start + sig_len;
            if sig_end > bytes.len() {
                return Err(CrablockError::InvalidFormat(
                    "Signature extends beyond package".to_string(),
                ));
            }

            if sig_len > 0 {
                signature = Some(bytes[sig_start..sig_end].to_vec());
            }

            let mut cursor = sig_end;
            if cursor < bytes.len() {
                let (algorithm, next) = Self::read_sized_string(bytes, cursor, "algorithm")?;
                signature_algorithm = algorithm;
                cursor = next;
            }
            if cursor < bytes.len() {
                let (fingerprint, next) = Self::read_sized_string(bytes, cursor, "fingerprint")?;
                signing_pubkey_fingerprint = fingerprint;
                cursor = next;
            }
            if cursor != bytes.len() {
                return Err(CrablockError::InvalidFormat(
                    "Unexpected trailing bytes after signature metadata".to_string(),
                ));
            }
        }

        let mut manifest = manifest;
        if manifest.signature_algorithm.is_none() {
            manifest.signature_algorithm = signature_algorithm.clone();
        }
        if manifest.signing_pubkey_fingerprint.is_none() {
            manifest.signing_pubkey_fingerprint = signing_pubkey_fingerprint.clone();
        }

        Ok(Self {
            header,
            manifest,
            payload,
            embedded_env_payload,
            signature,
            signature_algorithm,
            signing_pubkey_fingerprint,
        })
    }

    fn read_sized_string(
        bytes: &[u8],
        cursor: usize,
        field_name: &str,
    ) -> Result<(Option<String>, usize)> {
        if cursor >= bytes.len() {
            return Err(CrablockError::InvalidFormat(format!(
                "Missing signature {field_name} length"
            )));
        }
        let len = bytes[cursor] as usize;
        let start = cursor + 1;
        let end = start + len;
        if end > bytes.len() {
            return Err(CrablockError::InvalidFormat(format!(
                "Signature {field_name} metadata extends beyond package"
            )));
        }
        if len == 0 {
            return Ok((None, end));
        }

        let value = std::str::from_utf8(&bytes[start..end]).map_err(|e| {
            CrablockError::InvalidFormat(format!(
                "Signature {field_name} metadata is not UTF-8: {e}"
            ))
        })?;
        Ok((Some(value.to_string()), end))
    }

    pub fn verify_hashes(&self) -> Result<()> {
        // This check protects the encrypted payload before we even try to decrypt it.
        let computed_payload_hash = compute_sha256(&self.payload);
        if computed_payload_hash != self.manifest.payload_hash_sha256 {
            return Err(CrablockError::HashMismatch {
                expected: self.manifest.payload_hash_sha256.clone(),
                computed: computed_payload_hash,
            });
        }

        if let Some(embedded_env) = self.manifest.embedded_env.as_ref() {
            // Embedded env metadata and payload must either both exist or both be absent.
            let embedded_env_payload = self.embedded_env_payload.as_ref().ok_or_else(|| {
                CrablockError::InvalidFormat(
                    "Manifest declares embedded env but package payload is missing".to_string(),
                )
            })?;
            let computed_env_hash = compute_sha256(embedded_env_payload);
            if computed_env_hash != embedded_env.payload_hash_sha256 {
                return Err(CrablockError::HashMismatch {
                    expected: embedded_env.payload_hash_sha256.clone(),
                    computed: computed_env_hash,
                });
            }
        }

        Ok(())
    }

    pub fn write_to_file(&self, path: &Path) -> Result<()> {
        validate_package_path(path)?;
        let bytes = self.to_bytes()?;
        std::fs::write(path, bytes)?;
        Ok(())
    }

    pub fn read_from_file(path: &Path) -> Result<Self> {
        validate_package_path(path)?;
        let bytes = std::fs::read(path)?;
        Self::from_bytes(&bytes)
    }
}

pub struct PackageBuilder {
    // This builder is used mainly in tests and helper code to construct a package step by step.
    manifest: Option<Manifest>,
    payload: Option<Vec<u8>>,
    signature: Option<Vec<u8>>,
    embedded_env_payload: Option<Vec<u8>>,
    signature_algorithm: Option<String>,
    signing_pubkey_fingerprint: Option<String>,
}

impl PackageBuilder {
    pub fn new() -> Self {
        Self {
            manifest: None,
            payload: None,
            signature: None,
            embedded_env_payload: None,
            signature_algorithm: None,
            signing_pubkey_fingerprint: None,
        }
    }

    pub fn manifest(mut self, manifest: Manifest) -> Self {
        self.manifest = Some(manifest);
        self
    }

    pub fn payload(mut self, payload: Vec<u8>) -> Self {
        self.payload = Some(payload);
        self
    }

    pub fn signature(mut self, signature: Vec<u8>) -> Self {
        self.signature = Some(signature);
        self
    }

    pub fn embedded_env_payload(mut self, embedded_env_payload: Vec<u8>) -> Self {
        self.embedded_env_payload = Some(embedded_env_payload);
        self
    }

    pub fn signature_algorithm(mut self, signature_algorithm: String) -> Self {
        self.signature_algorithm = Some(signature_algorithm);
        self
    }

    pub fn signing_pubkey_fingerprint(mut self, fingerprint: String) -> Self {
        self.signing_pubkey_fingerprint = Some(fingerprint);
        self
    }

    pub fn build(self) -> Result<Package> {
        let manifest = self
            .manifest
            .ok_or_else(|| CrablockError::InvalidFormat("Missing manifest".to_string()))?;
        let payload = self
            .payload
            .ok_or_else(|| CrablockError::InvalidFormat("Missing payload".to_string()))?;

        let mut package = Package::new(manifest, payload, self.signature);
        package.embedded_env_payload = self.embedded_env_payload;
        package.signature_algorithm = self.signature_algorithm;
        package.signing_pubkey_fingerprint = self.signing_pubkey_fingerprint;
        Ok(package)
    }
}

impl Default for PackageBuilder {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::crypto::EncryptionAlgorithm;
    use std::path::Path;

    #[test]
    fn test_package_roundtrip() {
        let manifest = Manifest::new(
            "test_app".to_string(),
            100,
            EncryptionAlgorithm::Aes256Gcm,
            &[0u8; 12],
            "artifact_hash",
            "payload_hash",
        );

        let package = Package::new(manifest, vec![1, 2, 3, 4, 5], None);
        let bytes = package.to_bytes().unwrap();
        let restored = Package::from_bytes(&bytes).unwrap();

        assert_eq!(package.manifest.package_id, restored.manifest.package_id);
        assert_eq!(package.payload, restored.payload);
    }

    #[test]
    fn test_invalid_magic() {
        let bytes = vec![0u8; 20];
        let result = Package::from_bytes(&bytes);
        assert!(matches!(result, Err(CrablockError::InvalidMagic)));
    }

    #[test]
    fn test_package_with_signature() {
        let manifest = Manifest::new(
            "test_app".to_string(),
            100,
            EncryptionAlgorithm::ChaCha20Poly1305,
            &[1u8; 12],
            "hash1",
            "hash2",
        );

        let signature = vec![0xAB; 64];
        let mut package = Package::new(manifest, vec![1, 2, 3], Some(signature.clone()));
        package.signature_algorithm = Some("ed25519".to_string());
        package.signing_pubkey_fingerprint = Some("abc123".to_string());
        let bytes = package.to_bytes().unwrap();
        let restored = Package::from_bytes(&bytes).unwrap();

        assert_eq!(restored.signature, Some(signature));
        assert_eq!(restored.signature_algorithm.as_deref(), Some("ed25519"));
        assert_eq!(
            restored.signing_pubkey_fingerprint.as_deref(),
            Some("abc123")
        );
    }

    #[test]
    fn test_package_with_embedded_env_roundtrip() {
        let manifest = Manifest::new(
            "test_app".to_string(),
            100,
            EncryptionAlgorithm::Aes256Gcm,
            &[2u8; 12],
            "hash1",
            "hash2",
        );

        let mut package = Package::new(manifest, vec![1, 2, 3], None);
        package.embedded_env_payload = Some(vec![9, 9, 9]);

        let bytes = package.to_bytes().unwrap();
        let restored = Package::from_bytes(&bytes).unwrap();

        assert_eq!(restored.embedded_env_payload, Some(vec![9, 9, 9]));
    }

    #[test]
    fn test_canonical_bytes_exclude_signature() {
        let manifest = Manifest::new(
            "test_app".to_string(),
            100,
            EncryptionAlgorithm::Aes256Gcm,
            &[0u8; 12],
            "hash1",
            "hash2",
        );
        let mut package = Package::new(manifest.clone(), vec![9, 8, 7], Some(vec![0xAB; 64]));
        package.signature_algorithm = Some("ed25519".to_string());
        package.signing_pubkey_fingerprint = Some("fingerprint".to_string());

        let canonical = package.canonical_signing_bytes().unwrap();

        let mut expected = Package::new(manifest, vec![9, 8, 7], None);
        expected.signature_algorithm = Some("ed25519".to_string());
        expected.signing_pubkey_fingerprint = Some("fingerprint".to_string());
        let expected_bytes = expected.to_bytes().unwrap();

        assert_eq!(canonical, expected_bytes);
    }

    #[test]
    fn test_legacy_signed_package_without_metadata_parses() {
        let manifest = Manifest::new(
            "test_app".to_string(),
            100,
            EncryptionAlgorithm::Aes256Gcm,
            &[0u8; 12],
            "hash1",
            "hash2",
        );
        let payload = vec![1, 2, 3];
        let signature = vec![0xAA; 64];
        let manifest_bytes = manifest.to_cbor().unwrap();
        let header = PackageHeader {
            magic: MAGIC_BYTES.to_vec(),
            version: LEGACY_VERSION,
            manifest_len: manifest_bytes.len() as u32,
        };

        let mut bytes = Vec::new();
        bytes.extend_from_slice(&header.to_bytes());
        bytes.extend_from_slice(&manifest_bytes);
        bytes.extend_from_slice(&(payload.len() as u64).to_be_bytes());
        bytes.extend_from_slice(&payload);
        bytes.extend_from_slice(&(signature.len() as u16).to_be_bytes());
        bytes.extend_from_slice(&signature);

        let restored = Package::from_bytes(&bytes).unwrap();
        assert_eq!(restored.signature, Some(signature));
        assert!(restored.signature_algorithm.is_none());
        assert!(restored.signing_pubkey_fingerprint.is_none());
    }

    #[test]
    fn test_validate_package_path_accepts_crablock_extension() {
        assert!(validate_package_path(Path::new("app.crablock")).is_ok());
    }

    #[test]
    fn test_validate_package_path_rejects_non_crablock_extension() {
        let result = validate_package_path(Path::new("app.pkg"));

        assert!(matches!(
            result,
            Err(CrablockError::InvalidPackageExtension { .. })
        ));
    }
}