use super::*;

use std::{
    io::{Read, Write},
    os::{
        raw::c_char,
        unix::fs::PermissionsExt,
    },
    ffi::CStr,
    path::Path,
    fs::{File, OpenOptions},
    sync::{Arc, Mutex},
    mem::size_of,
};
use openssl::{
    bn::BigNum,
    hash::Hasher,
    nid::Nid,
};
use tss_esapi_sys::{
    TSS2_RC,
    TPMT_TK_HASHCHECK,
    TPM2_HANDLE,
    TPM2_MAX_CAP_BUFFER,
    TPM2_CAP
};
use tss_esapi::{
    Context,
    constants::{
        session_type::SessionType,
        response_code::Tss2ResponseCode,
        tss::{TPM2_RH_NULL, TPM2_PERSISTENT_FIRST},
        CapabilityType,
    },
    tcti_ldr::{TctiNameConf, TabrmdConfig},
    interface_types::{
        resource_handles::Hierarchy,
        algorithm::HashingAlgorithm,
        key_bits::RsaKeyBits,
        dynamic_handles::Persistent,
    },
    structures::{
        Digest,
        HashScheme,
        Public,
        PublicRsaParameters,
        SymmetricDefinition,
        SymmetricDefinitionObject,
        PublicKeyRsa,
        RsaScheme,
        RsaExponent,
        HashcheckTicket,
        Ticket,
        SignatureScheme,
        RsaDecryptionScheme,
        Data,
        CapabilityData,
    },
    attributes::{
        object::ObjectAttributes,
    },
    handles::{KeyHandle, PersistentTpmHandle},
};

impl From<tss_esapi::Error> for Error {
    fn from(err: tss_esapi::Error) -> Self {
        match err {
            tss_esapi::Error::WrapperError(err) => Error::custom(err),
            tss_esapi::Error::Tss2Error(err) => {
               let rc = err.tss2_rc();
               let text = tss2_rc_decode(err);
               let string = format!("response code: {}, response text: {}", rc, text);
               Error::custom(string)
            }
        }
    }
}

impl<Guard> From<std::sync::PoisonError<Guard>> for Error {
    fn from(error: std::sync::PoisonError<Guard>) -> Self {
        Error::custom(error.to_string())
    }
}

trait ContextExt {
    fn persistent_handles(&mut self) -> Result<Vec<PersistentTpmHandle>>;

    /// Returns the first free persistent handle available in the TPM.
    fn first_free_persistent(&mut self) -> Result<Persistent> {
        let mut handles = self.persistent_handles()?;
        let handle = handles.pop()
            .map_or_else(
                || PersistentTpmHandle::try_from(TPM2_PERSISTENT_FIRST),
                |v| PersistentTpmHandle::new(TPM2_HANDLE::from(v)+ 1))
            .conv_err()?;
        Ok(Persistent::Persistent(handle))
    }
}

impl ContextExt for Context {
    fn persistent_handles(&mut self) -> Result<Vec<PersistentTpmHandle>> {
        let capability= CapabilityType::Handles;
        let property = TPM2_PERSISTENT_FIRST;
        let max = usize::try_from(TPM2_MAX_CAP_BUFFER).unwrap();
        let count = (max - size_of::<TPM2_CAP>() - size_of::<u32>()) / size_of::<TPM2_HANDLE>();
        let property_count = u32::try_from(count).unwrap();
        let mut all_handles = Vec::new();
        loop {
            let (handles, more) = self.get_capability(capability, property, property_count)
                .unwrap();
            let list = match handles {
                CapabilityData::Handles(list) => list.into_inner(),
                _ => panic!("Unexpected capability type returned by TPM: {:?}", handles),
            };
            for handle in list {
                all_handles.push(PersistentTpmHandle::try_from(handle).conv_err()?);
            }
            if !more {
                break;
            }
        }
        Ok(all_handles)
    }
}

const COOKIE_LEN: usize = RSA_KEY_BYTES;

struct Cookie([u8; COOKIE_LEN]);

impl Cookie {
    fn random() -> Result<Cookie> {
        Ok(Cookie(rand_array()?))
    }

    fn empty() -> Cookie {
        Cookie([0; COOKIE_LEN])
    }

    fn load_or_init<P: AsRef<Path>>(cookie_path: P) -> Result<Cookie> {
        let cookie = match OpenOptions::new().read(true).open(&cookie_path) {
            Ok(mut file) => {
                Cookie::load(&mut file)?
            },
            Err(error) => {
                if std::io::ErrorKind::NotFound != error.kind() {
                    return Err(Error::from(error));
                }
                let cookie = Cookie::random()?;
                let mut file = OpenOptions::new().write(true).create_new(true).open(&cookie_path)
                    .conv_err()?;
                cookie.save(&mut file)?; 
                cookie
            }
        };
        Ok(cookie)
    }

