image4-pki 0.1.0

use crate::{error::DecodeError, Signature};
use const_oid::db::rfc5912::{ID_EC_PUBLIC_KEY, RSA_ENCRYPTION};
use der::{Document, SecretDocument};
use pkcs8::DecodePrivateKey;
use rsa::{
    pkcs1::{DecodeRsaPrivateKey, DecodeRsaPublicKey},
    pkcs8::{AssociatedOid, EncodePrivateKey, PrivateKeyInfo},
};
use sec1::DecodeEcPrivateKey;
use signature::{Keypair, Signer};
use spki::{
    AlgorithmIdentifierOwned, DecodePublicKey, DynSignatureAlgorithmIdentifier, EncodePublicKey,
    SubjectPublicKeyInfoRef,
};
use std::path::Path;

macro_rules! def_key {
    (
        $name:ident, $rsa_base:ident, $ec_base:ident, $op_key:ident, $which:literal,
        $dec_res:ty, $dec_info:ty,
        $enc_trait:ident, $enc_method:ident, $enc_res:ty
    ) => {
        #[doc = concat!(
            "A ", $which, " key for any algorithm supported by the crate.\n\n",
            "This key is split from the [`", stringify!($op_key), "`] type so that it could be ",
            "decoded and encoded using the [`", stringify!($dec_trait), "`], [`",
            stringify!($dec_rsa_trait), "`] and [`", stringify!($enc_trait), "`] traits."
        )]
        #[derive(Clone, Eq, PartialEq, Debug)]
        #[non_exhaustive]
        // TODO: find a good way to remove this
        #[allow(clippy::large_enum_variant)]
        pub enum $name {
            #[doc = concat!("PKCS#1 v1.5 RSA ", $which, " key.")]
            Rsa(::rsa::$rsa_base),
            #[doc = concat!("ECDSA with NIST P-256 curve ", $which, " key.")]
            EcP256(::p256::$ec_base),
            #[doc = concat!("ECDSA with NIST P-384 curve ", $which, " key.")]
            EcP384(::p384::$ec_base),
        }

        impl From<::rsa::$rsa_base> for $name {
            fn from(value: ::rsa::$rsa_base) -> Self {
                Self::Rsa(value)
            }
        }

        impl From<::p256::$ec_base> for $name {
            fn from(value: ::p256::$ec_base) -> Self {
                Self::EcP256(value)
            }
        }

        impl From<::p384::$ec_base> for $name {
            fn from(value: ::p384::$ec_base) -> Self {
                Self::EcP384(value)
            }
        }

        impl ::core::convert::TryFrom<$dec_info> for $name {
            type Error = $dec_res;

            fn try_from(info: $dec_info) -> Result<Self, $dec_res> {
                if info.algorithm.oid == RSA_ENCRYPTION {
                    ::core::convert::TryFrom::try_from(info).map(Self::Rsa)
                } else if info.algorithm.oid == ID_EC_PUBLIC_KEY {
                    let params: ::spki::ObjectIdentifier = info
                        .algorithm
                        .parameters
                        .map(::der::asn1::AnyRef::decode_as)
                        .transpose()?
                        .ok_or(::spki::Error::AlgorithmParametersMissing)?;

                    if params == ::p256::NistP256::OID {
                        ::core::convert::TryFrom::try_from(info).map(Self::EcP256)
                    } else if params == ::p384::NistP384::OID {
                        ::core::convert::TryFrom::try_from(info).map(Self::EcP384)
                    } else {
                        ::core::result::Result::Err(::spki::Error::OidUnknown {
                            oid: params
                        }.into())
                    }
                } else {
                    ::core::result::Result::Err(::spki::Error::OidUnknown {
                        oid: info.algorithm.oid,
                    }
                    .into())
                }
            }
        }

        impl $enc_trait for $name {
            fn $enc_method(&self) -> $enc_res {
                match self {
                    Self::Rsa(key) => key.$enc_method(),
                    Self::EcP256(key) => key.$enc_method(),
                    Self::EcP384(key) => key.$enc_method(),
                }
            }
        }
    };
}

def_key!(
    PrivateKey,
    RsaPrivateKey,
    SecretKey,
    SigningKey,
    "private",
    pkcs8::Error,
    PrivateKeyInfo<'_>,
    EncodePrivateKey,
    to_pkcs8_der,
    pkcs8::Result<SecretDocument>
);

impl PrivateKey {
    fn from_pem_label_and_doc(label: String, doc: &SecretDocument) -> Result<Self, DecodeError> {
        match label.as_str() {
            "RSA PRIVATE KEY" => Ok(PrivateKey::from_pkcs1_der(doc.as_bytes())?),
            "PRIVATE KEY" => Ok(PrivateKey::from_pkcs8_der(doc.as_bytes())?),
            "EC PRIVATE KEY" => Ok(PrivateKey::from_sec1_der(doc.as_bytes())?),
            _ => Err(DecodeError::UnexpectedLabel(label)),
        }
    }

