//! Identities and composite Ed25519 + ML-DSA-65 signatures (spec §14.2, §14.3).

//!

//! Both halves sign the same message, which binds the composite `KeyId`, and

//! verification requires both. A signature with either half missing, wrong,

//! or made for a different KeyId is invalid.

use crate::cbor::{self, Value};

use crate::hash::{hash_object, KeyId, ObjectRef, ObjectType};

use crate::objects::{

    array, bytes, check_version, field, oref, text, uint, ObjectError, FORMAT_VERSION,

};

use ed25519_dalek::ed25519::signature::Signer as _;

use ed25519_dalek::{

    Signature as EdSignature, SigningKey as EdSigningKey, VerifyingKey as EdVerifyingKey,

};

use ml_dsa::{

    EncodedSignature, EncodedVerifyingKey, MlDsa65, Seed, Signature as MlSignature,

    SigningKey as MlSigningKey, VerifyingKey as MlVerifyingKey,

};

use rand::RngCore;

pub const ED25519_PK_LEN: usize = 32;

pub const MLDSA65_PK_LEN: usize = 1952;

pub const ED25519_SIG_LEN: usize = 64;

pub const MLDSA65_SIG_LEN: usize = 3309;

/// `alg` value for Ed25519 + ML-DSA-65.

pub const ALG_COMPOSITE: u64 = 1;

#[derive(Clone, Debug, PartialEq, Eq)]

pub struct Identity {

    pub ed25519_pk: [u8; ED25519_PK_LEN],

    pub mldsa65_pk: Vec<u8>,

    pub name: String,

    pub email: String,

    pub created: u64,

    pub expires: Option<u64>,

    pub previous: Option<KeyId>,

    pub delegates: Vec<KeyId>,

}

impl Identity {

    pub fn to_value(&self) -> Value {

        Value::map([

            (0, Value::Uint(FORMAT_VERSION)),

            (1, Value::Bytes(self.ed25519_pk.to_vec())),

            (2, Value::Bytes(self.mldsa65_pk.clone())),

            (3, Value::Text(self.name.clone())),

            (4, Value::Text(self.email.clone())),

            (5, Value::Uint(self.created)),

            (6, self.expires.map(Value::Uint).unwrap_or(Value::Null)),

            (

                7,

                self.previous

                    .map(|p| Value::Bytes(p.0.to_vec()))

                    .unwrap_or(Value::Null),

            ),

            (

                8,

                Value::Array(

                    self.delegates

                        .iter()

                        .map(|d| Value::Bytes(d.0.to_vec()))

                        .collect(),

                ),

            ),

        ])

    }

    pub fn encode(&self) -> Vec<u8> {

        cbor::encode(&self.to_value())

    }

    pub fn hash(&self) -> KeyId {

        hash_object(ObjectType::Identity, &self.encode())

    }

    pub fn decode(b: &[u8]) -> Result<Identity, ObjectError> {

        let v = cbor::decode(b)?;

        let m = v

            .as_map()

            .ok_or_else(|| ObjectError::Schema("identity: expected map".into()))?;

        check_version(m)?;

        let ed = bytes(field(m, 1, "ed25519_pk")?, "ed25519_pk")?;

        let ml = bytes(field(m, 2, "mldsa65_pk")?, "mldsa65_pk")?;

        if ed.len() != ED25519_PK_LEN || ml.len() != MLDSA65_PK_LEN {

            return Err(ObjectError::Schema(

                "identity: wrong public key length".into(),

            ));

        }

        let opt_u = |v: &Value| -> Result<Option<u64>, ObjectError> {

            if v.is_null() {

                Ok(None)

            } else {

                uint(v, "expires").map(Some)

            }

        };

        let opt_r = |v: &Value| -> Result<Option<KeyId>, ObjectError> {

            if v.is_null() {

                Ok(None)

            } else {

                oref(v, "previous").map(Some)

            }

        };

        let mut delegates = array(field(m, 8, "delegates")?, "delegates")?

            .iter()

            .map(|d| oref(d, "delegate"))

            .collect::<Result<Vec<_>, _>>()?;

        delegates.sort();

        delegates.dedup();

        Ok(Identity {

            ed25519_pk: ed.try_into().unwrap(),

            mldsa65_pk: ml.to_vec(),

            name: text(field(m, 3, "name")?, "name")?.to_owned(),

            email: text(field(m, 4, "email")?, "email")?.to_owned(),

            created: uint(field(m, 5, "created")?, "created")?,

            expires: opt_u(field(m, 6, "expires")?)?,

            previous: opt_r(field(m, 7, "previous")?)?,

            delegates,

        })

    }

