qssh 0.4.3

//! Certificate-Based Authentication for QSSH
//!
//! Implements SSH certificate support for scalable enterprise authentication

use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH, Duration};
use chrono::{DateTime, Utc};
use serde::{Serialize, Deserialize};
use crate::{Result, QsshError, PqAlgorithm};
use sha2::{Sha256, Digest};
use fn_dsa::{KeyPairGenerator as _, SigningKey as FnSigningKeyTrait, VerifyingKey as FnVerifyingKeyTrait};
use signature::Keypair as _;

/// SSH Certificate Type
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CertificateType {
    /// User certificate for client authentication
    User = 1,
    /// Host certificate for server authentication
    Host = 2,
}

/// SSH Certificate
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SshCertificate {
    /// Certificate type
    pub cert_type: CertificateType,
    /// Serial number
    pub serial: u64,
    /// Certificate key ID
    pub key_id: String,
    /// Valid principals (users or hostnames)
    pub valid_principals: Vec<String>,
    /// Valid from timestamp
    pub valid_after: u64,
    /// Valid until timestamp
    pub valid_before: u64,
    /// Critical options
    pub critical_options: HashMap<String, String>,
    /// Extensions
    pub extensions: HashMap<String, String>,
    /// Public key being certified
    pub public_key: PublicKey,
    /// Signature key (CA key)
    pub signature_key: PublicKey,
    /// Certificate signature
    pub signature: Vec<u8>,
    /// Nonce for replay protection
    pub nonce: Vec<u8>,
}

/// Public key in certificate
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PublicKey {
    /// Algorithm type
    pub algorithm: PqAlgorithm,
    /// Key data
    pub key_data: Vec<u8>,
    /// Key fingerprint
    pub fingerprint: String,
}

/// Certificate Authority
pub struct CertificateAuthority {
    /// CA private key
    private_key: Vec<u8>,
    /// CA public key
    public_key: PublicKey,
    /// CA key algorithm
    algorithm: PqAlgorithm,
    /// Serial number counter
    serial_counter: u64,
    /// Default certificate lifetime
    default_lifetime: Duration,
}

impl CertificateAuthority {
    /// Create new CA
    pub fn new(algorithm: PqAlgorithm) -> Result<Self> {
        let (private_key, public_key_data) = generate_ca_keypair(algorithm)?;

        let public_key = PublicKey {
            algorithm,
            key_data: public_key_data.clone(),
            fingerprint: generate_fingerprint(&public_key_data),
        };

        Ok(Self {
            private_key,
            public_key,
            algorithm,
            serial_counter: 1,
            default_lifetime: Duration::from_secs(365 * 24 * 3600), // 1 year
        })
    }

    /// Load CA from files
    pub fn load(private_key_path: &str, public_key_path: &str) -> Result<Self> {
        let private_key = std::fs::read(private_key_path)
            .map_err(QsshError::Io)?;

        let public_key_data = std::fs::read(public_key_path)
            .map_err(QsshError::Io)?;

        // Detect algorithm from key size
        let algorithm = detect_algorithm(&public_key_data)?;

        let public_key = PublicKey {
            algorithm,
            key_data: public_key_data.clone(),
            fingerprint: generate_fingerprint(&public_key_data),
        };

        Ok(Self {
            private_key,
            public_key,
            algorithm,
            serial_counter: 1,
            default_lifetime: Duration::from_secs(365 * 24 * 3600),
        })
    }

    /// Issue user certificate
    pub fn issue_user_certificate(
        &mut self,
        user_public_key: Vec<u8>,
        key_id: String,
        principals: Vec<String>,
        options: CertificateOptions,
    ) -> Result<SshCertificate> {
        self.issue_certificate(
            CertificateType::User,
            user_public_key,
            key_id,
            principals,
            options,
        )
    }

    /// Issue host certificate
    pub fn issue_host_certificate(
        &mut self,
        host_public_key: Vec<u8>,
        key_id: String,
        hostnames: Vec<String>,
        options: CertificateOptions,
    ) -> Result<SshCertificate> {
        self.issue_certificate(
            CertificateType::Host,
            host_public_key,
            key_id,
            hostnames,
            options,
        )
    }

