/*
 * Copyright 2025 Yubico AB
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use strum_macros::EnumIter;
use yubihsmrs::object::{ObjectAlgorithm, ObjectCapability, ObjectDescriptor, ObjectType};
use yubihsmrs::Session;
use crate::traits::operation_traits::YubihsmOperations;
use crate::traits::command_traits::Command;
use crate::common::error::MgmError;
use crate::common::util::{get_op_keys, get_object_descriptors};
use crate::common::algorithms::MgmAlgorithm;
use crate::common::types::{NewObjectSpec, EXIT_LABEL};


#[derive(Debug, Clone, Copy, PartialEq, Eq, EnumIter)]
pub enum SymCommand {
    List,
    GetKeyProperties,
    Generate,
    Import,
    Delete,
    Encrypt,
    Decrypt,
    GetRandom,
    Exit,
}

impl Command for SymCommand {
    fn label(&self) -> &'static str {
        match self {
            Self::List => "List",
            Self::GetKeyProperties => "Get Object Properties",
            Self::Generate => "Generate",
            Self::Import => "Import",
            Self::Delete => "Delete",
            Self::Encrypt => "Encrypt",
            Self::Decrypt => "Decrypt",
            Self::GetRandom => "Generate pseudo random number",
            Self::Exit => EXIT_LABEL,
        }
    }

    fn description(&self) -> &'static str {
        match self {
            Self::List => "List all symmetric keys stored on the YubiHSM",
            Self::GetKeyProperties => "Get properties of a symmetric key stored on the YubiHSM",
            Self::Generate => "Generate a new symmetric key inside the YubiHSM",
            Self::Import => "Import a symmetric key into the YubiHSM",
            Self::Delete => "Delete a symmetric key from the YubiHSM",
            Self::Encrypt => "Encrypt data using an AES key stored on the YubiHSM",
            Self::Decrypt => "Decrypt data using an AES key stored on the YubiHSM",
            Self::GetRandom => "Get pseudo random bytes generated by the YubiHSM",
            Self::Exit => "",
        }
    }

    fn required_capabilities(&self) -> &'static [ObjectCapability] {
        match self {
            Self::List | Self::GetKeyProperties | Self::Exit => &[],
            Self::Generate => &[ObjectCapability::GenerateSymmetricKey],
            Self::Import => &[ObjectCapability::PutSymmetricKey],
            Self::Delete => &[ObjectCapability::DeleteSymmetricKey],
            Self::Encrypt => &[ObjectCapability::EncryptEcb, ObjectCapability::EncryptCbc],
            Self::Decrypt => &[ObjectCapability::DecryptEcb, ObjectCapability::DecryptCbc],
            Self::GetRandom => &[ObjectCapability::GetPseudoRandom],
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq,  Eq, Default)]
pub enum AesMode {
    #[default]
    Ecb,
    Cbc,
}

#[derive(Debug, Clone, Copy, PartialEq,  Eq, Default)]
pub enum EncryptionMode {
    #[default]
    Encrypt,
    Decrypt,
}

pub struct AesOperationSpec {
    pub operation_key: ObjectDescriptor,
    pub aes_mode: AesMode,
    pub enc_mode: EncryptionMode,
    pub iv: Vec<u8>,
    pub data: Vec<u8>,
}

pub struct SymmetricOperations;

impl YubihsmOperations for SymmetricOperations {

    fn context(&self) -> &'static str {
        SymmetricOperations::SYM_CONTEXT
    }

    fn get_all_objects(&self, session: &Session) -> Result<Vec<ObjectDescriptor>, MgmError> {
        let keys = session.list_objects_with_filter(
            0,
            ObjectType::SymmetricKey,
            "",
            ObjectAlgorithm::ANY,
            &Vec::new())?;
        get_object_descriptors(session, &keys)
    }

    fn get_generation_algorithms(&self) -> Vec<MgmAlgorithm> {
        MgmAlgorithm::AES_KEY_ALGORITHMS.to_vec()
    }

    fn get_object_capabilities(
        &self,
        _object_type: Option<ObjectType>,
        _object_algorithm: Option<ObjectAlgorithm>) -> Result<Vec<ObjectCapability>, MgmError> {
        Ok(Self::AES_KEY_CAPABILITIES.to_vec())
    }

    fn generate(&self, session: &Session, spec: &NewObjectSpec) -> Result<u16, MgmError> {
        Ok(session
            .generate_aes_key(
                spec.id,
                &spec.label,
                &spec.capabilities,
                &spec.domains,
                spec.algorithm)?)
    }

    fn import(&self, session: &Session, spec: &NewObjectSpec) -> Result<u16, MgmError> {
        match Self::get_symkey_algorithm_from_keylen(spec.data[0].len()) {
                Err(e) => return Err(MgmError::Error(format!("Failed to import symmetric key: {}", e))),
                Ok(algo) => {
                    if algo != spec.algorithm {
                        return Err(MgmError::Error(format!("Algorithm mismatch: key data length implies {:?} but spec algorithm is {:?}", algo, spec.algorithm)));
                    }
                }
        }
        Ok(session
            .import_aes_key(
                spec.id,
                &spec.label,
                &spec.domains,
                &spec.capabilities,
                spec.algorithm,
                &spec.data[0])?)
    }
}

impl SymmetricOperations {

    pub const SYM_CONTEXT: &'static str = "sym";

    const AES_KEY_CAPABILITIES: [ObjectCapability; 5] = [
        ObjectCapability::EncryptCbc,
        ObjectCapability::DecryptCbc,
        ObjectCapability::EncryptEcb,
        ObjectCapability::DecryptEcb,
        ObjectCapability::ExportableUnderWrap];

    pub fn is_aes_algorithm(algorithm: &ObjectAlgorithm) -> bool {
        MgmAlgorithm::AES_KEY_ALGORITHMS.iter().any(|a| a.algorithm() == *algorithm)
    }

    pub fn get_symkey_algorithm_from_keylen(keylen: usize) -> Result<ObjectAlgorithm, MgmError> {
        match keylen {
            16 => Ok(ObjectAlgorithm::Aes128),
            24 => Ok(ObjectAlgorithm::Aes192),
            32 => Ok(ObjectAlgorithm::Aes256),
            _ => Err(MgmError::Error(format!("Unsupported key length {}", keylen))),
        }
    }

    pub fn get_operation_keys(session:&Session, authkey: &ObjectDescriptor, enc_mode: EncryptionMode, aes_mode: AesMode) -> Result<Vec<ObjectDescriptor>, MgmError> {
        let caps = match enc_mode {
            EncryptionMode::Encrypt => match aes_mode {
                AesMode::Ecb => [ObjectCapability::EncryptEcb],
                AesMode::Cbc => [ObjectCapability::EncryptCbc],
            },
            EncryptionMode::Decrypt => match aes_mode {
                AesMode::Ecb => [ObjectCapability::DecryptEcb],
                AesMode::Cbc => [ObjectCapability::DecryptCbc],
            },
        };
        let keys = get_op_keys(
            session,
            authkey,
            &caps,
            ObjectType::SymmetricKey,
            Some(&MgmAlgorithm::extract_algorithms(&MgmAlgorithm::AES_KEY_ALGORITHMS))
        )?;
        Ok(keys)
    }

    pub fn operate(session:&Session, op_spec: AesOperationSpec) -> Result<Vec<u8>, MgmError> {
        let out_data = match op_spec.enc_mode {
            EncryptionMode::Encrypt => {
                match op_spec.aes_mode {
                    AesMode::Ecb => session.encrypt_aes_ecb(op_spec.operation_key.object_id(), &op_spec.data)?,
                    AesMode::Cbc => session.encrypt_aes_cbc(op_spec.operation_key.object_id(), &op_spec.iv, &op_spec.data)?
                }
            },
            EncryptionMode::Decrypt => {
                match op_spec.aes_mode {
                    AesMode::Ecb => session.decrypt_aes_ecb(op_spec.operation_key.object_id(), &op_spec.data)?,
                    AesMode::Cbc => session.decrypt_aes_cbc(op_spec.operation_key.object_id(), &op_spec.iv, &op_spec.data)?
                }
            }
        };
        Ok(out_data)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::traits::operation_traits::YubihsmOperations;

    // ══════════════════════════════════════════════
    //  get_object_capabilities
    // ══════════════════════════════════════════════

    #[test]
    fn test_sym_object_capabilities() {
        let caps = SymmetricOperations.get_object_capabilities(None, None).unwrap();
        assert_eq!(caps.len(), SymmetricOperations::AES_KEY_CAPABILITIES.len());
        assert!(caps.iter().all(|c| SymmetricOperations::AES_KEY_CAPABILITIES.contains(c)));

        // neither the object type nor the algorithm have any relevance
        let caps = SymmetricOperations.get_object_capabilities(Some(ObjectType::WrapKey), Some(ObjectAlgorithm::Rsa2048)).unwrap();
        assert_eq!(caps.len(), SymmetricOperations::AES_KEY_CAPABILITIES.len());
        assert!(caps.iter().all(|c| SymmetricOperations::AES_KEY_CAPABILITIES.contains(c)));
    }

    // ══════════════════════════════════════════════
    //  is_aes_algorithm
    // ══════════════════════════════════════════════

    #[test]
    fn test_is_aes_algorithm() {
        for algo in MgmAlgorithm::AES_KEY_ALGORITHMS.iter() {
            assert!(SymmetricOperations::is_aes_algorithm(&algo.algorithm()));
        }

        assert!(!SymmetricOperations::is_aes_algorithm(&ObjectAlgorithm::Rsa2048));
        assert!(!SymmetricOperations::is_aes_algorithm(&ObjectAlgorithm::Aes128CcmWrap));
    }

    // ══════════════════════════════════════════════
    //  get_symkey_algorithm_from_keylen
    // ══════════════════════════════════════════════

    #[test]
    fn test_algorithm_from_keylen() {
        assert_eq!(
            SymmetricOperations::get_symkey_algorithm_from_keylen(16).unwrap(),
            ObjectAlgorithm::Aes128
        );
        assert_eq!(
            SymmetricOperations::get_symkey_algorithm_from_keylen(24).unwrap(),
            ObjectAlgorithm::Aes192
        );
        assert_eq!(
            SymmetricOperations::get_symkey_algorithm_from_keylen(32).unwrap(),
            ObjectAlgorithm::Aes256
        );

        assert!(SymmetricOperations::get_symkey_algorithm_from_keylen(0).is_err());
        assert!(SymmetricOperations::get_symkey_algorithm_from_keylen(15).is_err());
        assert!(SymmetricOperations::get_symkey_algorithm_from_keylen(64).is_err());
    }

    // ══════════════════════════════════════════════
    //  Enum defaults
    // ══════════════════════════════════════════════

    #[test]
    fn test_aes_mode_default() {
        assert_eq!(AesMode::default(), AesMode::Ecb);
    }

    #[test]
    fn test_encryption_mode_default() {
        assert_eq!(EncryptionMode::default(), EncryptionMode::Encrypt);
    }
}