//! Encryption key support

use core_foundation::{declare_TCFType, impl_TCFType};
use crate::cvt;
use core_foundation::{
    base::TCFType, string::{CFStringRef, CFString},
    dictionary::CFMutableDictionary,
};
use core_foundation::base::ToVoid;
use core_foundation::boolean::CFBoolean;
use core_foundation::data::CFData;
use core_foundation::dictionary::CFDictionary;
use core_foundation::number::CFNumber;
use core_foundation::error::{CFError, CFErrorRef};

use security_framework_sys::item::{kSecAttrIsPermanent, kSecAttrLabel, kSecAttrKeyType, kSecAttrKeySizeInBits, kSecAttrAccessControl};
#[allow(deprecated)]
use security_framework_sys::item::{kSecAttrKeyTypeEC, kSecAttrKeyTypeECSECPrimeRandom, kSecAttrKeyTypeRSA, kSecValueRef};
use security_framework_sys::keychain_item::SecItemDelete;
#[cfg(target_os = "macos")]
use security_framework_sys::item::{
    kSecAttrKeyType3DES, kSecAttrKeyTypeDSA, kSecAttrKeyTypeAES,
    kSecAttrKeyTypeDES, kSecAttrKeyTypeRC4, kSecAttrKeyTypeCAST,
};

use security_framework_sys::base::SecKeyRef;
use security_framework_sys::key::SecKeyGetTypeID;

pub use security_framework_sys::key::Algorithm;

use security_framework_sys::key::{
    SecKeyCopyAttributes, SecKeyCopyExternalRepresentation,
    SecKeyCreateSignature, SecKeyCreateRandomKey,
    SecKeyCopyPublicKey,
    SecKeyCreateDecryptedData, SecKeyCreateEncryptedData,
};
use security_framework_sys::item::kSecAttrApplicationLabel;
use std::fmt;


use crate::base::Error;
use crate::item::Location;
use crate::access_control::SecAccessControl;
/// Types of `SecKey`s.
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct KeyType(CFStringRef);

#[allow(missing_docs)]
impl KeyType {
    #[inline(always)]
    #[must_use]
    pub fn rsa() -> Self {
        unsafe { Self(kSecAttrKeyTypeRSA) }
    }

    #[cfg(target_os = "macos")]
    #[inline(always)]
    #[must_use]
    pub fn dsa() -> Self {
        unsafe { Self(kSecAttrKeyTypeDSA) }
    }

    #[cfg(target_os = "macos")]
    #[inline(always)]
    #[must_use]
    pub fn aes() -> Self {
        unsafe { Self(kSecAttrKeyTypeAES) }
    }

    #[cfg(target_os = "macos")]
    #[inline(always)]
    #[must_use]
    pub fn des() -> Self {
        unsafe { Self(kSecAttrKeyTypeDES) }
    }

    #[cfg(target_os = "macos")]
    #[inline(always)]
    #[must_use]
    pub fn triple_des() -> Self {
        unsafe { Self(kSecAttrKeyType3DES) }
    }

    #[cfg(target_os = "macos")]
    #[inline(always)]
    #[must_use]
    pub fn rc4() -> Self {
        unsafe { Self(kSecAttrKeyTypeRC4) }
    }

    #[cfg(target_os = "macos")]
    #[inline(always)]
    #[must_use]
    pub fn cast() -> Self {
        unsafe { Self(kSecAttrKeyTypeCAST) }
    }

    #[inline(always)]
    #[must_use]
    #[allow(deprecated)]
    #[deprecated(note = "deprecated by Apple")]
    pub fn ec() -> Self {
        unsafe { Self(kSecAttrKeyTypeEC) }
    }

    #[inline(always)]
    #[must_use]
    pub fn ec_sec_prime_random() -> Self {
        unsafe { Self(kSecAttrKeyTypeECSECPrimeRandom) }
    }

    pub(crate) fn as_cfstring(self) -> CFStringRef {
        self.0
    }

