tss-keyfile 0.1.3

use alloc::vec::Vec;

#[cfg(feature = "pem")]
use alloc::string::String;

use hex::ToHex;
use rasn::prelude::*;

/// A TPM loadable key, to be loaded with `TPM2_Load`
pub const OID_TPMKEY_LOADABLE: &Oid = Oid::const_new(&[2, 23, 133, 10, 1, 3]);

/// A TPM importable key, to be imported with `TPM2_Import`
pub const OID_TPMKEY_IMPORTABLE: &Oid = Oid::const_new(&[2, 23, 133, 10, 1, 4]);

/// A sealed data key, to be extracted with `TPM2_Unseal`
pub const OID_TPMKEY_SEALED: &Oid = Oid::const_new(&[2, 23, 133, 10, 1, 5]);

/// TPM Key Type (OID encoding)
#[derive(AsnType, Clone, Debug, Decode, Encode, PartialEq, Eq, Hash)]
#[rasn(delegate)]
pub struct TPMKeyType(ObjectIdentifier);

impl TPMKeyType {
    pub fn new_loadable() -> Self {
        Self(ObjectIdentifier::from(OID_TPMKEY_LOADABLE))
    }

    pub fn new_importable() -> Self {
        Self(ObjectIdentifier::from(OID_TPMKEY_IMPORTABLE))
    }

    pub fn new_sealed() -> Self {
        Self(ObjectIdentifier::from(OID_TPMKEY_SEALED))
    }

    pub fn validate(&self) -> bool {
        // validate OID
        self.0 == OID_TPMKEY_LOADABLE
            || self.0 == OID_TPMKEY_IMPORTABLE
            || self.0 == OID_TPMKEY_SEALED
    }
}

impl AsRef<ObjectIdentifier> for TPMKeyType {
    fn as_ref(&self) -> &ObjectIdentifier {
        &self.0
    }
}

impl core::fmt::Display for TPMKeyType {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let typestr = if self.0 == OID_TPMKEY_LOADABLE {
            "Loadable Key"
        } else if self.0 == OID_TPMKEY_IMPORTABLE {
            "Importable Key"
        } else if self.0 == OID_TPMKEY_SEALED {
            "Sealed Data"
        } else {
            "Unknown"
        };
        write!(f, "{typestr} (OID:{})", &self.0)
    }
}

#[derive(AsnType, Clone, Debug, Decode, Encode, PartialEq, Eq, Hash)]
pub struct TPMPolicy {
    /// This is the integer representation of the TPM command code for the policy statement.
    #[rasn(tag(explicit(0)), default)]
    pub command_code: u32,

    /// This is a binary string representing a fully marshalled, TPM ordered, command body
    /// for the TPM policy command. Therefore to send the command, the implementation simply
    /// marshals the command code and appends this octet string as the body.
    ///
    /// Commands which have no body, such as TPM2_AuthVal, MUST be specified as a zero length
    /// OCTET STRING
    ///
    /// Note that there are some commands for which the simple body of the TPM policy command
    /// does not provide enough information to execute the policy command. A classic example
    /// is TPM2_PolicyAuthorize, whose body consists of a key name, a policyRef nonce and a
    /// signature. However, the implementation needs to know the actual key, not just the name,
    /// to implement the policy
    #[rasn(tag(explicit(1)), default)]
    pub command_policy: OctetString,
}

impl core::fmt::Display for TPMPolicy {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(
            f,
            "TPMPolicy {{ command_code: 0x{:08X}, command_policy: {} bytes }}",
            self.command_code,
            self.command_policy.len()
        )
    }
}

#[derive(AsnType, Clone, Debug, Decode, Encode, PartialEq, Eq, Hash)]
pub struct TPMAuthPolicy {
    /// An optional string name for the current policy which is only used for display purposes,
    /// MAY be used as a user visible mnemonic for the actual policy.
    #[rasn(tag(explicit(0)), default)]
    pub name: Option<Utf8String>,

    /// A sequence of TPMPolicy statements which MUST end with a PolicyAuthorize statement whose
    /// signature is over the hash of the current policy excluding the PolicyAuthorize statement
    /// and the nonce specified at the beginning of the TPMKey policy.
    /// There MUST be no other TPM2_PolicyAuthorize statements in the intermediate policy steps.
    /// There MAY be an initial TPM2_PolicyAuthorize statement containing a different public key
    /// (because this causes complexity building signed policy chains, implementations MAY choose
    /// to allow only a single policy signing key, in which case there MAY NOT be an initial
    /// TPM2_PolicyAuthorize statement).
    #[rasn(tag(explicit(1)), default)]
    pub policy: SequenceOf<TPMPolicy>,
}

