wolfcose 0.1.0

use crate::{
    decrypt0_into, decrypt_into, encrypt0_detached_into, encrypt0_into, encrypt0_to_vec,
    encrypt_detached_into, encrypt_into, encrypt_to_vec, mac0_into, mac0_to_vec, mac_into,
    mac_to_vec, sign1_into, sign1_to_vec, sign_into, sign_to_vec, verify1, verify_mac, verify_mac0,
    verify_sign, Algorithm, CoseKey, DecryptOutput, Encrypt0DetachedOutput, Error, MacVerifyOutput,
    PayloadMode, Recipient, Result, Signature, VerifyOutput,
};
use alloc::{vec, vec::Vec};
use core::ptr::NonNull;

const DEFAULT_SCRATCH: usize = 4096;

/// Builder for COSE_Sign1 operations.
pub struct Sign1Builder<'a> {
    key: Option<&'a CoseKey>,
    algorithm: Option<Algorithm>,
    kid: Option<&'a [u8]>,
    payload: Option<PayloadMode<'a>>,
    external_aad: &'a [u8],
    scratch: Vec<u8>,
    rng: Option<NonNull<crate::raw::WC_RNG>>,
}

impl<'a> Sign1Builder<'a> {
    /// Create an empty builder.
    pub fn new() -> Self {
        Self {
            key: None,
            algorithm: None,
            kid: None,
            payload: None,
            external_aad: &[],
            scratch: vec![0; DEFAULT_SCRATCH],
            rng: None,
        }
    }

    /// Set signing key.
    pub fn key(mut self, key: &'a CoseKey) -> Self {
        self.key = Some(key);
        self
    }

    /// Set signature algorithm.
    pub fn algorithm(mut self, algorithm: Algorithm) -> Self {
        self.algorithm = Some(algorithm);
        self
    }

    /// Set key identifier.
    pub fn kid(mut self, kid: &'a [u8]) -> Self {
        self.kid = Some(kid);
        self
    }

    /// Set payload mode.
    pub fn payload(mut self, payload: PayloadMode<'a>) -> Self {
        self.payload = Some(payload);
        self
    }

    /// Set external AAD.
    pub fn external_aad(mut self, external_aad: &'a [u8]) -> Self {
        self.external_aad = external_aad;
        self
    }

    /// Replace the internal scratch buffer size.
    pub fn scratch_len(mut self, len: usize) -> Self {
        self.scratch.resize(len, 0);
        self
    }

    /// Set wolfCrypt RNG.
    pub fn rng(mut self, rng: NonNull<crate::raw::WC_RNG>) -> Self {
        self.rng = Some(rng);
        self
    }

    /// Sign into caller output.
    pub fn sign_into<'out>(&mut self, out: &'out mut [u8]) -> Result<&'out [u8]> {
        sign1_into(
            self.key.ok_or(Error::InvalidArgument)?,
            self.algorithm.ok_or(Error::InvalidArgument)?,
            self.kid,
            self.payload.ok_or(Error::InvalidArgument)?,
            self.external_aad,
            &mut self.scratch,
            out,
            self.rng,
        )
    }

    /// Sign into a new vector.
    pub fn sign_to_vec(&mut self) -> Result<Vec<u8>> {
        sign1_to_vec(
            self.key.ok_or(Error::InvalidArgument)?,
            self.algorithm.ok_or(Error::InvalidArgument)?,
            self.kid,
            self.payload.ok_or(Error::InvalidArgument)?,
            self.external_aad,
            &mut self.scratch,
            self.rng,
        )
    }

    /// Verify a COSE_Sign1 message.
    pub fn verify<'inbuf>(
        &mut self,
        input: &'inbuf [u8],
        detached_payload: Option<&[u8]>,
    ) -> Result<VerifyOutput<'inbuf>> {
        verify1(
            self.key.ok_or(Error::InvalidArgument)?,
            input,
            detached_payload,
            self.external_aad,
            &mut self.scratch,
        )
    }
}

impl Default for Sign1Builder<'_> {
    fn default() -> Self {
        Self::new()
    }
}

