m-bus-core 0.5.0

use crate::{DeviceType, ManufacturerCode, SecurityMode};

#[cfg(feature = "decryption")]
use aes::Aes128Dec;
#[cfg(feature = "decryption")]
use cbc::{
    Decryptor,
    cipher::{BlockModeDecrypt, KeyIvInit},
};

#[derive(Debug, Clone, PartialEq)]
pub struct KeyContext {
    pub manufacturer: ManufacturerCode,
    pub identification_number: u32,
    pub version: u8,
    pub device_type: DeviceType,
    pub security_mode: SecurityMode,
    pub access_number: u8,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum DecryptionError {
    UnsupportedMode(SecurityMode),
    KeyNotFound,
    DecryptionFailed,
    InvalidKeyLength,
    InvalidDataLength,
    NotEncrypted,
    UnknownEncryptionState,
}

impl core::fmt::Display for DecryptionError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::UnsupportedMode(mode) => write!(f, "Unsupported security mode: {:?}", mode),
            Self::KeyNotFound => write!(f, "Decryption key not found"),
            Self::DecryptionFailed => write!(f, "Decryption operation failed"),
            Self::InvalidKeyLength => write!(f, "Invalid key length"),
            Self::InvalidDataLength => write!(f, "Invalid data length"),
            Self::NotEncrypted => write!(f, "Data is not encrypted"),
            Self::UnknownEncryptionState => {
                write!(f, "Unknown encryption state for this data block type")
            }
        }
    }
}

impl core::error::Error for DecryptionError {}

pub trait KeyProvider {
    fn get_key(&self, context: &KeyContext) -> Result<&[u8], DecryptionError>;
}

#[derive(Debug)]
pub struct EncryptedPayload<'a> {
    pub data: &'a [u8],
    pub context: KeyContext,
}

impl<'a> EncryptedPayload<'a> {
    pub fn new(data: &'a [u8], context: KeyContext) -> Self {
        Self { data, context }
    }

    /// Return plaintext bytes lazily without a payload-sized destination buffer.
    #[cfg(feature = "decryption")]
    pub fn decrypted_bytes<K: KeyProvider>(
        &self,
        provider: &K,
    ) -> Result<DecryptedBytes<'a>, DecryptionError> {
        let key = provider.get_key(&self.context)?;
        match self.context.security_mode {
            SecurityMode::NoEncryption => Ok(DecryptedBytes::plaintext(self.data)),
            SecurityMode::AesCbc128IvZero => DecryptedBytes::cbc(self.data, key, &[0u8; 16]),
            SecurityMode::AesCbc128IvNonZero => {
                DecryptedBytes::cbc(self.data, key, &self._derive_iv())
            }
            mode => Err(DecryptionError::UnsupportedMode(mode)),
        }
    }

    /// Decrypt into caller-provided contiguous storage.
    #[cfg(feature = "decryption")]
    pub fn decrypt_into<K: KeyProvider>(
        &self,
        provider: &K,
        output: &mut [u8],
    ) -> Result<usize, DecryptionError> {
        let key = provider.get_key(&self.context)?;

        match self.context.security_mode {
            SecurityMode::NoEncryption => {
                let len = self.data.len();
                let dest = output
                    .get_mut(..len)
                    .ok_or(DecryptionError::InvalidDataLength)?;
                dest.copy_from_slice(self.data);
                Ok(len)
            }
            SecurityMode::AesCbc128IvZero => {
                decrypt_aes_cbc_into(self.data, key, &[0u8; 16], output)
            }
            SecurityMode::AesCbc128IvNonZero => {
                let iv = self._derive_iv();
                decrypt_aes_cbc_into(self.data, key, &iv, output)
            }
            mode => Err(DecryptionError::UnsupportedMode(mode)),
        }
    }

    #[cfg(feature = "decryption")]
    fn _derive_iv(&self) -> [u8; 16] {
        let mut iv = [0u8; 16];
        // Bytes 0-1: Manufacturer ID (numeric, little-endian)
        let mfr_id = self.context.manufacturer.to_id();
        iv[0..2].copy_from_slice(&mfr_id.to_le_bytes());
        // Bytes 2-5: Identification number as BCD bytes (how they appear in frame)
        let bcd_bytes = decimal_to_bcd(self.context.identification_number);
        iv[2..6].copy_from_slice(&bcd_bytes);
        // Byte 6: Version
        iv[6] = self.context.version;
        // Byte 7: Device type
        iv[7] = self.context.device_type.into();
        // Bytes 8-15: Access number repeated 8 times
        iv[8..16].fill(self.context.access_number);
        iv
    }
}