    /// Assumed based on the key type
    pub fn default_size_in_bits(&self) -> Option<u32> {
        #[cfg(target_os = "macos")]
        unsafe {
            if self.0 == kSecAttrKeyTypeAES {
                return Some(256);
            } else if self.0 == kSecAttrKeyTypeDES {
                return Some(64);
            } else if self.0 == kSecAttrKeyType3DES {
                return Some(192);
            } else if self.0 == kSecAttrKeyTypeRC4 {
                return Some(128);
            } else if self.0 == kSecAttrKeyTypeCAST {
                return Some(128);
            } else if self.0 == kSecAttrKeyTypeDSA {
                return Some(2048);
            }
        }
        #[allow(deprecated)]
        unsafe {
            if self.0 == kSecAttrKeyTypeRSA {
                Some(2048)
            } else if self.0 == kSecAttrKeyTypeEC {
                Some(256)
            } else if self.0 == kSecAttrKeyTypeECSECPrimeRandom {
                Some(256)
            } else {
                None
            }
        }
    }
}

declare_TCFType! {
    /// A type representing an encryption key.
    SecKey, SecKeyRef
}
impl_TCFType!(SecKey, SecKeyRef, SecKeyGetTypeID);

unsafe impl Sync for SecKey {}
unsafe impl Send for SecKey {}

impl SecKey {
    /// Translates to `SecKeyCreateRandomKey`
    #[allow(deprecated)]
    #[doc(alias = "SecKeyCreateRandomKey")]
    pub fn new(options: &GenerateKeyOptions) -> Result<Self, CFError> {
        Self::generate(options.to_dictionary())
    }

    /// Translates to `SecKeyCreateRandomKey`
    /// `GenerateKeyOptions` provides a helper to create an attribute `CFDictionary`.
    #[deprecated(note = "Use SecKey::new")]
    pub fn generate(attributes: CFDictionary) -> Result<Self, CFError> {
        let mut error: CFErrorRef = ::std::ptr::null_mut();
        let sec_key = unsafe { SecKeyCreateRandomKey(attributes.as_concrete_TypeRef(), &mut error) };
        if error.is_null() {
            Ok(unsafe { Self::wrap_under_create_rule(sec_key) })
        } else {
            Err(unsafe { CFError::wrap_under_create_rule(error) })
        }
    }

    /// Returns the programmatic identifier for the key. For keys of class
    /// kSecAttrKeyClassPublic and kSecAttrKeyClassPrivate, the value is the
    /// hash of the public key.
    #[must_use]
    pub fn application_label(&self) -> Option<Vec<u8>> {
        self.attributes()
            .find(unsafe { kSecAttrApplicationLabel.to_void() })
            .map(|v| unsafe { CFData::wrap_under_get_rule(v.cast()) }.to_vec())
    }

    /// Translates to `SecKeyCopyAttributes`
    // TODO: deprecate and remove. CFDictionary should not be exposed in public Rust APIs.
    #[must_use]
    pub fn attributes(&self) -> CFDictionary {
        let pka = unsafe { SecKeyCopyAttributes(self.to_void() as _) };
        unsafe { CFDictionary::wrap_under_create_rule(pka) }
    }

    /// Translates to `SecKeyCopyExternalRepresentation`
    // TODO: deprecate and remove. CFData should not be exposed in public Rust APIs.
    #[must_use]
    pub fn external_representation(&self) -> Option<CFData> {
        let mut error: CFErrorRef = ::std::ptr::null_mut();
        let data = unsafe { SecKeyCopyExternalRepresentation(self.to_void() as _, &mut error) };
        if data.is_null() {
            return None;
        }
        Some(unsafe { CFData::wrap_under_create_rule(data) })
    }

    /// Translates to `SecKeyCopyPublicKey`
    #[must_use]
    pub fn public_key(&self) -> Option<Self> {
        let pub_seckey = unsafe { SecKeyCopyPublicKey(self.0.cast()) };
        if pub_seckey.is_null() {
            return None;
        }

        Some(unsafe { Self::wrap_under_create_rule(pub_seckey) })
    }