    fn as_slice(&self) -> &[u8] {
        self.0.as_slice()
    }

    fn as_mut_slice(&mut self) -> &mut [u8] {
        self.0.as_mut_slice()
    }

    fn save(&self, to: &mut File) -> Result<()> {
        to.write_all(self.as_slice()).conv_err()?;
        // Only allow read access from the current user.
        let metadata = to.metadata().conv_err()?;
        let mut permissions = metadata.permissions();
        permissions.set_mode(0o400);
        to.set_permissions(permissions).conv_err()?;
        Ok(())
    }

    fn load(from: &mut File) -> Result<Cookie> {
        let mut cookie = Cookie::empty();
        from.read_exact(cookie.as_mut_slice()).conv_err()?;
        Ok(cookie)
    }
}

pub(crate) struct TpmCredStore {
    context: Arc<Mutex<Context>>,
    cookie: Cookie,
}

impl TpmCredStore {
    /// The public exponent to use for generated RSA keys.
    const RSA_EXPONENT: u32 = 65537; // 2**16 + 1

    const RSA_KEY_BITS: RsaKeyBits = RsaKeyBits::Rsa3072;

    pub(crate) fn new<P: AsRef<Path>>(
        mut context: Context, cookie_path: P
    ) -> Result<TpmCredStore> {
        let cookie = Cookie::load_or_init(cookie_path)?;
        let session = context.start_auth_session(
            None,
            None,
            None,
            SessionType::Hmac,
            SymmetricDefinition::AES_256_CFB,
            HashingAlgorithm::Sha256,
        )
        .conv_err()?
        .ok_or_else(|| Error::custom("Received invalid session handle"))?;
        context.set_sessions((Some(session), None, None));
        let context = Arc::new(Mutex::new(context));
        Ok(TpmCredStore { context, cookie })
    }

    fn gen_node_creds(&self) -> Result<TpmCreds> {
        let template = {
            let object_attributes = ObjectAttributes::builder()
                .with_fixed_tpm(true)
                .with_fixed_parent(true)
                .with_sensitive_data_origin(true)
                .with_user_with_auth(true)
                .with_decrypt(true)
                .with_sign_encrypt(true)
                .with_restricted(false)
                .build()
                .conv_err()?;
            let name_hashing_algorithm = HashingAlgorithm::Sha256;
            let empty = [0u8; 0];
            let auth_policy = Digest::try_from(empty.as_slice()).conv_err()?;
            let parameters = PublicRsaParameters::new(
                SymmetricDefinitionObject::Null,
                RsaScheme::Null,
                TpmCredStore::RSA_KEY_BITS,
                RsaExponent::try_from(TpmCredStore::RSA_EXPONENT).conv_err()?, 
            );
            let unique = PublicKeyRsa::try_from(self.cookie.as_slice())
                .conv_err()?;
            Public::Rsa {
                object_attributes,
                name_hashing_algorithm,
                auth_policy,
                parameters,
                unique,
            }
        };
        let result = {
            let mut context = self.context.lock().conv_err()?;
            context.create_primary(
                Hierarchy::Endorsement,
                template,
                None,
                None,
                None,
                None,
            )
            .conv_err()?
        };
        let public = AsymKeyPub::try_from(result.out_public)?;

        Ok(TpmCreds { public, handle: result.key_handle , context: self.context.clone() })
    }
}

impl TryFrom<Public> for AsymKeyPub {
    type Error = Error;

    fn try_from(public: Public) -> Result<AsymKeyPub> {
        match public {
            Public::Rsa { parameters, unique, .. } => {
                let exponent_value = parameters.exponent().value();
                let exponent = BigNum::from_u32(exponent_value).conv_err()?;
                let modulus = BigNum::from_slice(unique.as_slice()).conv_err()?;
                let rsa = Rsa::from_public_components(modulus, exponent).conv_err()?;
                let pkey = PKey::from_rsa(rsa).conv_err()?;
                Ok(AsymKeyPub { pkey, kind: AsymKeyKind::Rsa })
            },
            _ => Err(Error::custom("Unsupported key type returned by TPM")),
        }
    }
}

trait NidExt {
    fn sha3_256() -> Nid {
        Nid::from_raw(1097)
    }

    fn sha3_384() -> Nid {
        Nid::from_raw(1098)
    }