/// Convert a decimal number to BCD bytes (little-endian, 4 bytes)
/// e.g., 14639203 -> [0x03, 0x92, 0x63, 0x14]
#[cfg(feature = "decryption")]
fn decimal_to_bcd(mut value: u32) -> [u8; 4] {
    let mut bcd = [0u8; 4];
    for byte in &mut bcd {
        let low = (value % 10) as u8;
        value /= 10;
        let high = (value % 10) as u8;
        value /= 10;
        *byte = (high << 4) | low;
    }
    bcd
}

pub struct StaticKeyProvider<const N: usize> {
    entries: [(u64, [u8; 16]); N],
    count: usize,
}

impl<const N: usize> Default for StaticKeyProvider<N> {
    fn default() -> Self {
        Self {
            entries: [(0, [0u8; 16]); N],
            count: 0,
        }
    }
}

impl<const N: usize> StaticKeyProvider<N> {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn add_key(
        &mut self,
        manufacturer_id: u16,
        identification_number: u32,
        key: [u8; 16],
    ) -> Result<(), DecryptionError> {
        let key_id = Self::compute_key_id(manufacturer_id, identification_number);
        let entry = self
            .entries
            .get_mut(self.count)
            .ok_or(DecryptionError::InvalidDataLength)?;
        *entry = (key_id, key);
        self.count += 1;
        Ok(())
    }

    const fn compute_key_id(manufacturer_id: u16, identification_number: u32) -> u64 {
        ((manufacturer_id as u64) << 32) | (identification_number as u64)
    }
}

impl<const N: usize> KeyProvider for StaticKeyProvider<N> {
    fn get_key(&self, context: &KeyContext) -> Result<&[u8], DecryptionError> {
        let manufacturer_id = context.manufacturer.to_id();

        let key_id = Self::compute_key_id(manufacturer_id, context.identification_number);

        self.entries[..self.count]
            .iter()
            .find(|(id, _)| *id == key_id)
            .map(|(_, key)| key.as_slice())
            .ok_or(DecryptionError::KeyNotFound)
    }
}

/// Lazily decrypted payload bytes, borrowing the original ciphertext.
///
/// AES-CBC retains its key schedule, chaining state, and one 16-byte working
/// block, not a payload-sized plaintext buffer. Trailing incomplete blocks are
/// yielded unchanged, matching [`EncryptedPayload::decrypt_into`].
///
/// Contiguous-slice parsers still require a destination buffer; callers that
/// consume bytes incrementally can use this iterator directly.
#[cfg(feature = "decryption")]
pub struct DecryptedBytes<'a> {
    input: &'a [u8],
    decryptor: Option<Decryptor<Aes128Dec>>,
    block: cipher::Block<Aes128Dec>,
    position: usize,
    remaining: usize,
}

#[cfg(feature = "decryption")]
impl<'a> DecryptedBytes<'a> {
    fn plaintext(input: &'a [u8]) -> Self {
        Self {
            input,
            decryptor: None,
            block: Default::default(),
            position: 16,
            remaining: input.len(),
        }
    }

    fn cbc(input: &'a [u8], key: &[u8], iv: &[u8]) -> Result<Self, DecryptionError> {
        if key.len() != 16 {
            return Err(DecryptionError::InvalidKeyLength);
        }
        if iv.len() != 16 || input.is_empty() {
            return Err(DecryptionError::InvalidDataLength);
        }
        let decryptor = Decryptor::<Aes128Dec>::new_from_slices(key, iv)
            .map_err(|_| DecryptionError::InvalidKeyLength)?;
        Ok(Self {
            decryptor: Some(decryptor),
            ..Self::plaintext(input)
        })
    }
}