    /// Encrypts a block of data using a public key and specified algorithm
    pub fn encrypt_data(&self, algorithm: Algorithm, input: &[u8]) -> Result<Vec<u8>, CFError> {
        let mut error: CFErrorRef = std::ptr::null_mut();

        let output = unsafe {
            SecKeyCreateEncryptedData(self.as_concrete_TypeRef(), algorithm.into(), CFData::from_buffer(input).as_concrete_TypeRef(), &mut error)
        };

        if error.is_null() {
            let output = unsafe { CFData::wrap_under_create_rule(output) };
            Ok(output.to_vec())
        } else {
            Err(unsafe { CFError::wrap_under_create_rule(error) })
        }
    }

    /// Decrypts a block of data using a private key and specified algorithm
    pub fn decrypt_data(&self, algorithm: Algorithm, input: &[u8]) -> Result<Vec<u8>, CFError> {
        let mut error: CFErrorRef = std::ptr::null_mut();

        let output = unsafe {
            SecKeyCreateDecryptedData(self.as_concrete_TypeRef(), algorithm.into(), CFData::from_buffer(input).as_concrete_TypeRef(), &mut error)
        };

        if error.is_null() {
            let output = unsafe { CFData::wrap_under_create_rule(output) };
            Ok(output.to_vec())
        } else {
            Err(unsafe { CFError::wrap_under_create_rule(error) })
        }
    }

    /// Creates the cryptographic signature for a block of data using a private
    /// key and specified algorithm.
    pub fn create_signature(&self, algorithm: Algorithm, input: &[u8]) -> Result<Vec<u8>, CFError> {
        let mut error: CFErrorRef = std::ptr::null_mut();

        let output = unsafe {
            SecKeyCreateSignature(
                self.as_concrete_TypeRef(),
                algorithm.into(),
                CFData::from_buffer(input).as_concrete_TypeRef(),
                &mut error,
            )
        };

        if error.is_null() {
            let output = unsafe { CFData::wrap_under_create_rule(output) };
            Ok(output.to_vec())
        } else {
            Err(unsafe { CFError::wrap_under_create_rule(error) })
        }
    }

    /// Verifies the cryptographic signature for a block of data using a public
    /// key and specified algorithm.
    pub fn verify_signature(&self, algorithm: Algorithm, signed_data: &[u8], signature: &[u8]) -> Result<bool, CFError> {
        use security_framework_sys::key::SecKeyVerifySignature;
        let mut error: CFErrorRef = std::ptr::null_mut();

        let valid = unsafe {
            SecKeyVerifySignature(
                self.as_concrete_TypeRef(),
                algorithm.into(),
                CFData::from_buffer(signed_data).as_concrete_TypeRef(),
                CFData::from_buffer(signature).as_concrete_TypeRef(),
                &mut error,
            )
        };

        if !error.is_null() {
            return Err(unsafe { CFError::wrap_under_create_rule(error) })?;
        }
        Ok(valid != 0)
    }

    /// Performs the Diffie-Hellman style of key exchange.
    pub fn key_exchange(
        &self,
        algorithm: Algorithm,
        public_key: &Self,
        requested_size: usize,
        shared_info: Option<&[u8]>,
    ) -> Result<Vec<u8>, CFError> {
        use core_foundation::data::CFData;
        use security_framework_sys::item::{
            kSecKeyKeyExchangeParameterRequestedSize, kSecKeyKeyExchangeParameterSharedInfo,
        };

        unsafe {
            let mut params = vec![(
                CFString::wrap_under_get_rule(kSecKeyKeyExchangeParameterRequestedSize),
                CFNumber::from(requested_size as i64).into_CFType(),
            )];

            if let Some(shared_info) = shared_info {
                params.push((
                    CFString::wrap_under_get_rule(kSecKeyKeyExchangeParameterSharedInfo),
                    CFData::from_buffer(shared_info).as_CFType(),
                ));
            }

            let parameters = CFDictionary::from_CFType_pairs(&params);

            let mut error: CFErrorRef = std::ptr::null_mut();

            let output = security_framework_sys::key::SecKeyCopyKeyExchangeResult(
                self.as_concrete_TypeRef(),
                algorithm.into(),
                public_key.as_concrete_TypeRef(),
                parameters.as_concrete_TypeRef(),
                &mut error,
            );

            if error.is_null() {
                let output = CFData::wrap_under_create_rule(output);
                Ok(output.to_vec())
            } else {
                Err(CFError::wrap_under_create_rule(error))
            }
        }
    }

