// Copyright 2020 Patrick Uiterwijk
//
// Licensed under the EUPL-1.2-or-later
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://joinup.ec.europa.eu/collection/eupl/eupl-text-eupl-12
//
// 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 crate::error::{Error, Result};

use std::convert::{TryFrom, TryInto};

use serde::{Deserialize, Serialize};
use tss_esapi::{interface_types::algorithm::HashingAlgorithm, structures::SignatureScheme};

fn serialize_as_base64<S>(bytes: &[u8], serializer: S) -> std::result::Result<S::Ok, S::Error>
where
    S: serde::Serializer,
{
    serializer.serialize_str(&base64::encode(bytes))
}

fn deserialize_as_base64<'de, D>(deserializer: D) -> std::result::Result<Vec<u8>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    String::deserialize(deserializer)
        .and_then(|string| base64::decode(&string).map_err(serde::de::Error::custom))
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub enum SignedPolicyStep {
    PCRs {
        pcr_ids: Vec<u16>,
        hash_algorithm: String,
        #[serde(
            deserialize_with = "deserialize_as_base64",
            serialize_with = "serialize_as_base64"
        )]
        value: Vec<u8>,
    },
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct SignedPolicy {
    // policy_ref contains the policy_ref used in the aHash, used to determine the policy to use from a list
    #[serde(
        deserialize_with = "deserialize_as_base64",
        serialize_with = "serialize_as_base64"
    )]
    pub policy_ref: Vec<u8>,
    // steps contains the policy steps that are signed
    pub steps: Vec<SignedPolicyStep>,
    // signature contains the signature over aHash
    #[serde(
        deserialize_with = "deserialize_as_base64",
        serialize_with = "serialize_as_base64"
    )]
    pub signature: Vec<u8>,
}

pub type SignedPolicyList = Vec<SignedPolicy>;

#[derive(Debug, Clone)]
pub enum TPMPolicyStep {
    NoStep,
    PCRs(HashingAlgorithm, Vec<u64>, Box<TPMPolicyStep>),
    Authorized {
        signkey: PublicKey,
        policy_ref: Vec<u8>,
        policies: Option<SignedPolicyList>,
        next: Box<TPMPolicyStep>,
    },
    Or([Box<TPMPolicyStep>; 8]),
}

#[derive(Debug, Serialize, Deserialize, Clone, Copy)]
pub enum RSAPublicKeyScheme {
    RSAPSS,
    RSASSA,
}

impl RSAPublicKeyScheme {
    fn to_scheme(&self, hash_algo: HashAlgo) -> SignatureScheme {
        match self {
            RSAPublicKeyScheme::RSAPSS => SignatureScheme::RsaPss {
                hash_scheme: hash_algo.into(),
            },
            RSAPublicKeyScheme::RSASSA => SignatureScheme::RsaSsa {
                hash_scheme: hash_algo.into(),
            },
        }
    }

    fn to_rsa_scheme(&self, hash_algo: HashAlgo) -> Result<tss_esapi::structures::RsaScheme> {
        match self {
            RSAPublicKeyScheme::RSAPSS => {
                Ok(tss_esapi::structures::RsaScheme::RsaPss(hash_algo.into()))
            }
            RSAPublicKeyScheme::RSASSA => {
                Ok(tss_esapi::structures::RsaScheme::RsaSsa(hash_algo.into()))
            }
        }
    }
}

#[derive(Debug, Serialize, Deserialize, Clone, Copy)]
pub enum HashAlgo {
    SHA1,
    SHA256,
    SHA384,
    SHA512,
    SM3_256,
    SHA3_256,
    SHA3_384,
    SHA3_512,
}