    /// Whether the identity is past its expiry at `now_ms`.

    pub fn is_expired_at(&self, now_ms: u64) -> bool {

        self.expires.is_some_and(|e| now_ms > e)

    }

    /// A successor identity with the same keys, `previous` set to this

    /// identity's KeyId, and the given delegates and expiry (spec §14.2).

    pub fn successor(&self, delegates: Vec<KeyId>, expires: Option<u64>, created: u64) -> Identity {

        let mut delegates = delegates;

        delegates.sort();

        delegates.dedup();

        Identity {

            previous: Some(self.hash()),

            delegates,

            expires,

            created,

            ..self.clone()

        }

    }

    fn ed_key(&self) -> Result<EdVerifyingKey, VerifyError> {

        EdVerifyingKey::from_bytes(&self.ed25519_pk).map_err(|_| VerifyError::BadKey)

    }

    fn ml_key(&self) -> Result<MlVerifyingKey<MlDsa65>, VerifyError> {

        let enc = EncodedVerifyingKey::<MlDsa65>::try_from(self.mldsa65_pk.as_slice())

            .map_err(|_| VerifyError::BadKey)?;

        Ok(MlVerifyingKey::decode(&enc))

    }

}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]

pub enum Domain {

    Change,

    Attest,

    Policy,

    /// A server-issued challenge, proving possession of a key to a host

    /// (hub spec §5.3, §5.4). Never a stored object.

    Challenge,

}

impl Domain {

    pub fn name(self) -> &'static str {

        match self {

            Domain::Change => "change",

            Domain::Attest => "attest",

            Domain::Policy => "policy",

            Domain::Challenge => "challenge",

        }

    }

}

/// `"oxy1 sig " || domain || 0x00 || KeyId || target` (spec §14.3).

pub fn signing_message(domain: Domain, key: &KeyId, target: &ObjectRef) -> Vec<u8> {

    let mut m = Vec::with_capacity(80);

    m.extend_from_slice(b"oxy1 sig ");

    m.extend_from_slice(domain.name().as_bytes());

    m.push(0);

    m.extend_from_slice(&key.0);

    m.extend_from_slice(&target.0);

    m

}

#[derive(Clone, PartialEq, Eq)]

pub struct Signature {

    pub key: KeyId,

    pub ed: [u8; ED25519_SIG_LEN],

    pub ml: Vec<u8>,

}

impl std::fmt::Debug for Signature {

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        write!(f, "Signature(key {:?})", self.key)

    }

}

impl Signature {

    pub fn to_value(&self) -> Value {

        Value::map([

            (0, Value::Bytes(self.key.0.to_vec())),

            (1, Value::Uint(ALG_COMPOSITE)),

            (2, Value::Bytes(self.ed.to_vec())),

            (3, Value::Bytes(self.ml.clone())),

        ])

    }

    pub fn from_value(v: &Value) -> Result<Signature, ObjectError> {

        let m = v

            .as_map()

            .ok_or_else(|| ObjectError::Schema("signature: expected map".into()))?;

        let alg = uint(field(m, 1, "alg")?, "alg")?;

        if alg != ALG_COMPOSITE {

            return Err(ObjectError::Schema(format!(

                "unknown signature algorithm {alg}"

            )));

        }

        let ed = bytes(field(m, 2, "ed")?, "ed")?;

        let ml = bytes(field(m, 3, "ml")?, "ml")?;

        if ed.len() != ED25519_SIG_LEN || ml.len() != MLDSA65_SIG_LEN {

            return Err(ObjectError::Schema("signature: wrong length".into()));

        }

        Ok(Signature {

            key: oref(field(m, 0, "key")?, "key")?,

            ed: ed.try_into().unwrap(),

            ml: ml.to_vec(),

        })

    }