    /// Issue a certificate
    fn issue_certificate(
        &mut self,
        cert_type: CertificateType,
        public_key_data: Vec<u8>,
        key_id: String,
        principals: Vec<String>,
        options: CertificateOptions,
    ) -> Result<SshCertificate> {
        let serial = self.next_serial();
        let now = current_timestamp();

        let valid_after = options.valid_after.unwrap_or(now);
        let valid_before = options.valid_before
            .unwrap_or(now + self.default_lifetime.as_secs());

        let public_key = PublicKey {
            algorithm: self.algorithm,
            key_data: public_key_data.clone(),
            fingerprint: generate_fingerprint(&public_key_data),
        };

        // Set default extensions based on certificate type
        let mut extensions = options.extensions.unwrap_or_default();
        if cert_type == CertificateType::User {
            apply_default_user_extensions(&mut extensions);
        }

        let mut cert = SshCertificate {
            cert_type,
            serial,
            key_id,
            valid_principals: principals,
            valid_after,
            valid_before,
            critical_options: options.critical_options.unwrap_or_default(),
            extensions,
            public_key,
            signature_key: self.public_key.clone(),
            signature: Vec::new(),
            nonce: generate_nonce(),
        };

        // Sign the certificate
        cert.signature = self.sign_certificate(&cert)?;

        Ok(cert)
    }

    /// Sign certificate data
    fn sign_certificate(&self, cert: &SshCertificate) -> Result<Vec<u8>> {
        let cert_data = serialize_for_signing(cert)?;

        match self.algorithm {
            PqAlgorithm::Falcon512 => {
                sign_falcon(&cert_data, &self.private_key)
            },
            PqAlgorithm::SphincsPlus => {
                sign_sphincs(&cert_data, &self.private_key)
            },
            _ => Err(QsshError::Protocol("Unsupported CA algorithm".into())),
        }
    }

    /// Get next serial number
    fn next_serial(&mut self) -> u64 {
        let serial = self.serial_counter;
        self.serial_counter += 1;
        serial
    }

    /// Get CA public key
    pub fn public_key(&self) -> &PublicKey {
        &self.public_key
    }
}

/// Certificate options
#[derive(Debug, Clone, Default)]
pub struct CertificateOptions {
    /// Valid from timestamp
    pub valid_after: Option<u64>,
    /// Valid until timestamp
    pub valid_before: Option<u64>,
    /// Critical options
    pub critical_options: Option<HashMap<String, String>>,
    /// Extensions
    pub extensions: Option<HashMap<String, String>>,
}

/// Certificate validator
pub struct CertificateValidator {
    /// Trusted CA keys
    trusted_cas: Vec<PublicKey>,
    /// Revoked certificates
    revoked_serials: Vec<u64>,
    /// Allow expired certificates (for testing)
    allow_expired: bool,
}

impl Default for CertificateValidator {
    fn default() -> Self {
        Self::new()
    }
}

impl CertificateValidator {
    /// Create new validator
    pub fn new() -> Self {
        Self {
            trusted_cas: Vec::new(),
            revoked_serials: Vec::new(),
            allow_expired: false,
        }
    }

    /// Add trusted CA
    pub fn add_trusted_ca(&mut self, ca_public_key: PublicKey) {
        self.trusted_cas.push(ca_public_key);
    }

    /// Add revoked certificate
    pub fn revoke_certificate(&mut self, serial: u64) {
        self.revoked_serials.push(serial);
    }