#[cfg(feature = "decryption")]
impl Iterator for DecryptedBytes<'_> {
    type Item = u8;

    fn next(&mut self) -> Option<Self::Item> {
        if self.remaining == 0 {
            return None;
        }
        if self.position == 16
            && let Some(decryptor) = &mut self.decryptor
            && let Some(block) = self.input.get(..16)
        {
            self.block.copy_from_slice(block);
            decryptor.decrypt_block(&mut self.block);
            self.input = &self.input[16..];
            self.position = 0;
        }
        let byte = if let Some(byte) = self.block.get(self.position) {
            self.position += 1;
            *byte
        } else {
            let (byte, rest) = self.input.split_first()?;
            self.input = rest;
            *byte
        };
        self.remaining -= 1;
        Some(byte)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.remaining, Some(self.remaining))
    }
}

#[cfg(feature = "decryption")]
impl ExactSizeIterator for DecryptedBytes<'_> {}
#[cfg(feature = "decryption")]
impl core::iter::FusedIterator for DecryptedBytes<'_> {}

// Keep the AES key schedule off the unencrypted caller's stack.
#[cfg(feature = "decryption")]
#[inline(never)]
fn decrypt_aes_cbc_into(
    data: &[u8],
    key: &[u8],
    iv: &[u8],
    output: &mut [u8],
) -> Result<usize, DecryptionError> {
    if key.len() != 16 {
        return Err(DecryptionError::InvalidKeyLength);
    }
    if iv.len() != 16 || data.is_empty() {
        return Err(DecryptionError::InvalidDataLength);
    }
    let target = output
        .get_mut(..data.len())
        .ok_or(DecryptionError::InvalidDataLength)?;
    target.copy_from_slice(data);
    let encrypted_len = data.len() / 16 * 16;
    if encrypted_len != 0 {
        let decryptor = Decryptor::<Aes128Dec>::new_from_slices(key, iv)
            .map_err(|_| DecryptionError::InvalidKeyLength)?;
        decryptor
            .decrypt_padded::<cipher::block_padding::NoPadding>(&mut target[..encrypted_len])
            .map_err(|_| DecryptionError::DecryptionFailed)?;
    }
    Ok(data.len())
}

#[cfg(all(test, feature = "decryption"))]
mod tests {
    use super::*;

    #[test]
    fn test_decrypt_aes_cbc_basic() {
        // Test vector: simple AES-128-CBC decryption
        let key = [0u8; 16];
        let iv = [0u8; 16];
        let encrypted = [
            0x66, 0xe9, 0x4b, 0xd4, 0xef, 0x8a, 0x2c, 0x3b, 0x88, 0x4c, 0xfa, 0x59, 0xca, 0x34,
            0x2b, 0x2e,
        ];
        let mut output = [0u8; 16];

        let result = decrypt_aes_cbc_into(&encrypted, &key, &iv, &mut output);
        assert!(result.is_ok());
        let len = result.unwrap();
        assert_eq!(len, 16);
        assert_eq!(output, [0; 16]);
        let mut lazy = DecryptedBytes::cbc(&encrypted, &key, &iv).unwrap();
        assert_eq!(lazy.len(), 16);
        for remaining in (0..16).rev() {
            assert_eq!(lazy.next(), Some(0));
            assert_eq!(lazy.len(), remaining);
        }
        assert_eq!(lazy.next(), None);
        assert_eq!(lazy.next(), None);
    }

    #[test]
    fn test_key_provider_basic() {
        let mut provider = StaticKeyProvider::<10>::new();

        // Add a test key
        let key = [
            0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E,
            0x0F, 0x10,
        ];

        // Manufacturer code "ABC" -> ID calculation
        let manufacturer_code = ManufacturerCode::from_id(0x0421).unwrap(); // ABC
        let identification_number = 12345678u32;

        // Add the key
        provider
            .add_key(0x0421, identification_number, key)
            .unwrap();

        // Create context
        let context = KeyContext {
            manufacturer: manufacturer_code,
            identification_number,
            version: 0x01,
            device_type: DeviceType::WaterMeter,
            security_mode: crate::SecurityMode::AesCbc128IvZero,
            access_number: 0x00,
        };

        // Retrieve the key
        let retrieved_key = provider.get_key(&context).unwrap();
        assert_eq!(retrieved_key, &key);
    }