/// Builder for COSE_Encrypt0 operations.
pub struct Encrypt0Builder<'a> {
    key: Option<&'a CoseKey>,
    algorithm: Option<Algorithm>,
    iv: &'a [u8],
    payload: Option<PayloadMode<'a>>,
    external_aad: &'a [u8],
    scratch: Vec<u8>,
}

impl<'a> Encrypt0Builder<'a> {
    /// Create an empty builder.
    pub fn new() -> Self {
        Self {
            key: None,
            algorithm: None,
            iv: &[],
            payload: None,
            external_aad: &[],
            scratch: vec![0; DEFAULT_SCRATCH],
        }
    }

    /// Set key.
    pub fn key(mut self, key: &'a CoseKey) -> Self {
        self.key = Some(key);
        self
    }

    /// Set content algorithm.
    pub fn algorithm(mut self, algorithm: Algorithm) -> Self {
        self.algorithm = Some(algorithm);
        self
    }

    /// Set IV.
    pub fn iv(mut self, iv: &'a [u8]) -> Self {
        self.iv = iv;
        self
    }

    /// Set payload mode.
    pub fn payload(mut self, payload: PayloadMode<'a>) -> Self {
        self.payload = Some(payload);
        self
    }

    /// Set external AAD.
    pub fn external_aad(mut self, external_aad: &'a [u8]) -> Self {
        self.external_aad = external_aad;
        self
    }

    /// Replace the internal scratch buffer size.
    pub fn scratch_len(mut self, len: usize) -> Self {
        self.scratch.resize(len, 0);
        self
    }

    /// Encrypt an attached payload into caller output.
    pub fn encrypt_into<'out>(&mut self, out: &'out mut [u8]) -> Result<&'out [u8]> {
        let payload = match self.payload.ok_or(Error::InvalidArgument)? {
            PayloadMode::Attached(payload) => payload,
            PayloadMode::Detached(_) => return Err(Error::InvalidArgument),
        };
        encrypt0_into(
            self.key.ok_or(Error::InvalidArgument)?,
            self.algorithm.ok_or(Error::InvalidArgument)?,
            self.iv,
            payload,
            self.external_aad,
            &mut self.scratch,
            out,
        )
    }

    /// Encrypt an attached payload into a new vector.
    pub fn encrypt_to_vec(&mut self) -> Result<Vec<u8>> {
        let payload = match self.payload.ok_or(Error::InvalidArgument)? {
            PayloadMode::Attached(payload) => payload,
            PayloadMode::Detached(_) => return Err(Error::InvalidArgument),
        };
        encrypt0_to_vec(
            self.key.ok_or(Error::InvalidArgument)?,
            self.algorithm.ok_or(Error::InvalidArgument)?,
            self.iv,
            payload,
            self.external_aad,
            &mut self.scratch,
        )
    }

    /// Encrypt a detached payload into caller message and ciphertext buffers.
    pub fn encrypt_detached_into(
        &mut self,
        out: &mut [u8],
        detached_ciphertext: &mut [u8],
    ) -> Result<Encrypt0DetachedOutput> {
        let payload = match self.payload.ok_or(Error::InvalidArgument)? {
            PayloadMode::Detached(payload) => payload,
            PayloadMode::Attached(_) => return Err(Error::InvalidArgument),
        };
        encrypt0_detached_into(
            self.key.ok_or(Error::InvalidArgument)?,
            self.algorithm.ok_or(Error::InvalidArgument)?,
            self.iv,
            payload,
            self.external_aad,
            &mut self.scratch,
            out,
            detached_ciphertext,
        )
    }

    /// Decrypt into caller plaintext.
    pub fn decrypt_into(
        &mut self,
        input: &[u8],
        detached_ciphertext: Option<&[u8]>,
        plaintext: &mut [u8],
    ) -> Result<DecryptOutput> {
        decrypt0_into(
            self.key.ok_or(Error::InvalidArgument)?,
            input,
            detached_ciphertext,
            self.external_aad,
            &mut self.scratch,
            plaintext,
        )
    }
}