    fn sha3_512() -> Nid {
        Nid::from_raw(1099)
    }
}

impl NidExt for Nid {}

trait MessageDigestExt {
    fn hash_algo(&self) -> Result<HashingAlgorithm>;
    fn hash_scheme(&self) -> Result<HashScheme> {
        Ok(HashScheme::new(self.hash_algo()?))
    }
}

impl MessageDigestExt for MessageDigest {
    fn hash_algo(&self) -> Result<HashingAlgorithm> {
        let nid = self.type_();
        let algo =
            if Nid::SHA1 == nid {
                HashingAlgorithm::Sha1
            }
            else if Nid::SHA256 == nid {
                HashingAlgorithm::Sha256
            }
            else if Nid::SHA384 == nid {
                HashingAlgorithm::Sha384
            }
            else if Nid::SHA512 == nid {
                HashingAlgorithm::Sha512
            }
            else if Nid::sha3_256() == nid {
                HashingAlgorithm::Sha3_256
            }
            else if Nid::sha3_384() == nid {
                HashingAlgorithm::Sha3_384
            }
            else if Nid::sha3_512() == nid {
                HashingAlgorithm::Sha3_512
            }
            else {
                return Err(Error::custom(
                    format!("Unsupported hash algorithm with NID: {:?}", nid)));
            };
        Ok(algo)
    }
}

trait HashcheckTicketExt {
    fn null() -> HashcheckTicket;
}

impl HashcheckTicketExt for HashcheckTicket {
    /// Returns the NULL Ticket of the hashcheck type, as defined in part 1 of the TPM spec,
    /// clause 4.47.
    fn null() -> HashcheckTicket {
        let tk = TPMT_TK_HASHCHECK {
            tag: HashcheckTicket::POSSIBLE_TAGS[0].into(),
            digest: Default::default(),
            hierarchy: TPM2_RH_NULL,
        };
        HashcheckTicket::try_from(tk).unwrap()
    }
}

pub(crate) struct TpmCreds {
    public: AsymKeyPub,
    context: Arc<Mutex<Context>>,
    handle: KeyHandle,
}

impl Owned for TpmCreds {
    fn owner_of_kind(&self, kind: HashKind) -> Principal {
        self.public.owner_of_kind(kind)
    } 
}

impl Verifier for TpmCreds {
    fn verify<'a, I: Iterator<Item=&'a [u8]>>(&self, parts: I, signature: &[u8]) -> Result<bool> {
        self.public.verify(parts, signature)
    }
}

impl Encrypter for TpmCreds {
    fn encrypt(&self, slice: &[u8]) -> Result<Vec<u8>> {
        self.public.encrypt(slice)
    }
}

impl CredsPub for TpmCreds {}

impl Signer for TpmCreds {
    fn sign<'a, I: Iterator<Item=&'a [u8]>>(&self, parts: I) -> Result<Signature> {
        let msg_digest = self.public.digest();
        let digest = {
            let mut hasher = Hasher::new(msg_digest).conv_err()?;
            for part in parts {
                hasher.update(part).conv_err()?;
            }
            let bytes = hasher.finish().conv_err()?;
            let slice: &[u8] = &bytes;
            Digest::try_from(slice).conv_err()?
        };
        let validation = HashcheckTicket::null();
        let scheme = SignatureScheme::RsaSsa { hash_scheme: msg_digest.hash_scheme()? };
        let sig = {
            let mut context = self.context.lock().conv_err()?;
            context.sign(self.handle, digest, scheme, validation)
                .conv_err()?
        };
        let buf = match sig {
            tss_esapi::structures::Signature::RsaSsa(inner) => {
                let mut buf = [0u8; RSA_KEY_BYTES];
                let slice: &[u8] = inner.signature();
                buf.as_mut_slice().write_all(slice).conv_err()?;
                buf
            },
            _ => return Err(Error::custom(format!("Unexpected signature type: {:?}", sig))),
        };

        Ok(Signature::Rsa(buf))
    }
}

impl Decrypter for TpmCreds {
    fn decrypt(&self, slice: &[u8]) -> Result<Vec<u8>> {
        let cipher_text = PublicKeyRsa::try_from(slice).conv_err()?;
        let in_scheme = RsaDecryptionScheme::RsaEs;
        let empty = [0u8; 0];
        let label = Data::try_from(empty.as_slice()).conv_err()?;
        let plain_text = {
            let mut lock = self.context.lock().conv_err()?;
            lock.rsa_decrypt(self.handle, cipher_text, in_scheme, label)?
        };
        Ok(Vec::from(plain_text.value()))
    } 
}