    #[test]
    fn test_derive_iv() {
        let manufacturer_code = ManufacturerCode::from_id(0x1ee6).unwrap(); // GWF
        let context = KeyContext {
            manufacturer: manufacturer_code,
            identification_number: 12345678,
            version: 0x42,
            device_type: DeviceType::WaterMeter,
            security_mode: crate::SecurityMode::AesCbc128IvNonZero,
            access_number: 0x50,
        };

        let payload = EncryptedPayload { data: &[], context };

        let iv = payload._derive_iv();

        // Check manufacturer ID (first 2 bytes should be 0x1ee6 in little-endian)
        assert_eq!(iv[0], 0xe6);
        assert_eq!(iv[1], 0x1e);

        // Check identification number (bytes 2-5, as BCD)
        // 12345678 decimal -> BCD [0x78, 0x56, 0x34, 0x12]
        assert_eq!(&iv[2..6], &[0x78, 0x56, 0x34, 0x12]);

        // Check version (byte 6)
        assert_eq!(iv[6], 0x42);

        // Check device type (byte 7) - WaterMeter should be 0x07
        assert_eq!(iv[7], 0x07);

        // Check access number repeated (bytes 8-15)
        for &byte in iv.iter().skip(8) {
            assert_eq!(byte, 0x50);
        }
    }