    pub fn encode_list(sigs: &[Signature]) -> Vec<u8> {

        cbor::encode(&Value::Array(sigs.iter().map(|s| s.to_value()).collect()))

    }

    pub fn decode_list(b: &[u8]) -> Result<Vec<Signature>, ObjectError> {

        let v = cbor::decode(b)?;

        array(&v, "signatures")?

            .iter()

            .map(Signature::from_value)

            .collect()

    }

}

#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]

pub enum VerifyError {

    #[error("signature names a different key than the identity")]

    KeyMismatch,

    #[error("identity holds a malformed public key")]

    BadKey,

    #[error("Ed25519 signature invalid")]

    Ed25519,

    #[error("ML-DSA-65 signature invalid")]

    MlDsa,

    #[error("a challenge is at least 16 bytes")]

    ShortChallenge,

}

/// Verify a composite signature by `identity` over `target` (spec §14.3).

pub fn verify(

    identity: &Identity,

    sig: &Signature,

    domain: Domain,

    target: &ObjectRef,

) -> Result<(), VerifyError> {

    let key = identity.hash();

    if sig.key != key {

        return Err(VerifyError::KeyMismatch);

    }

    let msg = signing_message(domain, &key, target);

    let ed_sig = EdSignature::from_bytes(&sig.ed);

    identity

        .ed_key()?

        .verify_strict(&msg, &ed_sig)

        .map_err(|_| VerifyError::Ed25519)?;

    let enc =

        EncodedSignature::<MlDsa65>::try_from(sig.ml.as_slice()).map_err(|_| VerifyError::MlDsa)?;

    let ml_sig = MlSignature::<MlDsa65>::decode(&enc).ok_or(VerifyError::MlDsa)?;

    if !identity.ml_key()?.verify_with_context(&msg, b"", &ml_sig) {

        return Err(VerifyError::MlDsa);

    }

    Ok(())

}

/// The seeds of a composite key pair. Stored only in the user's key directory.

#[derive(Clone)]

pub struct SecretKey {

    pub ed_seed: [u8; 32],

    pub ml_seed: [u8; 32],

}

impl std::fmt::Debug for SecretKey {

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        f.write_str("SecretKey(..)")

    }

}

impl SecretKey {

    pub fn generate() -> SecretKey {

        let mut ed_seed = [0u8; 32];

        let mut ml_seed = [0u8; 32];

        rand::rngs::OsRng.fill_bytes(&mut ed_seed);

        rand::rngs::OsRng.fill_bytes(&mut ml_seed);

        SecretKey { ed_seed, ml_seed }

    }

    fn ed(&self) -> EdSigningKey {

        EdSigningKey::from_bytes(&self.ed_seed)

    }

    fn ml(&self) -> MlSigningKey<MlDsa65> {

        MlSigningKey::from_seed(&Seed::from(self.ml_seed))

    }

    /// The identity object for these keys.

    pub fn identity(&self, name: &str, email: &str, created: u64) -> Identity {

        Identity {

            ed25519_pk: self.ed().verifying_key().to_bytes(),

            mldsa65_pk: self

                .ml()

                .expanded_key()

                .verifying_key()

                .encode()

                .as_slice()

                .to_vec(),

            name: name.into(),

            email: email.into(),

            created,

            expires: None,

            previous: None,

            delegates: vec![],

        }

    }

    pub fn sign(&self, domain: Domain, key: &KeyId, target: &ObjectRef) -> Signature {

        let msg = signing_message(domain, key, target);

        let ed = self.ed().sign(&msg).to_bytes();

        let ml = self

            .ml()

            .expanded_key()

            .sign_deterministic(&msg, b"")

            .expect("empty context is valid")

            .encode()

            .as_slice()

            .to_vec();

        Signature { key: *key, ed, ml }

    }

    pub fn encode(&self, key: &KeyId) -> Vec<u8> {

        cbor::encode(&Value::map([

            (0, Value::Uint(FORMAT_VERSION)),

            (1, Value::Bytes(key.0.to_vec())),

            (2, Value::Bytes(self.ed_seed.to_vec())),

            (3, Value::Bytes(self.ml_seed.to_vec())),

        ]))

    }