impl Default for Encrypt0Builder<'_> {
    fn default() -> Self {
        Self::new()
    }
}

/// Builder for COSE_Mac0 operations.
pub struct Mac0Builder<'a> {
    key: Option<&'a CoseKey>,
    algorithm: Option<Algorithm>,
    kid: Option<&'a [u8]>,
    payload: Option<PayloadMode<'a>>,
    external_aad: &'a [u8],
    scratch: Vec<u8>,
}

impl<'a> Mac0Builder<'a> {
    /// Create an empty builder.
    pub fn new() -> Self {
        Self {
            key: None,
            algorithm: None,
            kid: None,
            payload: None,
            external_aad: &[],
            scratch: vec![0; DEFAULT_SCRATCH],
        }
    }

    /// Set key.
    pub fn key(mut self, key: &'a CoseKey) -> Self {
        self.key = Some(key);
        self
    }

    /// Set MAC algorithm.
    pub fn algorithm(mut self, algorithm: Algorithm) -> Self {
        self.algorithm = Some(algorithm);
        self
    }

    /// Set key identifier.
    pub fn kid(mut self, kid: &'a [u8]) -> Self {
        self.kid = Some(kid);
        self
    }

    /// Set payload mode.
    pub fn payload(mut self, payload: PayloadMode<'a>) -> Self {
        self.payload = Some(payload);
        self
    }

    /// Set external AAD.
    pub fn external_aad(mut self, external_aad: &'a [u8]) -> Self {
        self.external_aad = external_aad;
        self
    }

    /// Replace the internal scratch buffer size.
    pub fn scratch_len(mut self, len: usize) -> Self {
        self.scratch.resize(len, 0);
        self
    }

    /// Create MAC into caller output.
    pub fn mac_into<'out>(&mut self, out: &'out mut [u8]) -> Result<&'out [u8]> {
        mac0_into(
            self.key.ok_or(Error::InvalidArgument)?,
            self.algorithm.ok_or(Error::InvalidArgument)?,
            self.kid,
            self.payload.ok_or(Error::InvalidArgument)?,
            self.external_aad,
            &mut self.scratch,
            out,
        )
    }

    /// Create MAC into a new vector.
    pub fn mac_to_vec(&mut self) -> Result<Vec<u8>> {
        mac0_to_vec(
            self.key.ok_or(Error::InvalidArgument)?,
            self.algorithm.ok_or(Error::InvalidArgument)?,
            self.kid,
            self.payload.ok_or(Error::InvalidArgument)?,
            self.external_aad,
            &mut self.scratch,
        )
    }

    /// Verify MAC.
    pub fn verify<'inbuf>(
        &mut self,
        input: &'inbuf [u8],
        detached_payload: Option<&[u8]>,
    ) -> Result<MacVerifyOutput<'inbuf>> {
        verify_mac0(
            self.key.ok_or(Error::InvalidArgument)?,
            input,
            detached_payload,
            self.external_aad,
            &mut self.scratch,
        )
    }
}

impl Default for Mac0Builder<'_> {
    fn default() -> Self {
        Self::new()
    }
}

/// Builder for multi-signer COSE_Sign operations.
pub struct SignBuilder<'a> {
    signers: Vec<Signature<'a>>,
    payload: Option<PayloadMode<'a>>,
    external_aad: &'a [u8],
    scratch: Vec<u8>,
    rng: Option<NonNull<crate::raw::WC_RNG>>,
}

impl<'a> SignBuilder<'a> {
    /// Create an empty builder.
    pub fn new() -> Self {
        Self {
            signers: Vec::new(),
            payload: None,
            external_aad: &[],
            scratch: vec![0; DEFAULT_SCRATCH],
            rng: None,
        }
    }

    /// Add a signer.
    pub fn signer(mut self, signer: Signature<'a>) -> Self {
        self.signers.push(signer);
        self
    }

    /// Set payload mode.
    pub fn payload(mut self, payload: PayloadMode<'a>) -> Self {
        self.payload = Some(payload);
        self
    }