impl core::fmt::Display for TPMAuthPolicy {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "TPMAuthPolicy {{")?;
        if let Some(name) = &self.name {
            write!(f, " name: \"{}\",", name)?;
        }
        write!(f, " policy: {} entries }}", self.policy.len())
    }
}

/// The design of the TPMkey ASN.1 format is that it should have a distinguishing OID
/// at the beginning so the DER form of the key can be easily recognized.
/// In PEM form, the key MUST have the following PEM guards:
/// "-----BEGIN TSS2 PRIVATE KEY-----" and "-----END TSS2 PRIVATE KEY-----"
///
/// All additional information that may be needed to load the key is specified as
/// optional explicit elements, which can be extended by later specifications,
/// which is why the TPMkey is not versioned.
#[derive(AsnType, Clone, Debug, Decode, Encode, PartialEq, Eq, Hash)]
pub struct TPMKey {
    /// A unique OID specifying the key type:
    /// - TPMKEY_LOADABLE
    /// - TPMKEY_IMPORTABLE
    /// - TPMKEY_SEALED
    pub r#type: TPMKeyType,

    /// An implementation needs to know as it formulates the TPM2_Load/Import/Unseal
    /// command whether it must also send down an authorization. So this parameter
    /// gives that indication.
    /// emptyAuth MUST be true if authorization is NOT required and MUST be either
    /// false or absent if authorization is required.
    ///
    /// Since this element has three states (one representing true and two representing false)
    /// it is RECOMMENDED that implementations emitting TPMkey representations use absence of
    /// the tag to represent false. However, implementations reading TPMKey MUST be able to
    /// process all three possible states.
    #[rasn(tag(explicit(0)), default)]
    pub empty_auth: Option<bool>,

    /// This MUST be present if the TPM key has a policy hash because it describes to the
    /// implementation how to construct the policy
    #[rasn(tag(explicit(1)), default)]
    pub policy: Option<SequenceOf<TPMPolicy>>,

    /// secret contains the additional cryptographic secret used to specify the
    /// outer wrapping of an importable object. For keys, it MUST be present for key
    /// type id-importablekey and MUST NOT be present for key type id-loadablekey.
    /// For sealed data objects of type id-sealedkey, it MAY be present and if present
    /// indicates the object is importable.
    ///
    /// Importable objects (designed to be processed by TPM2_Import) MUST have an
    /// unencrypted inner wrapper (symmetricAlg MUST be TPM_ALG_NULL and encryptionKey
    /// MUST be empty) and an outer wrapper encrypted to the parent key using inSymSeed.
    /// The secret parameter is the fully marshalled TPM2B_ENCRYPTED_SECRET form of inSymSeed.
    #[rasn(tag(explicit(2)), default)]
    pub secret: Option<OctetString>,

    /// This SHOULD be present if the TPMkey policy contains a TPM2_PolicyAuthorize
    /// statement because it contains signed policies that could be used to satisfy
    /// the TPM key policy.
    /// If the TPM key has no policy hash then this MUST NOT be present.
    #[rasn(tag(explicit(3)), default)]
    pub auth_policy: Option<SequenceOf<TPMAuthPolicy>>,

    /// An optional string description for the key which is only used for display purposes,
    /// MAY be used as a user visible mnemonic for the key.
    #[rasn(tag(explicit(4)), default)]
    pub description: Option<Utf8String>,

    /// This MUST be present and true if the parent is a permanent handle (MSO 0x40)
    /// and RSA 2048 is used for the primary key.
    /// If the parent is not a permanent handle then this MUST NOT be present.
    /// If the parent is a permanent handle and if P-256 is used for the primary then
    /// this MUST NOT be present. Given that P-256 primary keys are easier to generate,
    /// implementations SHOULD NOT set this flag.
    #[rasn(tag(explicit(5)), default)]
    pub rsa_parent: Option<bool>,

    /// This MUST be present for all keys and specifies the handle of the parent key.
    /// The parent key SHOULD be either a persistent handle (MSO 0x81) or a permanent
    /// handle (MSO 0x40).
    /// Since volatile handle numbering can change unexpectedly depending on key load order,
    /// the parent SHOULD NOT be a volatile handle (MSO 0x80).
    /// The parent MUST NOT have any other MSO.
    pub parent: u32,

    /// This MUST be present and MUST correspond to the fully marshalled
    /// TPM2B_PUBLIC structure of the TPM Key.
    pub pubkey: OctetString,

    /// This MUST be present and MUST correspond to the fully marshalled
    /// TPM2B_PRIVATE structure of the TPM Key.
    /// For importable keys, this must be the duplicate parameter that would be input to TPM2_Import.
    pub privkey: OctetString,
}