    pub fn decode(b: &[u8]) -> Result<(KeyId, SecretKey), ObjectError> {

        let v = cbor::decode(b)?;

        let m = v

            .as_map()

            .ok_or_else(|| ObjectError::Schema("secret key: expected map".into()))?;

        check_version(m)?;

        let seed = |k: u64| -> Result<[u8; 32], ObjectError> {

            bytes(field(m, k, "seed")?, "seed")?

                .try_into()

                .map_err(|_| ObjectError::Schema("seed: expected 32 bytes".into()))

        };

        Ok((

            oref(field(m, 1, "key")?, "key")?,

            SecretKey {

                ed_seed: seed(2)?,

                ml_seed: seed(3)?,

            },

        ))

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    fn fixed_key() -> SecretKey {

        SecretKey {

            ed_seed: [1; 32],

            ml_seed: [2; 32],

        }

    }

    #[test]

    fn identity_round_trip_and_sizes() {

        let sk = fixed_key();

        let id = sk.identity("Ada", "[email protected]", 1_700_000_000_000);

        assert_eq!(id.mldsa65_pk.len(), MLDSA65_PK_LEN);

        let d = Identity::decode(&id.encode()).unwrap();

        assert_eq!(d, id);

        assert_eq!(d.hash(), id.hash());

        // Deterministic from seeds.

        assert_eq!(

            fixed_key()

                .identity("Ada", "[email protected]", 1_700_000_000_000)

                .hash(),

            id.hash()

        );

        let mut other = id.clone();

        other.delegates.push(ObjectRef([9; 32]));

        assert_ne!(other.hash(), id.hash());

    }

    #[test]

    fn sign_and_verify() {

        let sk = fixed_key();

        let id = sk.identity("Ada", "[email protected]", 1);

        let key = id.hash();

        let target = ObjectRef([7; 32]);

        let sig = sk.sign(Domain::Change, &key, &target);

        assert_eq!(sig.ml.len(), MLDSA65_SIG_LEN);

        assert_eq!(verify(&id, &sig, Domain::Change, &target), Ok(()));

        // Different target, different domain: invalid.

        assert_eq!(

            verify(&id, &sig, Domain::Change, &ObjectRef([8; 32])),

            Err(VerifyError::Ed25519)

        );

        assert_eq!(

            verify(&id, &sig, Domain::Attest, &target),

            Err(VerifyError::Ed25519)

        );

        // Deterministic signatures.

        assert_eq!(sk.sign(Domain::Change, &key, &target), sig);

        // Signature list round trip.

        let l =

            Signature::decode_list(&Signature::encode_list(std::slice::from_ref(&sig))).unwrap();

        assert_eq!(l, vec![sig]);

    }

    #[test]

    fn stripping_or_forging_either_half_fails() {

        let sk = fixed_key();

        let id = sk.identity("Ada", "a@b", 1);

        let key = id.hash();

        let target = ObjectRef([7; 32]);

        let good = sk.sign(Domain::Change, &key, &target);

        let mut no_ml = good.clone();

        no_ml.ml = vec![0; MLDSA65_SIG_LEN];

        assert_eq!(

            verify(&id, &no_ml, Domain::Change, &target),

            Err(VerifyError::MlDsa)

        );

        let mut bad_ed = good.clone();

        bad_ed.ed[10] ^= 1;

        assert_eq!(

            verify(&id, &bad_ed, Domain::Change, &target),

            Err(VerifyError::Ed25519)

        );

        // A signature made for another KeyId over the same target does not transfer.

        let other = SecretKey {

            ed_seed: [3; 32],

            ml_seed: [4; 32],

        };

        let other_id = other.identity("Bob", "b@b", 1);

        let theirs = other.sign(Domain::Change, &other_id.hash(), &target);

        assert_eq!(

            verify(&id, &theirs, Domain::Change, &target),

            Err(VerifyError::KeyMismatch)

        );

        let mut relabelled = theirs.clone();

        relabelled.key = key;

        assert_eq!(

            verify(&id, &relabelled, Domain::Change, &target),

            Err(VerifyError::Ed25519)