impl CredsPriv for TpmCreds {}

impl Creds for TpmCreds {
    fn public(&self) -> &AsymKeyPub {
        &self.public
    }
}

trait TctiNameConfExt {
    fn default() -> TctiNameConf;
}

impl TctiNameConfExt for TctiNameConf {
    /// Returns a configuration which specifies that Tabrmd should be connected to using the system
    /// DBus.
    fn default() -> TctiNameConf {
        TctiNameConf::Tabrmd(TabrmdConfig::default())
    }
}

#[link(name = "tss2-rc")]
extern {
    fn Tss2_RC_Decode(rc: TSS2_RC) -> *const c_char;
}

/// Interface for types which can be converted to TSS2 response codes.
trait HasResponseCode {
    /// Returns the TSS2 response code associated with this instance.
    fn tss2_rc(&self) -> TSS2_RC;
}

/// Returns the error message associated with the given TSS2 response code.
fn tss2_rc_decode<E: HasResponseCode>(err: E) -> &'static str {
    let c_str = unsafe {
        let ptr = Tss2_RC_Decode(err.tss2_rc());
        CStr::from_ptr(ptr)
    };
    // We're relying on Tss2_RC_Decode to return valid C strings.
    c_str.to_str().unwrap()
}

impl HasResponseCode for Tss2ResponseCode {
    fn tss2_rc(&self) -> TSS2_RC {
        match self {
            Tss2ResponseCode::Success => 0,
            Tss2ResponseCode::FormatZero(code) => code.0,
            Tss2ResponseCode::FormatOne(code) => code.0,
        }
    }
}

impl HasResponseCode for TSS2_RC {
    fn tss2_rc(&self) -> TSS2_RC {
        *self
    }
}

#[cfg(test)]
mod test {
    use super::*;

    use tempdir::TempDir;
    use std::{
        fs::File,
    };
    use tss_esapi::{
        interface_types::{
            ecc::EccCurve,
        },
        structures::{
            EccPoint,
            EccScheme,
            KeyDerivationFunctionScheme,
            PublicEccParameters,
        }
    };
    use ctor::ctor;

    trait TestContextExt {
        fn for_test() -> Result<Context>;
    }

    impl TestContextExt for Context {
        fn for_test() -> Result<Context> {
            let config = TabrmdConfig::from_str("bus_type=session").conv_err()?;
            Context::new(TctiNameConf::Tabrmd(config)).conv_err()
        }
    }

    #[ctor]
    fn ctor() {
        env_logger::init();
    }

    /// Displays the message associated with a TSS2 return code.
    //#[test]
    fn print_error_message() {
        const RC: TSS2_RC = 0x00000101;
        let msg = tss2_rc_decode(RC);
        println!("{}", msg);
    }

    #[test]
    fn create_context() {
        let mut context = Context::for_test().unwrap();
        context.self_test(true).unwrap();
    }

    #[test]
    fn create_primary_key() {
        let mut context = Context::for_test().unwrap();

        let public = {
            let object_attributes = ObjectAttributes::builder()
                .with_fixed_tpm(true)
                .with_fixed_parent(true)
                .with_sensitive_data_origin(true)
                .with_user_with_auth(true)
                .with_decrypt(false)
                .with_sign_encrypt(true)
                .with_restricted(false)
                .build()
                .expect("ObjectAttributesBuilder failed");
            let name_hashing_algorithm = HashingAlgorithm::Sha256;
            let empty = [0u8; 0];
            let auth_policy = Digest::try_from(empty.as_slice()).unwrap();
            let parameters = PublicEccParameters::new(
                SymmetricDefinitionObject::Null,
                EccScheme::EcDsa(HashScheme::new(HashingAlgorithm::Sha256)),
                EccCurve::NistP256,
                KeyDerivationFunctionScheme::Null
            );
            Public::Ecc {
                object_attributes,
                name_hashing_algorithm,
                auth_policy,
                parameters,
                unique: EccPoint::default(),
            }
        };

        let session = context.start_auth_session(
            None,
            None,
            None,
            SessionType::Hmac,
            SymmetricDefinition::AES_256_CFB,
            HashingAlgorithm::Sha256,
        )
        .conv_err()
        .expect("Failed to create session")
        .expect("Received invalid handle");

        context.execute_with_session(Some(session), |ctx| {
            let primary = ctx.create_primary(
                Hierarchy::Null,
                public,
                None,
                None,
                None,
                None,
            )
            .conv_err()
            .expect("create_primary failed")
            .key_handle;

            ctx.flush_context(primary.into()).expect("flush_context failed");
        });
    }