    /// Translates to `SecItemDelete`, passing in the `SecKeyRef`
    pub fn delete(&self) -> Result<(), Error> {
        let query = CFMutableDictionary::from_CFType_pairs(&[(
            unsafe { kSecValueRef }.to_void(),
            self.to_void(),
        )]);

        cvt(unsafe { SecItemDelete(query.as_concrete_TypeRef()) })
    }
}

/// Where to generate the key.
#[derive(Debug)]
pub enum Token {
    /// Generate the key in software, compatible with all `KeyType`s.
    Software,
    /// Generate the key in the Secure Enclave such that the private key is not
    /// extractable. Only compatible with `KeyType::ec()`.
    SecureEnclave,
}

/// Helper for creating `CFDictionary` attributes for `SecKey::generate`
/// Recommended reading:
/// <https://developer.apple.com/documentation/technotes/tn3137-on-mac-keychains>
#[derive(Debug, Default)]
pub struct GenerateKeyOptions {
    /// kSecAttrKeyType
    #[deprecated(note = "use set_key_type()")]
    pub key_type: Option<KeyType>,
    /// kSecAttrKeySizeInBits
    #[deprecated(note = "use set_size_in_bits()")]
    pub size_in_bits: Option<u32>,
    /// kSecAttrLabel
    #[deprecated(note = "use set_label()")]
    pub label: Option<String>,
    /// kSecAttrTokenID
    #[deprecated(note = "use set_token()")]
    pub token: Option<Token>,
    /// Which keychain to store the key in, if any.
    #[deprecated(note = "use set_location()")]
    pub location: Option<Location>,
    /// Access control
    #[deprecated(note = "use set_access_control()")]
    pub access_control: Option<SecAccessControl>,
    /// `kSecAttrSynchronizable`
    #[cfg(feature = "sync-keychain")]
    synchronizable: Option<bool>,
}

#[allow(deprecated)]
impl GenerateKeyOptions {
    /// `kSecAttrKeyType`
    pub fn set_key_type(&mut self, key_type: KeyType) -> &mut Self {
        self.key_type = Some(key_type);
        self
    }

    /// `kSecAttrKeySizeInBits`
    pub fn set_size_in_bits(&mut self, size_in_bits: u32) -> &mut Self {
        self.size_in_bits = Some(size_in_bits);
        self
    }

    /// `kSecAttrLabel`
    pub fn set_label(&mut self, label: impl Into<String>) -> &mut Self {
        self.label = Some(label.into());
        self
    }

    /// `kSecAttrTokenID`
    pub fn set_token(&mut self, token: Token) -> &mut Self {
        self.token = Some(token);
        self
    }

    /// Which keychain to store the key in, if any.
    pub fn set_location(&mut self, location: Location) -> &mut Self {
        self.location = Some(location);
        self
    }

    /// Set `access_control`
    pub fn set_access_control(&mut self, access_control: SecAccessControl) -> &mut Self {
        self.access_control = Some(access_control);
        self
    }

    /// Set `synchronizable` (`kSecAttrSynchronizable`)
    #[cfg(feature = "sync-keychain")]
    pub fn set_synchronizable(&mut self, synchronizable: bool) -> &mut Self {
        self.synchronizable = Some(synchronizable);
        self
    }