    /// Validate certificate
    pub fn validate(&self, cert: &SshCertificate) -> Result<ValidationResult> {
        // Check if certificate is revoked
        if self.revoked_serials.contains(&cert.serial) {
            return Ok(ValidationResult::Revoked);
        }

        // Check validity period
        let now = current_timestamp();
        if !self.allow_expired {
            if now < cert.valid_after {
                return Ok(ValidationResult::NotYetValid);
            }
            if now > cert.valid_before {
                return Ok(ValidationResult::Expired);
            }
        }

        // Check if CA is trusted
        let ca_trusted = self.trusted_cas.iter()
            .any(|ca| ca.fingerprint == cert.signature_key.fingerprint);

        if !ca_trusted {
            return Ok(ValidationResult::UntrustedCA);
        }

        // Verify signature
        let cert_data = serialize_for_signing(cert)?;
        let valid = match cert.signature_key.algorithm {
            PqAlgorithm::Falcon512 => {
                verify_falcon(&cert_data, &cert.signature, &cert.signature_key.key_data)?
            },
            PqAlgorithm::SphincsPlus => {
                verify_sphincs(&cert_data, &cert.signature, &cert.signature_key.key_data)?
            },
            _ => false,
        };

        if valid {
            Ok(ValidationResult::Valid)
        } else {
            Ok(ValidationResult::InvalidSignature)
        }
    }

    /// Check if principal is authorized
    pub fn check_principal(&self, cert: &SshCertificate, principal: &str) -> bool {
        cert.valid_principals.is_empty() ||
        cert.valid_principals.contains(&principal.to_string())
    }

    /// Check critical options
    pub fn check_critical_options(&self, cert: &SshCertificate, required: &HashMap<String, String>) -> bool {
        for (key, value) in required {
            match cert.critical_options.get(key) {
                Some(v) if v == value => continue,
                _ => return false,
            }
        }
        true
    }
}

/// Certificate validation result
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ValidationResult {
    /// Certificate is valid
    Valid,
    /// Certificate has been revoked
    Revoked,
    /// Certificate is not yet valid
    NotYetValid,
    /// Certificate has expired
    Expired,
    /// CA is not trusted
    UntrustedCA,
    /// Signature verification failed
    InvalidSignature,
}

/// Certificate store for managing certificates
pub struct CertificateStore {
    /// User certificates by key ID
    user_certs: HashMap<String, SshCertificate>,
    /// Host certificates by hostname
    host_certs: HashMap<String, SshCertificate>,
    /// CA certificates
    ca_certs: Vec<PublicKey>,
}

impl Default for CertificateStore {
    fn default() -> Self {
        Self::new()
    }
}

impl CertificateStore {
    /// Create new certificate store
    pub fn new() -> Self {
        Self {
            user_certs: HashMap::new(),
            host_certs: HashMap::new(),
            ca_certs: Vec::new(),
        }
    }

    /// Add user certificate
    pub fn add_user_cert(&mut self, cert: SshCertificate) {
        self.user_certs.insert(cert.key_id.clone(), cert);
    }

    /// Add host certificate
    pub fn add_host_cert(&mut self, hostname: String, cert: SshCertificate) {
        self.host_certs.insert(hostname, cert);
    }

    /// Add CA certificate
    pub fn add_ca_cert(&mut self, ca_key: PublicKey) {
        self.ca_certs.push(ca_key);
    }

    /// Get user certificate
    pub fn get_user_cert(&self, key_id: &str) -> Option<&SshCertificate> {
        self.user_certs.get(key_id)
    }

    /// Get host certificate
    pub fn get_host_cert(&self, hostname: &str) -> Option<&SshCertificate> {
        self.host_certs.get(hostname)
    }

    /// List all certificates
    pub fn list_all(&self) -> Vec<CertificateInfo> {
        let mut info = Vec::new();

        for cert in self.user_certs.values() {
            info.push(CertificateInfo::from_cert(cert));
        }

        for cert in self.host_certs.values() {
            info.push(CertificateInfo::from_cert(cert));
        }

        info
    }
}

/// Certificate information
#[derive(Debug, Clone)]
pub struct CertificateInfo {
    pub cert_type: CertificateType,
    pub key_id: String,
    pub serial: u64,
    pub valid_from: DateTime<Utc>,
    pub valid_until: DateTime<Utc>,
    pub principals: Vec<String>,
    pub ca_fingerprint: String,
}