    /// Tries to parse a PEM file at the specified path as a private key.
    pub fn read_pem_file(path: impl AsRef<Path>) -> Result<Self, DecodeError> {
        let (label, doc) = SecretDocument::read_pem_file(path)?;

        Self::from_pem_label_and_doc(label, &doc)
    }

    /// Tries to parse an in-memory PEM file as a private key.
    pub fn from_pem(pem: &str) -> Result<Self, DecodeError> {
        let (label, doc) = SecretDocument::from_pem(pem)?;

        Self::from_pem_label_and_doc(label.to_string(), &doc)
    }

    /// Creates a public key from a private key.
    pub fn to_public_key(&self) -> PublicKey {
        match self {
            PrivateKey::Rsa(key) => PublicKey::Rsa(key.to_public_key()),
            PrivateKey::EcP256(key) => PublicKey::EcP256(key.public_key()),
            PrivateKey::EcP384(key) => PublicKey::EcP384(key.public_key()),
        }
    }
}

def_key!(
    PublicKey,
    RsaPublicKey,
    PublicKey,
    VerifyingKey,
    "public",
    spki::Error,
    SubjectPublicKeyInfoRef<'_>,
    EncodePublicKey,
    to_public_key_der,
    spki::Result<Document>
);

impl PublicKey {
    fn from_pem_label_and_doc(label: &str, doc: &Document) -> spki::Result<Self> {
        let bytes = doc.as_bytes();

        match label {
            "RSA PUBLIC KEY" => Ok(PublicKey::from_pkcs1_der(bytes)?),
            "PUBLIC KEY" => PublicKey::from_public_key_der(bytes),
            _ => Err(spki::Error::Asn1(
                der::pem::Error::UnexpectedTypeLabel {
                    expected: "PUBLIC KEY\" or \"RSA PUBLIC KEY",
                }
                .into(),
            )),
        }
    }

    /// Tries to parse a PEM file at the specified path as a public key.
    pub fn read_pem_file(path: impl AsRef<Path>) -> spki::Result<Self> {
        let (label, doc) = Document::read_pem_file(path)?;

        Self::from_pem_label_and_doc(&label, &doc)
    }

    /// Tries to parse an in-memory PEM file as a public key.
    pub fn from_pem(pem: &str) -> spki::Result<Self> {
        let (label, doc) = Document::from_pem(pem)?;

        Self::from_pem_label_and_doc(label, &doc)
    }
}