impl core::fmt::Display for TPMKey {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        writeln!(f, "TPMKey {{")?;
        writeln!(f, "  type:  \"{}\"", self.r#type)?;
        if let Some(empty_auth) = &self.empty_auth {
            writeln!(f, "  empty_auth:  \"{}\"", empty_auth)?;
        }
        if let Some(policy) = &self.policy {
            writeln!(f, "  policy: {} entries", policy.len())?;
        }
        if let Some(secret) = &self.secret {
            writeln!(f, "  secret: {} bytes", secret.len())?;
        }
        if let Some(auth_policy) = &self.auth_policy {
            writeln!(f, "  auth_policy: {} entries", auth_policy.len())?;
        }
        if let Some(description) = &self.description {
            writeln!(f, "  description: \"{}\"", description)?;
        }
        if let Some(rsa_parent) = &self.rsa_parent {
            writeln!(f, "  rsa_parent:  \"{}\"", rsa_parent)?;
        }
        writeln!(f, "  parent:  \"0x{:08X}\"", self.parent)?;
        write!(f, "  pubkey:  \"")?;
        for b in self.pubkey.iter() {
            write!(f, "{:02X}", b)?;
        }
        writeln!(f, "\"")?;
        write!(f, "  privkey:  \"")?;
        for b in self.privkey.iter() {
            write!(f, "{:02X}", b)?;
        }
        writeln!(f, "\"")?;
        write!(f, "}}")
    }
}

#[cfg(feature = "pem")]
impl TPMKey {
    pub fn try_from_pem(pem: &[u8]) -> Result<TPMKey, crate::Error> {
        let (label, der) = pem_rfc7468::decode_vec(pem)?;

        if label != "TSS2 PRIVATE KEY" {
            return Err(crate::Error::PemNoTss2PrivateKey);
        }

        TPMKey::try_from_der(&der)
    }

    pub fn try_to_pem(&self) -> Result<String, crate::Error> {
        let der = self.try_to_der()?;

        let pem =
            pem_rfc7468::encode_string("TSS2 PRIVATE KEY", pem_rfc7468::LineEnding::LF, &der)?;

        Ok(pem)
    }
}

impl TPMKey {
    pub fn try_from_der(der: &[u8]) -> Result<TPMKey, crate::Error> {
        let tpmkey: TPMKey = rasn::der::decode(der)?;
        tpmkey.validate()?;
        Ok(tpmkey)
    }

    pub fn try_to_der(&self) -> Result<Vec<u8>, crate::Error> {
        Ok(rasn::der::encode(self)?)
    }

    pub fn new_loadable(parent: u32, pubkey: &[u8], privkey: &[u8]) -> Self {
        Self {
            r#type: TPMKeyType::new_loadable(),
            parent,
            pubkey: pubkey.into(),
            privkey: privkey.into(),
            empty_auth: None,
            policy: None,
            secret: None,
            auth_policy: None,
            description: None,
            rsa_parent: None,
        }
    }

    pub fn new_importable(parent: u32, pubkey: &[u8], privkey: &[u8]) -> Self {
        Self {
            r#type: TPMKeyType::new_importable(),
            parent,
            pubkey: pubkey.into(),
            privkey: privkey.into(),
            empty_auth: None,
            policy: None,
            secret: None,
            auth_policy: None,
            description: None,
            rsa_parent: None,
        }
    }

    pub fn new_sealed(parent: u32, pubkey: &[u8], privkey: &[u8]) -> Self {
        Self {
            r#type: TPMKeyType::new_sealed(),
            parent,
            pubkey: pubkey.into(),
            privkey: privkey.into(),
            empty_auth: None,
            policy: None,
            secret: None,
            auth_policy: None,
            description: None,
            rsa_parent: None,
        }
    }

    pub fn is_loadable(&self) -> bool {
        self.r#type.0 == OID_TPMKEY_LOADABLE
    }

    pub fn is_importable(&self) -> bool {
        self.r#type.0 == OID_TPMKEY_IMPORTABLE
    }

    pub fn is_sealed(&self) -> bool {
        self.r#type.0 == OID_TPMKEY_SEALED
    }

    pub fn requires_auth(&self) -> bool {
        !self.empty_auth.unwrap_or(false)
    }

    pub fn validate(&self) -> Result<(), crate::Error> {
        // validate OID
        if !self.r#type.validate() {
            return Err(crate::Error::NoTss2PrivateKey);
        };

        // validate parent handle is permanent, not unknown
        // transient is discouraged but allowed
        let mso = (self.parent >> 24) & 0xFF;
        if mso != 0x40 && mso != 0x81 && mso != 0x80 {
            return Err(crate::Error::InvalidParentHandle(
                self.parent.to_be_bytes().encode_hex(),
            ));
        };

        // Validate secret field requirements based on key type
        if self.is_loadable() && self.secret.is_some() {
            return Err(crate::Error::UnexpectedSecret);
        }

        if self.is_importable() && self.secret.is_none() {
            return Err(crate::Error::MissingSecret);
        }

        Ok(())
    }
}