impl CertificateInfo {
    fn from_cert(cert: &SshCertificate) -> Self {
        Self {
            cert_type: cert.cert_type,
            key_id: cert.key_id.clone(),
            serial: cert.serial,
            valid_from: DateTime::from_timestamp(cert.valid_after as i64, 0).unwrap_or_default(),
            valid_until: DateTime::from_timestamp(cert.valid_before as i64, 0).unwrap_or_default(),
            principals: cert.valid_principals.clone(),
            ca_fingerprint: cert.signature_key.fingerprint.clone(),
        }
    }
}

// Helper functions

fn generate_ca_keypair(algorithm: PqAlgorithm) -> Result<(Vec<u8>, Vec<u8>)> {
    match algorithm {
        PqAlgorithm::Falcon512 => {
            let mut sk = vec![0u8; fn_dsa::sign_key_size(fn_dsa::FN_DSA_LOGN_512)];
            let mut pk = vec![0u8; fn_dsa::vrfy_key_size(fn_dsa::FN_DSA_LOGN_512)];
            fn_dsa::KeyPairGeneratorStandard::default()
                .keygen(fn_dsa::FN_DSA_LOGN_512, &mut aes_gcm::aead::OsRng, &mut sk, &mut pk);
            Ok((sk, pk))
        },
        PqAlgorithm::SphincsPlus => {
            let sk = slh_dsa::SigningKey::<slh_dsa::Sha2_128s>::new(&mut aes_gcm::aead::OsRng);
            let pk = sk.verifying_key().clone();
            Ok((sk.to_bytes().to_vec(), pk.to_bytes().to_vec()))
        },
        _ => Err(QsshError::Protocol("Unsupported algorithm for CA".into())),
    }
}

fn detect_algorithm(public_key: &[u8]) -> Result<PqAlgorithm> {
    // Detect based on key size
    match public_key.len() {
        897 => Ok(PqAlgorithm::Falcon512),
        32 => Ok(PqAlgorithm::SphincsPlus),
        _ => Ok(PqAlgorithm::Falcon512), // Default
    }
}

fn generate_fingerprint(data: &[u8]) -> String {
    let mut hasher = Sha256::new();
    hasher.update(data);
    let hash = hasher.finalize();
    format!("SHA256:{}", base64::Engine::encode(&base64::engine::general_purpose::STANDARD, hash))
}

fn generate_nonce() -> Vec<u8> {
    use rand::Rng;
    let mut rng = rand::thread_rng();
    (0..32).map(|_| rng.gen()).collect()
}

fn current_timestamp() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs()
}

fn serialize_for_signing(cert: &SshCertificate) -> Result<Vec<u8>> {
    // Create a copy without signature for signing
    let mut cert_copy = cert.clone();
    cert_copy.signature.clear();

    bincode::serialize(&cert_copy)
        .map_err(|e| QsshError::Protocol(format!("Failed to serialize certificate: {}", e)))
}

fn sign_falcon(data: &[u8], private_key: &[u8]) -> Result<Vec<u8>> {
    let mut sk = fn_dsa::SigningKeyStandard::decode(private_key)
        .ok_or_else(|| QsshError::Crypto("Invalid Falcon private key".into()))?;
    let mut sig = vec![0u8; fn_dsa::signature_size(fn_dsa::FN_DSA_LOGN_512)];
    sk.sign(&mut aes_gcm::aead::OsRng, &fn_dsa::DOMAIN_NONE, &fn_dsa::HASH_ID_RAW, data, &mut sig);
    Ok(sig)
}

fn verify_falcon(data: &[u8], signature: &[u8], public_key: &[u8]) -> Result<bool> {
    let vk = fn_dsa::VerifyingKeyStandard::decode(public_key)
        .ok_or_else(|| QsshError::Crypto("Invalid Falcon public key".into()))?;
    Ok(vk.verify(signature, &fn_dsa::DOMAIN_NONE, &fn_dsa::HASH_ID_RAW, data))
}