    /// Set external AAD.
    pub fn external_aad(mut self, external_aad: &'a [u8]) -> Self {
        self.external_aad = external_aad;
        self
    }

    /// Set RNG.
    pub fn rng(mut self, rng: NonNull<crate::raw::WC_RNG>) -> Self {
        self.rng = Some(rng);
        self
    }

    /// Create message into caller output.
    pub fn sign_into<'out>(&mut self, out: &'out mut [u8]) -> Result<&'out [u8]> {
        sign_into(
            &self.signers,
            self.payload.ok_or(Error::InvalidArgument)?,
            self.external_aad,
            &mut self.scratch,
            out,
            self.rng,
        )
    }

    /// Create message into a new vector.
    pub fn sign_to_vec(&mut self) -> Result<Vec<u8>> {
        sign_to_vec(
            &self.signers,
            self.payload.ok_or(Error::InvalidArgument)?,
            self.external_aad,
            &mut self.scratch,
            self.rng,
        )
    }

    /// Verify a signer.
    pub fn verify<'inbuf>(
        &mut self,
        key: &CoseKey,
        signer_index: usize,
        input: &'inbuf [u8],
        detached_payload: Option<&[u8]>,
    ) -> Result<VerifyOutput<'inbuf>> {
        verify_sign(
            key,
            signer_index,
            input,
            detached_payload,
            self.external_aad,
            &mut self.scratch,
        )
    }
}

impl Default for SignBuilder<'_> {
    fn default() -> Self {
        Self::new()
    }
}

/// Builder for multi-recipient COSE_Encrypt operations.
pub struct EncryptBuilder<'a> {
    recipients: Vec<Recipient<'a>>,
    algorithm: Option<Algorithm>,
    iv: &'a [u8],
    payload: Option<PayloadMode<'a>>,
    external_aad: &'a [u8],
    scratch: Vec<u8>,
    rng: Option<NonNull<crate::raw::WC_RNG>>,
}

impl<'a> EncryptBuilder<'a> {
    /// Create an empty builder.
    pub fn new() -> Self {
        Self {
            recipients: Vec::new(),
            algorithm: None,
            iv: &[],
            payload: None,
            external_aad: &[],
            scratch: vec![0; DEFAULT_SCRATCH],
            rng: None,
        }
    }

    /// Add recipient.
    pub fn recipient(mut self, recipient: Recipient<'a>) -> Self {
        self.recipients.push(recipient);
        self
    }

    /// Set content algorithm.
    pub fn algorithm(mut self, algorithm: Algorithm) -> Self {
        self.algorithm = Some(algorithm);
        self
    }

    /// Set IV.
    pub fn iv(mut self, iv: &'a [u8]) -> Self {
        self.iv = iv;
        self
    }

    /// Set payload mode.
    pub fn payload(mut self, payload: PayloadMode<'a>) -> Self {
        self.payload = Some(payload);
        self
    }

    /// Set external AAD.
    pub fn external_aad(mut self, external_aad: &'a [u8]) -> Self {
        self.external_aad = external_aad;
        self
    }

    /// Set RNG.
    pub fn rng(mut self, rng: NonNull<crate::raw::WC_RNG>) -> Self {
        self.rng = Some(rng);
        self
    }

