dat 2.4.1

use crate::error::DatError;
use aws_lc_rs::signature::{EcdsaKeyPair, UnparsedPublicKey, ECDSA_P256_SHA256_FIXED, ECDSA_P256_SHA256_FIXED_SIGNING, ECDSA_P384_SHA384_FIXED, ECDSA_P384_SHA384_FIXED_SIGNING, ECDSA_P521_SHA512_FIXED, ECDSA_P521_SHA512_FIXED_SIGNING};
use std::str::FromStr;
use strum_macros::{Display, EnumString};

#[repr(u8)]
#[derive(Debug, Display, Clone, Copy, Eq, PartialEq)]
pub enum DatSignatureAlgorithm {
    P256,
    P384,
    P521,
}

impl FromStr for DatSignatureAlgorithm {
    type Err = DatError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "P256" => Ok(DatSignatureAlgorithm::P256),
            "P384" => Ok(DatSignatureAlgorithm::P384),
            "P521" => Ok(DatSignatureAlgorithm::P521),
            _ => Err(DatError::UnknownSignatureAlgorithm),
        }
    }
}


#[repr(u8)]
#[derive(Debug, Display, EnumString, Clone, Copy, Eq, PartialEq)]
pub enum DatSignatureKeyExportOption {
    PAIR,
    SIGNING,
    VERIFYING,
}

pub struct DatSignatureKey {
    algorithm: DatSignatureAlgorithm,
    signing_key: Option<EcdsaKeyPair>,
    signing_key_bytes: Vec<u8>,
    verifying_key: UnparsedPublicKey<Vec<u8>>,
}

impl DatSignatureKey {
    pub fn generate(algorithm: DatSignatureAlgorithm) -> Result<Self, DatError> {
        Self::from_or_new(true, algorithm, &[], &[])
    }

    pub fn from_bytes(algorithm: DatSignatureAlgorithm, signing_key_bytes: &[u8], verifying_key_bytes: &[u8]) -> Result<DatSignatureKey, DatError> {
        Self::from_or_new(false, algorithm, signing_key_bytes, verifying_key_bytes)
    }
    fn from_or_new(new: bool, algorithm: DatSignatureAlgorithm, signing_key_bytes: &[u8], verifying_key_bytes: &[u8]) -> Result<DatSignatureKey, DatError> {
        let (sa, va) = match algorithm {
            DatSignatureAlgorithm::P256 => &(ECDSA_P256_SHA256_FIXED_SIGNING, ECDSA_P256_SHA256_FIXED),
            DatSignatureAlgorithm::P384 => &(ECDSA_P384_SHA384_FIXED_SIGNING, ECDSA_P384_SHA384_FIXED),
            DatSignatureAlgorithm::P521 => &(ECDSA_P521_SHA512_FIXED_SIGNING, ECDSA_P521_SHA512_FIXED),
        };
        let mut signing_key_bytes = Box::from(signing_key_bytes);
        let mut verifying_key_bytes = Box::from(verifying_key_bytes);

        if new {
            let pkcs8v1_pair = EcdsaKeyPair::generate(sa)
                .map_err(|_| DatError::GenerateSigningKeyError)?
                .to_pkcs8v1()
                .map_err(|_| DatError::GenerateSigningKeyError)?;
            let (sk, vk) = to_raw_pkcs8v1(algorithm, pkcs8v1_pair.as_ref())?;
            signing_key_bytes = sk;
            verifying_key_bytes = vk;
        } else if verifying_key_bytes.is_empty() {
            if signing_key_bytes.is_empty() {
                return Err(DatError::InvalidSignatureKey);
            }
            verifying_key_bytes = calc_public_key(algorithm, &*signing_key_bytes)?;
        }

        let signing_key: Option<EcdsaKeyPair> = if signing_key_bytes.is_empty() {
            None
        } else {
            Some(EcdsaKeyPair::from_private_key_and_public_key(sa, &*signing_key_bytes, &*verifying_key_bytes)
                .map_err(|_| { DatError::InvalidSignatureKey })?)
        };

        Ok(DatSignatureKey {
            algorithm,
            signing_key,
            signing_key_bytes: Vec::from(signing_key_bytes),
            verifying_key: UnparsedPublicKey::new(va, Vec::from(verifying_key_bytes)),
        })
    }

    #[inline]
    pub fn algorithm(&self) -> DatSignatureAlgorithm {
        self.algorithm
    }

    #[inline]
    pub fn signature_size(&self) -> usize {
        match self.algorithm {
            DatSignatureAlgorithm::P256 => 64,
            DatSignatureAlgorithm::P384 => 96,
            DatSignatureAlgorithm::P521 => 132,
        }
    }

    #[inline]
    pub fn raw_signing_key_size(&self) -> usize {
        match self.algorithm {
            DatSignatureAlgorithm::P256 => 32,
            DatSignatureAlgorithm::P384 => 48,
            DatSignatureAlgorithm::P521 => 66,
        }
    }

    #[inline]
    pub fn to_bytes(&self) -> (Box<[u8]>, Box<[u8]>) {
        (Box::from(&*self.signing_key_bytes), Box::from(&*self.verifying_key.as_ref()))
    }

    pub fn sign(&self, data: &[u8]) -> Result<Box<[u8]>, DatError> {
        self.signing_key.as_ref().ok_or_else(|| DatError::VerifyOnlyCertificate)?
            .sign(&aws_lc_rs::rand::SystemRandom::new(), data)
            .map(|x| Box::from(x.as_ref()))
            .map_err(|_| DatError::SignError)
    }

    pub fn verify(&self, body: &[u8], sign: &[u8]) -> Result<(), DatError> {
        self.verifying_key.verify(body, sign).map_err(|_| DatError::InvalidDat)
    }

    #[inline]
    pub fn has_signing_key(&self) -> bool {
        self.signing_key.is_some()
    }
}

impl Clone for DatSignatureKey {
    fn clone(&self) -> Self {
        let (k, v) = self.to_bytes();
        DatSignatureKey::from_bytes(self.algorithm(), &*k, &*v).unwrap()
    }
}

pub fn calc_public_key(algorithm: DatSignatureAlgorithm, signing_key: &[u8]) -> Result<Box<[u8]>, DatError> {
    match algorithm {
        DatSignatureAlgorithm::P256 => { p256::ecdsa::SigningKey::from_slice(signing_key).map(|x| x.verifying_key().to_sec1_bytes()) },
        DatSignatureAlgorithm::P384 => { p384::ecdsa::SigningKey::from_slice(signing_key).map(|x| x.verifying_key().to_sec1_bytes()) },
        DatSignatureAlgorithm::P521 => { p521::ecdsa::SigningKey::from_slice(signing_key).map(|x| x.verifying_key().to_sec1_bytes()) },
    }.map_err(|_| DatError::InvalidSignatureKey)
}

pub fn to_raw_pkcs8v1(algorithm: DatSignatureAlgorithm, pkcs8_der: &[u8]) -> Result<(Box<[u8]>, Box<[u8]>), DatError> {
    let (offset, private_len, public_len) = match algorithm {
        DatSignatureAlgorithm::P256 => (36, 68, 65),
        DatSignatureAlgorithm::P384 => (35, 83 ,97),
        DatSignatureAlgorithm::P521 => (35, 101, 133),
    };
    if pkcs8_der.len() > public_len {
        Ok((Box::from(&pkcs8_der[offset..private_len]), Box::from(&pkcs8_der[pkcs8_der.len() - public_len..])))
    } else {
        Err(DatError::InvalidSignatureKey)
    }
}