//
// Copyright (C) 2024 The Android Open Source Project
//
// 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 crate::error::Error;
use alloc::boxed::Box;
use alloc::vec::Vec;
use kmr_common::{crypto, km_err, vec_try};
use kmr_wire::keymint::Digest;
use optee_utee::{
    AlgorithmId, AttributeId, AttributeMemref, Mac, TransientObject, TransientObjectType,
};

struct MacOperation {
    mac: Mac,
    len: usize,
}

// SAFETY: The raw pointer to TEE_OperationHandle is held only by us, so
// it can be safely transferred to another thread.
unsafe impl Send for MacOperation {}

impl crypto::AccumulatingOperation for MacOperation {
    fn update(&mut self, data: &[u8]) -> Result<(), kmr_common::Error> {
        self.mac.update(data);
        Ok(())
    }

    fn finish(self: Box<Self>) -> Result<Vec<u8>, kmr_common::Error> {
        let mut output = vec_try![0u8; self.len]?;
        self.mac.compute_final(&[0u8; 0], output.as_mut_slice()).map_err(Error::kmerr)?;
        Ok(output)
    }
}

pub struct Hmac;

impl crypto::Hmac for Hmac {
    fn begin(
        &self,
        key: crypto::OpaqueOr<crypto::hmac::Key>,
        digest: Digest,
    ) -> Result<Box<dyn crypto::AccumulatingOperation>, kmr_common::Error> {
        let (algorithm, keytype, output_len) = match &digest {
            Digest::Md5 => Ok((AlgorithmId::HmacMd5, TransientObjectType::HmacMd5, 16)),
            Digest::Sha1 => Ok((AlgorithmId::HmacSha1, TransientObjectType::HmacSha1, 20)),
            Digest::Sha224 => Ok((AlgorithmId::HmacSha224, TransientObjectType::HmacSha224, 28)),
            Digest::Sha256 => Ok((AlgorithmId::HmacSha256, TransientObjectType::HmacSha256, 32)),
            Digest::Sha384 => Ok((AlgorithmId::HmacSha384, TransientObjectType::HmacSha384, 48)),
            Digest::Sha512 => Ok((AlgorithmId::HmacSha512, TransientObjectType::HmacSha512, 64)),
            digest => Err(km_err!(UnsupportedDigest, "Unknown digest {:?}", digest)),
        }?;
        let key = kmr_common::explicit!(key)?;
        let len = key.0.len() * 8;

        let mut keyobj = TransientObject::allocate(keytype, len).map_err(Error::kmerr)?;
        let attr = AttributeMemref::from_ref(AttributeId::SecretValue, key.0.as_ref());
        keyobj.populate(&[attr.into()]).map_err(Error::kmerr)?;

        let mac = Mac::allocate(algorithm, output_len * 8).map_err(Error::kmerr)?;
        mac.set_key(&keyobj).map_err(Error::kmerr)?;
        mac.init(&[0u8; 0]);

        Ok(Box::new(MacOperation { mac: mac, len: output_len }))
    }
}

pub struct AesCmac;

impl crypto::AesCmac for AesCmac {
    fn begin(
        &self,
        key: crypto::OpaqueOr<crypto::aes::Key>,
    ) -> Result<Box<dyn crypto::AccumulatingOperation>, kmr_common::Error> {
        let key = kmr_common::explicit!(key)?;
        let (keybuf, keysize_in_bits) = match &key {
            crypto::aes::Key::Aes128(k) => (&k[..], 128),
            crypto::aes::Key::Aes192(k) => (&k[..], 192),
            crypto::aes::Key::Aes256(k) => (&k[..], 256),
        };

        let mut keyobj = TransientObject::allocate(TransientObjectType::Aes, key.size().0 as usize)
            .map_err(Error::kmerr)?;
        let attr = AttributeMemref::from_ref(AttributeId::SecretValue, keybuf);
        keyobj.populate(&[attr.into()]).map_err(Error::kmerr)?;

        let mac = Mac::allocate(AlgorithmId::AesCmac, keysize_in_bits).map_err(Error::kmerr)?;
        mac.set_key(&keyobj).map_err(Error::kmerr)?;
        mac.init(&[0u8; 0]);

        Ok(Box::new(MacOperation { mac: mac, len: crypto::aes::BLOCK_SIZE }))
    }
}