    /// Tests that a TPM Credential Store can be created when a cookie does not already exist.
    #[test]
    fn tpm_cred_store_new() -> Result<()> {
        let dir = TempDir::new("btnode").conv_err()?;
        let cookie_path = dir.path().join("cookie.bin");
        let store = TpmCredStore::new(Context::for_test()?, &cookie_path)?;
        let cookie = File::open(&cookie_path).conv_err()?;
        let metadata = cookie.metadata().conv_err()?;
        let actual = metadata.permissions().mode();
        // Assert that the cookie can only be read by its owner.
        assert_eq!(0o400, 0o777 & actual);
        drop(store);
        dir.close()?;
        Ok(())
    }

    #[test]
    fn gen_creds() -> Result<()> {
        let dir = TempDir::new("btnode").conv_err()?;
        let cookie_path = dir.path().join("cookie.bin");
        let store = TpmCredStore::new(Context::for_test()?, &cookie_path)?;
        store.gen_node_creds()?;
        Ok(())
    }

    /// Displays the numeric identifiers used by the supported hash algorithms.
    fn show_nids() {
        fn show_nid(digest: MessageDigest) {
            let nid = digest.type_();
            println!("{}: {:?}", nid.long_name().unwrap(), nid);
        }

        show_nid(MessageDigest::sha1());
        show_nid(MessageDigest::sha256());
        show_nid(MessageDigest::sha384());
        show_nid(MessageDigest::sha512());
        show_nid(MessageDigest::sha3_256());
        show_nid(MessageDigest::sha3_384());
        show_nid(MessageDigest::sha3_512());
    }

    /// Verifies that the NIDs returned by the supported hash algorithms are as expected.
    #[test]
    fn verify_expected_nids() {
        fn assert_eq(digest: MessageDigest, nid: Nid) {
            assert_eq!(digest.type_(), nid);
        }

        assert_eq(MessageDigest::sha1(), Nid::SHA1);
        assert_eq(MessageDigest::sha256(), Nid::SHA256);
        assert_eq(MessageDigest::sha384(), Nid::SHA384);
        assert_eq(MessageDigest::sha512(), Nid::SHA512);
        assert_eq(MessageDigest::sha3_256(), Nid::sha3_256());
        assert_eq(MessageDigest::sha3_384(), Nid::sha3_384());
        assert_eq(MessageDigest::sha3_512(), Nid::sha3_512());
    }

    fn test_store() -> Result<TpmCredStore> {
        let dir = TempDir::new("btnode").conv_err()?;
        let cookie_path = dir.path().join("cookie.bin");
        let store = TpmCredStore::new(Context::for_test()?, &cookie_path)?;       
        dir.close()?;
        Ok(store)
    }

    #[test]
    fn tpm_sign_verify() -> Result<()> {
        let store = test_store()?;
        let handle = store.gen_node_creds()?;
        let data: [u8; 1024] = rand_array()?;
        let parts = [data.as_slice()];
        let sig = handle.sign(parts.into_iter())?;
        assert!(handle.verify(parts.into_iter(), sig.as_slice())?);
        Ok(())
    }

    #[test]
    fn tpm_encrypt_decrypt() -> Result<()> {
        let store = test_store()?;
        let handle = store.gen_node_creds()?;
        let expected: [u8; RSA_KEY_BYTES / 2] = rand_array()?;
        let ct = handle.encrypt(expected.as_slice())?;
        let actual = handle.decrypt(&ct)?;
        assert_eq!(expected.as_slice(), actual);
        Ok(())
    }

    /// Tests that `HashcheckTicket::null` doesn't panic.
    #[test]
    fn hashcheck_null() {
        HashcheckTicket::null();
    }

    /// Checks that the value of `TpmCredStore::RSA_KEY_BITS` matches the value of `RSA_KEY_BYTES`.
    #[test]
    fn rsa_key_bits_and_key_bytes_compatible() {
        let bytes = match TpmCredStore::RSA_KEY_BITS {
            RsaKeyBits::Rsa1024 => 128,
            RsaKeyBits::Rsa2048 => 256,
            RsaKeyBits::Rsa3072 => 384,
            RsaKeyBits::Rsa4096 => 512,
        };
        assert_eq!(RSA_KEY_BYTES, bytes)
    }

    fn list_persistent_handles() {
        let mut context = Context::for_test().unwrap();
        let all_handles = context.persistent_handles().unwrap();
        for handle in all_handles.iter() {
            println!("{:?}", handle);
        }
    }
}