    #[test]
    fn test_mode5_decryption_real_frame() {
        // Test vector from real wMBus Mode 5 encrypted frame
        // Frame: 6644496A3100015514377203926314496A00075000500598A78E0D...
        // Key: F8B24F12F9D113F680BEE765FDE67EC0
        // Expected IV: 496A0392631400075050505050505050

        let key = [
            0xF8, 0xB2, 0x4F, 0x12, 0xF9, 0xD1, 0x13, 0xF6, 0x80, 0xBE, 0xE7, 0x65, 0xFD, 0xE6,
            0x7E, 0xC0,
        ];

        // Encrypted payload (80 bytes after TPL header)
        let encrypted = [
            0x98, 0xA7, 0x8E, 0x0D, 0x71, 0xAA, 0x63, 0x58, 0xEE, 0xBD, 0x0B, 0x20, 0xBF, 0xDF,
            0x99, 0xED, 0xA2, 0xD2, 0x2F, 0xA2, 0x53, 0x14, 0xF3, 0xF1, 0xB8, 0x44, 0x70, 0x89,
            0x8E, 0x49, 0x53, 0x03, 0x92, 0x37, 0x70, 0xBA, 0x8D, 0xDA, 0x97, 0xC9, 0x64, 0xF0,
            0xEA, 0x6C, 0xE2, 0x4F, 0x56, 0x50, 0xC0, 0xA6, 0xCD, 0xF3, 0xDE, 0x37, 0xDE, 0x33,
            0xFB, 0xFB, 0xEB, 0xAC, 0xE4, 0x00, 0x9B, 0xB0, 0xD8, 0xEB, 0xA2, 0xCB, 0xE8, 0x04,
            0x33, 0xFF, 0x13, 0x13, 0x28, 0x20, 0x60, 0x20, 0xB1, 0xBF,
        ];

        // Expected decrypted data
        let expected = [
            0x2F, 0x2F, 0x0C, 0x13, 0x29, 0x73, 0x06, 0x00, 0x02, 0x6C, 0x94, 0x21, 0x82, 0x04,
            0x6C, 0x81, 0x21, 0x8C, 0x04, 0x13, 0x75, 0x44, 0x06, 0x00, 0x8D, 0x04, 0x93, 0x13,
            0x2C, 0xFB, 0xFE, 0x12, 0x44, 0x00, 0x51, 0x41, 0x00, 0x70, 0x35, 0x00, 0x77, 0x33,
            0x00, 0x75, 0x49, 0x00, 0x16, 0x36, 0x00, 0x73, 0x56, 0x00, 0x91, 0x55, 0x00, 0x95,
            0x57, 0x00, 0x31, 0x57, 0x00, 0x28, 0x42, 0x00, 0x18, 0x47, 0x00, 0x61, 0x39, 0x00,
            0x56, 0x42, 0x00, 0x02, 0xFD, 0x17, 0x00, 0x00, 0x2F, 0x2F,
        ];

        // Build context from TPL header:
        // Frame bytes [0x49, 0x6A] = manufacturer ID 0x6A49 when parsed as little-endian
        // ID bytes [0x03, 0x92, 0x63, 0x14] in BCD = 14639203 decimal
        let manufacturer = ManufacturerCode::from_id(0x6A49).unwrap();
        let context = KeyContext {
            manufacturer,
            identification_number: 14639203, // BCD 03926314 parsed as decimal
            version: 0x00,
            device_type: DeviceType::WaterMeter,
            security_mode: crate::SecurityMode::AesCbc128IvNonZero,
            access_number: 0x50,
        };

        // Verify IV derivation
        let payload = EncryptedPayload {
            data: &encrypted,
            context: context.clone(),
        };
        let iv = payload._derive_iv();
        // Expected IV: 496A0392631400075050505050505050
        let expected_iv = [
            0x49, 0x6A, 0x03, 0x92, 0x63, 0x14, 0x00, 0x07, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50,
            0x50, 0x50,
        ];
        assert_eq!(iv, expected_iv);

        // Test decryption
        let mut provider = StaticKeyProvider::<1>::new();
        provider.add_key(0x6A49, 14639203, key).unwrap();

        let mut output = [0u8; 80];
        let len = payload.decrypt_into(&provider, &mut output).unwrap();

        assert_eq!(len, 80);
        assert_eq!(&output[..80], &expected[..]);
        assert!(payload.decrypted_bytes(&provider).unwrap().eq(expected));

        // Every partial-block boundary preserves the historical plaintext tail.
        for length in 1..=encrypted.len() {
            let payload = EncryptedPayload::new(&encrypted[..length], context.clone());
            let mut buffered = [0xAA; 80];
            assert_eq!(payload.decrypt_into(&provider, &mut buffered), Ok(length));
            let mut bytes = payload.decrypted_bytes(&provider).unwrap();
            for index in 0..length {
                assert_eq!(bytes.len(), length - index);
                let want = if index < length / 16 * 16 {
                    expected[index]
                } else {
                    encrypted[index]
                };
                assert_eq!(bytes.next(), Some(want));
                assert_eq!(buffered[index], want);
            }
            assert_eq!(bytes.next(), None);
            assert_eq!(bytes.next(), None);
        }
        // The lazy result is independent of the temporary context/provider.
        let bytes = {
            let local_payload = EncryptedPayload::new(&encrypted, context.clone());
            let mut local_provider = StaticKeyProvider::<1>::new();
            local_provider.add_key(0x6A49, 14639203, key).unwrap();
            local_payload.decrypted_bytes(&local_provider).unwrap()
        };
        assert!(bytes.eq(expected));
        let mut too_short = [0xAA; 79];
        assert_eq!(
            payload.decrypt_into(&provider, &mut too_short),
            Err(DecryptionError::InvalidDataLength)
        );
        assert_eq!(too_short, [0xAA; 79]);

        let mut context = context;
        context.security_mode = SecurityMode::NoEncryption;
        let plaintext = EncryptedPayload::new(&encrypted, context);
        assert!(plaintext.decrypted_bytes(&provider).unwrap().eq(encrypted));
        let mut output = [0xAA; 81];
        assert_eq!(plaintext.decrypt_into(&provider, &mut output), Ok(80));
        assert_eq!(&output[..80], &encrypted);
        assert_eq!(output[80], 0xAA);
        assert_eq!(
            plaintext.decrypt_into(&provider, &mut too_short),
            Err(DecryptionError::InvalidDataLength)
        );
        assert_eq!(too_short, [0xAA; 79]);
        let missing_key = StaticKeyProvider::<0>::new();
        assert_eq!(
            plaintext.decrypt_into(&missing_key, &mut too_short),
            Err(DecryptionError::KeyNotFound)
        );
        assert_eq!(too_short, [0xAA; 79]);
        let empty = EncryptedPayload::new(&[], plaintext.context);
        assert_eq!(empty.decrypt_into(&provider, &mut []), Ok(0));
    }
    #[test]
    fn lazy_decryption_rejects_invalid_parameters() {
        assert_eq!(
            DecryptedBytes::cbc(&[0; 16], &[0; 15], &[0; 16]).err(),
            Some(DecryptionError::InvalidKeyLength)
        );
        assert_eq!(
            DecryptedBytes::cbc(&[0; 16], &[0; 16], &[0; 15]).err(),
            Some(DecryptionError::InvalidDataLength)
        );
        assert_eq!(
            DecryptedBytes::cbc(&[], &[0; 16], &[0; 16]).err(),
            Some(DecryptionError::InvalidDataLength)
        );
        assert_eq!(DecryptedBytes::plaintext(&[]).next(), None);
    }
}