    /// Encrypt attached payload into caller output.
    pub fn encrypt_into<'out>(&mut self, out: &'out mut [u8]) -> Result<&'out [u8]> {
        let payload = match self.payload.ok_or(Error::InvalidArgument)? {
            PayloadMode::Attached(payload) => payload,
            PayloadMode::Detached(_) => {
                return encrypt_detached_into(
                    &self.recipients,
                    self.algorithm.ok_or(Error::InvalidArgument)?,
                    self.iv,
                    match self.payload.ok_or(Error::InvalidArgument)? {
                        PayloadMode::Detached(payload) => payload,
                        PayloadMode::Attached(_) => unreachable!(),
                    },
                    self.external_aad,
                    &mut self.scratch,
                    out,
                    self.rng,
                )
            }
        };
        encrypt_into(
            &self.recipients,
            self.algorithm.ok_or(Error::InvalidArgument)?,
            self.iv,
            payload,
            self.external_aad,
            &mut self.scratch,
            out,
            self.rng,
        )
    }

    /// Encrypt attached payload into a new vector.
    pub fn encrypt_to_vec(&mut self) -> Result<Vec<u8>> {
        let payload = match self.payload.ok_or(Error::InvalidArgument)? {
            PayloadMode::Attached(payload) => payload,
            PayloadMode::Detached(_) => return Err(Error::InvalidArgument),
        };
        encrypt_to_vec(
            &self.recipients,
            self.algorithm.ok_or(Error::InvalidArgument)?,
            self.iv,
            payload,
            self.external_aad,
            &mut self.scratch,
            self.rng,
        )
    }

    /// Decrypt into caller plaintext.
    pub fn decrypt_into(
        &mut self,
        recipient: &Recipient<'_>,
        recipient_index: usize,
        input: &[u8],
        detached_ciphertext: Option<&[u8]>,
        plaintext: &mut [u8],
    ) -> Result<DecryptOutput> {
        decrypt_into(
            recipient,
            recipient_index,
            input,
            detached_ciphertext,
            self.external_aad,
            &mut self.scratch,
            plaintext,
        )
    }
}

impl Default for EncryptBuilder<'_> {
    fn default() -> Self {
        Self::new()
    }
}

/// Builder for multi-recipient COSE_Mac operations.
pub struct MacBuilder<'a> {
    recipients: Vec<Recipient<'a>>,
    algorithm: Option<Algorithm>,
    payload: Option<PayloadMode<'a>>,
    external_aad: &'a [u8],
    scratch: Vec<u8>,
}

impl<'a> MacBuilder<'a> {
    /// Create an empty builder.
    pub fn new() -> Self {
        Self {
            recipients: Vec::new(),
            algorithm: None,
            payload: None,
            external_aad: &[],
            scratch: vec![0; DEFAULT_SCRATCH],
        }
    }

    /// Add recipient.
    pub fn recipient(mut self, recipient: Recipient<'a>) -> Self {
        self.recipients.push(recipient);
        self
    }

    /// Set MAC algorithm.
    pub fn algorithm(mut self, algorithm: Algorithm) -> Self {
        self.algorithm = Some(algorithm);
        self
    }

    /// Set payload mode.
    pub fn payload(mut self, payload: PayloadMode<'a>) -> Self {
        self.payload = Some(payload);
        self
    }

    /// Set external AAD.
    pub fn external_aad(mut self, external_aad: &'a [u8]) -> Self {
        self.external_aad = external_aad;
        self
    }

    /// Create MAC into caller output.
    pub fn mac_into<'out>(&mut self, out: &'out mut [u8]) -> Result<&'out [u8]> {
        mac_into(
            &self.recipients,
            self.algorithm.ok_or(Error::InvalidArgument)?,
            self.payload.ok_or(Error::InvalidArgument)?,
            self.external_aad,
            &mut self.scratch,
            out,
        )
    }

    /// Create MAC into a new vector.
    pub fn mac_to_vec(&mut self) -> Result<Vec<u8>> {
        mac_to_vec(
            &self.recipients,
            self.algorithm.ok_or(Error::InvalidArgument)?,
            self.payload.ok_or(Error::InvalidArgument)?,
            self.external_aad,
            &mut self.scratch,
        )
    }

    /// Verify recipient MAC.
    pub fn verify<'inbuf>(
        &mut self,
        recipient: &Recipient<'_>,
        recipient_index: usize,
        input: &'inbuf [u8],
        detached_payload: Option<&[u8]>,
    ) -> Result<MacVerifyOutput<'inbuf>> {
        verify_mac(
            recipient,
            recipient_index,
            input,
            detached_payload,
            self.external_aad,
            &mut self.scratch,
        )
    }
}

impl Default for MacBuilder<'_> {
    fn default() -> Self {
        Self::new()
    }
}