    /// Collect options into a `CFDictioanry`
    // CFDictionary should not be exposed in public Rust APIs.
    #[deprecated(note = "Pass the options to SecKey::new")]
    pub fn to_dictionary(&self) -> CFDictionary {
        use security_framework_sys::item::{kSecAttrTokenID, kSecAttrTokenIDSecureEnclave};

        let is_permanent = CFBoolean::from(self.location.is_some());
        let mut private_attributes = CFMutableDictionary::from_CFType_pairs(&[(
            unsafe { kSecAttrIsPermanent }.to_void(),
            is_permanent.to_void(),
        )]);
        if let Some(access_control) = &self.access_control {
            private_attributes.set(unsafe { kSecAttrAccessControl }.to_void(), access_control.to_void());
        }

        #[cfg(target_os = "macos")]
        let public_attributes = CFMutableDictionary::from_CFType_pairs(&[(
            unsafe { kSecAttrIsPermanent }.to_void(),
            is_permanent.to_void(),
        )]);

        let key_type = self.key_type.unwrap_or_else(KeyType::rsa);
        let key_type_str = key_type.as_cfstring();

        let size_in_bits = self.size_in_bits.unwrap_or(key_type.default_size_in_bits().unwrap_or(256));
        let size_in_bits = CFNumber::from(size_in_bits as i32);

        let mut attribute_key_values = vec![
            (unsafe { kSecAttrKeyType }.to_void(), key_type_str.to_void()),
            (unsafe { kSecAttrKeySizeInBits }.to_void(), size_in_bits.to_void()),
        ];
        #[cfg(target_os = "macos")]
        if key_type != KeyType::aes() {
            attribute_key_values.push((unsafe { security_framework_sys::item::kSecPublicKeyAttrs }.to_void(), public_attributes.to_void()));
            attribute_key_values.push((unsafe { security_framework_sys::item::kSecPrivateKeyAttrs }.to_void(), private_attributes.to_void()));
        } else {
            attribute_key_values.push((unsafe { security_framework_sys::item::kSecAttrCanEncrypt }.to_void(), CFBoolean::true_value().to_void()));
            attribute_key_values.push((unsafe { security_framework_sys::item::kSecAttrCanDecrypt }.to_void(), CFBoolean::true_value().to_void()));
        }

        let label = self.label.as_deref().map(CFString::new);
        if let Some(label) = &label {
            attribute_key_values.push((unsafe { kSecAttrLabel }.to_void(), label.to_void()));
        }

        #[cfg(target_os = "macos")]
        match &self.location {
            #[cfg(feature = "OSX_10_15")]
            Some(Location::DataProtectionKeychain) => {
                use security_framework_sys::item::kSecUseDataProtectionKeychain;
                attribute_key_values.push((
                    unsafe { kSecUseDataProtectionKeychain }.to_void(),
                    CFBoolean::true_value().to_void(),
                ));
            }
            Some(Location::FileKeychain(keychain)) => {
                attribute_key_values.push((
                    unsafe { security_framework_sys::item::kSecUseKeychain }.to_void(),
                    keychain.as_concrete_TypeRef().to_void(),
                ));
            },
            _ => {},
        }

        match self.token.as_ref().unwrap_or(&Token::Software) {
            Token::Software => {},
            Token::SecureEnclave => {
                attribute_key_values.push((
                    unsafe { kSecAttrTokenID }.to_void(),
                    unsafe { kSecAttrTokenIDSecureEnclave }.to_void(),
                ));
            }
        }

        #[cfg(feature = "sync-keychain")]
        if let Some(synchronizable) = &self.synchronizable {
            attribute_key_values.push((
                unsafe { security_framework_sys::item::kSecAttrSynchronizable }.to_void(),
                CFBoolean::from(*synchronizable).to_void(),
            ));
        }

        CFMutableDictionary::from_CFType_pairs(&attribute_key_values).to_immutable()
    }
}

impl fmt::Debug for SecKey {
    #[cold]
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("SecKey").finish_non_exhaustive()
    }
}

#[cfg(test)]
mod test {
    use super::{KeyType, SecKey, GenerateKeyOptions};

    #[test]
    #[cfg(target_os = "macos")]
    fn generate_aes_key() {
        let key = SecKey::new(&GenerateKeyOptions::default()
            .set_key_type(KeyType::aes())
            .set_size_in_bits(256)).expect("Failed to generate AES key");

        // Verify the key has attributes
        let attrs = key.attributes();
        assert!(!attrs.is_empty(), "AES key should have attributes");
    }
}