fn sign_sphincs(data: &[u8], private_key: &[u8]) -> Result<Vec<u8>> {
    use signature::Signer;

    let sk = slh_dsa::SigningKey::<slh_dsa::Sha2_128s>::try_from(private_key)
        .map_err(|_| QsshError::Crypto("Invalid SPHINCS+ private key".into()))?;
    let sig: slh_dsa::Signature<slh_dsa::Sha2_128s> = Signer::try_sign(&sk, data)
        .map_err(|e| QsshError::Crypto(format!("SPHINCS+ signing failed: {}", e)))?;
    Ok(sig.to_bytes().to_vec())
}

fn verify_sphincs(data: &[u8], signature: &[u8], public_key: &[u8]) -> Result<bool> {
    use signature::Verifier;

    let pk = slh_dsa::VerifyingKey::<slh_dsa::Sha2_128s>::try_from(public_key)
        .map_err(|_| QsshError::Crypto("Invalid SPHINCS+ public key".into()))?;
    let sig = slh_dsa::Signature::<slh_dsa::Sha2_128s>::try_from(signature)
        .map_err(|_| QsshError::Crypto("Invalid SPHINCS+ signature".into()))?;
    Ok(Verifier::verify(&pk, data, &sig).is_ok())
}

fn apply_default_user_extensions(extensions: &mut HashMap<String, String>) {
    // Add default SSH user certificate extensions if not present
    extensions.entry("permit-X11-forwarding".to_string())
        .or_insert("".to_string());
    extensions.entry("permit-agent-forwarding".to_string())
        .or_insert("".to_string());
    extensions.entry("permit-port-forwarding".to_string())
        .or_insert("".to_string());
    extensions.entry("permit-pty".to_string())
        .or_insert("".to_string());
    extensions.entry("permit-user-rc".to_string())
        .or_insert("".to_string());
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_ca_creation() {
        let ca = CertificateAuthority::new(PqAlgorithm::Falcon512).unwrap();
        assert_eq!(ca.algorithm, PqAlgorithm::Falcon512);
        assert!(!ca.public_key.fingerprint.is_empty());
    }

    #[test]
    fn test_certificate_issuance() {
        let mut ca = CertificateAuthority::new(PqAlgorithm::Falcon512).unwrap();

        let user_key = vec![1, 2, 3, 4, 5];
        let cert = ca.issue_user_certificate(
            user_key,
            "user@example.com".to_string(),
            vec!["alice".to_string(), "bob".to_string()],
            CertificateOptions::default(),
        ).unwrap();

        assert_eq!(cert.cert_type, CertificateType::User);
        assert_eq!(cert.key_id, "user@example.com");
        assert_eq!(cert.valid_principals, vec!["alice", "bob"]);
        assert!(!cert.signature.is_empty());
    }

    #[test]
    fn test_certificate_validation() {
        let mut ca = CertificateAuthority::new(PqAlgorithm::Falcon512).unwrap();
        let mut validator = CertificateValidator::new();
        validator.add_trusted_ca(ca.public_key().clone());

        let user_key = vec![5, 4, 3, 2, 1];
        let cert = ca.issue_user_certificate(
            user_key,
            "test@example.com".to_string(),
            vec!["testuser".to_string()],
            CertificateOptions::default(),
        ).unwrap();

        let result = validator.validate(&cert).unwrap();
        assert_eq!(result, ValidationResult::Valid);

        // Test principal checking
        assert!(validator.check_principal(&cert, "testuser"));
        assert!(!validator.check_principal(&cert, "otheruser"));
    }

    #[test]
    fn test_certificate_expiration() {
        let mut ca = CertificateAuthority::new(PqAlgorithm::Falcon512).unwrap();
        let validator = CertificateValidator::new();

        let past = current_timestamp() - 3600;
        let options = CertificateOptions {
            valid_after: Some(past - 7200),
            valid_before: Some(past),
            ..Default::default()
        };

        let cert = ca.issue_user_certificate(
            vec![1, 2, 3],
            "expired@example.com".to_string(),
            vec!["user".to_string()],
            options,
        ).unwrap();

        let result = validator.validate(&cert).unwrap();
        assert_eq!(result, ValidationResult::Expired);
    }
}