impl From<HashAlgo> for HashingAlgorithm {
    fn from(halg: HashAlgo) -> Self {
        match halg {
            HashAlgo::SHA1 => HashingAlgorithm::Sha1,
            HashAlgo::SHA256 => HashingAlgorithm::Sha256,
            HashAlgo::SHA384 => HashingAlgorithm::Sha384,
            HashAlgo::SHA512 => HashingAlgorithm::Sha512,
            HashAlgo::SM3_256 => HashingAlgorithm::Sm3_256,
            HashAlgo::SHA3_256 => HashingAlgorithm::Sha3_256,
            HashAlgo::SHA3_384 => HashingAlgorithm::Sha3_384,
            HashAlgo::SHA3_512 => HashingAlgorithm::Sha3_512,
        }
    }
}

impl From<HashAlgo> for tss_esapi::structures::HashScheme {
    fn from(halg: HashAlgo) -> Self {
        tss_esapi::structures::HashScheme::new(halg.into())
    }
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub enum PublicKey {
    RSA {
        scheme: RSAPublicKeyScheme,
        hashing_algo: HashAlgo,
        exponent: u32,
        #[serde(
            deserialize_with = "deserialize_as_base64",
            serialize_with = "serialize_as_base64"
        )]
        modulus: Vec<u8>,
    },
}

impl PublicKey {
    pub(crate) fn get_signing_scheme(&self) -> SignatureScheme {
        match self {
            PublicKey::RSA {
                scheme,
                hashing_algo,
                exponent: _,
                modulus: _,
            } => scheme.to_scheme(*hashing_algo),
        }
    }
}

impl TryFrom<PublicKey> for tss_esapi::structures::Public {
    type Error = Error;

    fn try_from(publickey: PublicKey) -> Result<Self> {
        match publickey {
            PublicKey::RSA {
                scheme,
                hashing_algo,
                modulus,
                exponent,
            } => {
                let object_attributes =
                    tss_esapi::attributes::object::ObjectAttributesBuilder::new()
                        .with_fixed_tpm(false)
                        .with_fixed_parent(false)
                        .with_sensitive_data_origin(false)
                        .with_user_with_auth(true)
                        .with_decrypt(false)
                        .with_sign_encrypt(true)
                        .with_restricted(false);

                Ok(tss_esapi::structures::PublicBuilder::new()
                    .with_public_algorithm(tss_esapi::interface_types::algorithm::PublicAlgorithm::Rsa)
                    .with_name_hashing_algorithm(hashing_algo.into())
                    .with_rsa_parameters(tss_esapi::structures::PublicRsaParametersBuilder::new_unrestricted_signing_key(
                            scheme.to_rsa_scheme(hashing_algo)?,
                            ((modulus.len() * 8) as u16).try_into()?,
                            tss_esapi::structures::RsaExponent::create(exponent)?,
                        )
                        .build()?
                    )
                    .with_object_attributes(object_attributes.build()?)
                    .with_rsa_unique_identifier(modulus.clone().try_into()?)
                    .build()?)
            }
        }
    }
}

fn get_pcr_hash_alg_from_name(name: Option<&String>) -> HashingAlgorithm {
    match name {
        None => HashingAlgorithm::Sha256,
        Some(val) => match val.to_lowercase().as_str() {
            "sha1" => HashingAlgorithm::Sha1,
            "sha256" => HashingAlgorithm::Sha256,
            "sha384" => HashingAlgorithm::Sha384,
            "sha512" => HashingAlgorithm::Sha512,
            _ => panic!("Unsupported hash algo: {:?}", name),
        },
    }
}

impl TryFrom<&SignedPolicyStep> for TPMPolicyStep {
    type Error = Error;

    fn try_from(spolicy: &SignedPolicyStep) -> Result<Self> {
        match spolicy {
            SignedPolicyStep::PCRs {
                pcr_ids,
                hash_algorithm,
                value: _,
            } => Ok(TPMPolicyStep::PCRs(
                get_pcr_hash_alg_from_name(Some(&hash_algorithm)),
                pcr_ids.iter().map(|x| *x as u64).collect(),
                Box::new(TPMPolicyStep::NoStep),
            )),
        }
    }
}