macro_rules! def_op_key {
    ($name:ident, $base:ident, $what:literal, $enc_trait:ident, $enc_method:ident, $enc_res:ty) => {
        #[derive(Clone, Debug)]
        #[doc = concat!("A ", $what, " key for any algorithm supported by the crate.")]
        #[non_exhaustive]
        pub enum $name {
            #[doc = concat!("PKCS#1 v1.5 RSA ", $what, " key which uses SHA-1 as the digest.")]
            RsaWithSha1(::rsa::pkcs1v15::$name<::sha1::Sha1>),
            #[doc = concat!("PKCS#1 v1.5 RSA ", $what, " key which uses SHA-256 as the digest.")]
            RsaWithSha256(::rsa::pkcs1v15::$name<::sha2::Sha256>),
            #[doc = concat!("PKCS#1 v1.5 RSA ", $what, " key which uses SHA-384 as the digest.")]
            RsaWithSha384(::rsa::pkcs1v15::$name<::sha2::Sha384>),
            #[doc = concat!("ECDSA/P-256 ", $what, " key (uses SHA-256 as the digest primitive).")]
            EcP256(::p256::ecdsa::$name),
            #[doc = concat!("ECDSA/P-384 ", $what, " key (uses SHA-384 as the digest primitive).")]
            EcP384(::p384::ecdsa::$name),
        }

        impl $name {
            #[doc = concat!(
                "Creates a new [`", stringify!($name), "`] from a [`", stringify!($base), "`] and",
                " a [`DigestAlgo`]."
            )]
            /// If no digest algorithm is provided, a default one is used.
            ///
            /// For RSA PKCS#1 v1.5 signatures the default is SHA-384, for elliptic curve algorithms the
            /// default and the only supported digest algorithms are SHA-256 for NIST P-256 and SHA-384 for
            /// NIST P-384.
            ///
            /// [`DigestAlgo`]: crate::DigestAlgo
            pub fn from_key_and_algo(
                key: $base,
                algo: ::core::option::Option<$crate::DigestAlgo>
            ) -> ::core::result::Result<Self, $crate::error::UnsupportedDigest> {
                match (key, algo) {
                    ($base::Rsa(key), Some($crate::DigestAlgo::Sha1)) => Ok(Self::RsaWithSha1(::rsa::pkcs1v15::$name::new(key))),
                    ($base::Rsa(key), Some($crate::DigestAlgo::Sha256)) => Ok(Self::RsaWithSha256(::rsa::pkcs1v15::$name::new(key))),
                    ($base::Rsa(key), Some($crate::DigestAlgo::Sha384) | None) => Ok(Self::RsaWithSha384(::rsa::pkcs1v15::$name::new(key))),
                    ($base::EcP256(key), Some($crate::DigestAlgo::Sha256) | None) => Ok(Self::EcP256(key.into())),
                    ($base::EcP384(key), Some($crate::DigestAlgo::Sha384) | None) => Ok(Self::EcP384(key.into())),
                    ($base::EcP256(_), Some(other)) => Err($crate::error::UnsupportedDigest {
                        signing_algo: $crate::SigningAlgo::EcP256,
                        digest_algo: other,
                    }),
                    ($base::EcP384(_), Some(other)) => Err($crate::error::UnsupportedDigest {
                        signing_algo: $crate::SigningAlgo::EcP384,
                        digest_algo: other,
                    })
                }
            }
        }

        impl $enc_trait for $name {
            fn $enc_method(&self) -> $enc_res {
                match self {
                    Self::RsaWithSha1(key) => key.$enc_method(),
                    Self::RsaWithSha256(key) => key.$enc_method(),
                    Self::RsaWithSha384(key) => key.$enc_method(),
                    Self::EcP256(key) => key.$enc_method(),
                    Self::EcP384(key) => key.$enc_method(),
                }
            }
        }

        impl From<$name> for $base {
            fn from(value: $name) -> Self {
                match value {
                    $name::RsaWithSha1(key) => Self::Rsa(key.into()),
                    $name::RsaWithSha256(key) => Self::Rsa(key.into()),
                    $name::RsaWithSha384(key) => Self::Rsa(key.into()),
                    $name::EcP256(key) => Self::EcP256(key.into()),
                    $name::EcP384(key) => Self::EcP384(key.into()),
                }
            }
        }
    };
}

def_op_key!(
    SigningKey,
    PrivateKey,
    "signing",
    EncodePrivateKey,
    to_pkcs8_der,
    pkcs8::Result<SecretDocument>
);

def_op_key!(
    VerifyingKey,
    PublicKey,
    "verifying",
    EncodePublicKey,
    to_public_key_der,
    spki::Result<Document>
);

impl Signer<Signature> for SigningKey {
    fn try_sign(&self, msg: &[u8]) -> signature::Result<Signature> {
        match self {
            SigningKey::RsaWithSha1(key) => key.try_sign(msg).map(Into::into),
            SigningKey::RsaWithSha256(key) => key.try_sign(msg).map(Into::into),
            SigningKey::RsaWithSha384(key) => key.try_sign(msg).map(Into::into),
            SigningKey::EcP256(key) => {
                <p256::ecdsa::SigningKey as Signer<p256::ecdsa::Signature>>::try_sign(key, msg)
                    .map(Into::into)
            }
            SigningKey::EcP384(key) => {
                <p384::ecdsa::SigningKey as Signer<p384::ecdsa::Signature>>::try_sign(key, msg)
                    .map(Into::into)
            }
        }
    }
}

impl Keypair for SigningKey {
    type VerifyingKey = VerifyingKey;

    fn verifying_key(&self) -> Self::VerifyingKey {
        match self {
            SigningKey::RsaWithSha1(key) => VerifyingKey::RsaWithSha1(key.verifying_key()),
            SigningKey::RsaWithSha256(key) => VerifyingKey::RsaWithSha256(key.verifying_key()),
            SigningKey::RsaWithSha384(key) => VerifyingKey::RsaWithSha384(key.verifying_key()),
            SigningKey::EcP256(key) => VerifyingKey::EcP256(*key.verifying_key()),
            SigningKey::EcP384(key) => VerifyingKey::EcP384(*key.verifying_key()),
        }
    }
}

impl DynSignatureAlgorithmIdentifier for SigningKey {
    fn signature_algorithm_identifier(&self) -> spki::Result<AlgorithmIdentifierOwned> {
        match self {
            SigningKey::RsaWithSha1(key) => key.signature_algorithm_identifier(),
            SigningKey::RsaWithSha256(key) => key.signature_algorithm_identifier(),
            SigningKey::RsaWithSha384(key) => key.signature_algorithm_identifier(),
            SigningKey::EcP256(key) => key.signature_algorithm_identifier(),
            SigningKey::EcP384(key) => key.signature_algorithm_identifier(),
        }
    }
}