// SSH加密模块（Phase 3：密钥交换）
// 参考OpenSSH curve25519.c, kex.c

use anyhow::{Result, anyhow};
use x25519_dalek::{EphemeralSecret, PublicKey, SharedSecret};
use ed25519_dalek::{SigningKey, VerifyingKey, Signature, Signer};
use sha2::{Sha256, Digest};
use log::{info, debug};
use rand::rngs::OsRng;

/// Curve25519密钥交换处理器（参考OpenSSH curve25519.c）
pub struct Curve25519Kex {
    secret: Option<EphemeralSecret>,  // 使用Option包装（一次性使用类型）
    public: PublicKey,
}

impl Curve25519Kex {
    /// 创建新的Curve25519密钥交换实例
    pub fn new() -> Self {
        // 参考OpenSSH curve25519.c: curve25519_make_key()
        // x25519-dalek 2.0标准API：使用random_from_rng
        let secret = EphemeralSecret::random_from_rng(OsRng);
        let public = PublicKey::from(&secret);
        
        Self { secret: Some(secret), public }  // Some包装
    }
    
    /// 获取公钥（用于SSH_MSG_KEX_ECDH_INIT）
    pub fn public_key(&self) -> &[u8] {
        self.public.as_bytes()
    }
    
    /// 计算共享密钥（参考OpenSSH curve25519_shared_secret()）
    /// 使用&mut self（消耗模式，符合OpenSSH设计）
    pub fn compute_shared_secret(&mut self, client_public: &[u8]) -> Result<[u8; 32]> {
        if client_public.len() != 32 {
            return Err(anyhow!("Invalid client public key length"));
        }
        
        // 参考OpenSSH：curve25519共享密钥计算
        let client_public = PublicKey::from(<[u8; 32]>::try_from(client_public)?);
        
        // 使用take()取出secret（Rust标准模式）
        if let Some(secret) = self.secret.take() {
            let shared_secret = secret.diffie_hellman(&client_public);
            Ok(shared_secret.as_bytes().clone())
        } else {
            Err(anyhow!("Secret already used"))
        }
    }
}

/// SSH会话密钥计算（参考OpenSSH kex.c: derive_keys()）
pub struct SessionKeys {
    pub session_id: Vec<u8>,
    pub encryption_key_ctos: Vec<u8>,
    pub encryption_key_stoc: Vec<u8>,
    pub mac_key_ctos: Vec<u8>,
    pub mac_key_stoc: Vec<u8>,
}

impl SessionKeys {
    /// 计算会话密钥（参考OpenSSH kex.c: kex_derive_keys()）
    pub fn derive(
        shared_secret: &[u8],
        hash_algo: &str,
        server_public_key: &[u8],
        client_public_key: &[u8],
        server_host_key: &[u8],
    ) -> Result<Self> {
        // 参考OpenSSH：SHA256 hash计算
        // Hash = SHA256(共享密钥 + 其他数据)
        
        // 会话ID计算（参考OpenSSH kex.c）
        let mut hasher = Sha256::new();
        hasher.update(shared_secret);
        hasher.update(server_public_key);
        hasher.update(client_public_key);
        hasher.update(server_host_key);
        let hash = hasher.finalize();
        
        let session_id = hash.to_vec();
        
        // 加密密钥计算（简化实现，参考OpenSSH）
        let encryption_key_ctos = Self::derive_key(&session_id, shared_secret, 'A')?;
        let encryption_key_stoc = Self::derive_key(&session_id, shared_secret, 'B')?;
        let mac_key_ctos = Self::derive_key(&session_id, shared_secret, 'C')?;
        let mac_key_stoc = Self::derive_key(&session_id, shared_secret, 'D')?;
        
        Ok(Self {
            session_id,
            encryption_key_ctos,
            encryption_key_stoc,
            mac_key_ctos,
            mac_key_stoc,
        })
    }
    
    /// 密钥派生函数（参考OpenSSH kex.c: kex_derive_key()）
    fn derive_key(session_id: &[u8], shared_secret: &[u8], char: char) -> Result<Vec<u8>> {
        // OpenSSH key derivation: KDF(session_id, shared_secret, char)
        // 简化实现：SHA256(session_id + shared_secret + char)
        
        let mut hasher = Sha256::new();
        hasher.update(session_id);
        hasher.update(shared_secret);
        hasher.update(&[char as u8]);
        
        Ok(hasher.finalize().to_vec())
    }
}

/// Ed25519服务器主机密钥（参考OpenSSH sshkey.c）
pub struct Ed25519HostKey {
    signing_key: SigningKey,
}

impl Ed25519HostKey {
    /// 加载或生成主机密钥（参考OpenSSH hostfile.c）
    pub fn load_or_generate(key_path: &str) -> Result<Self> {
        // 简化实现：生成临时密钥（实际应从文件加载）
        // 参考OpenSSH ssh-keygen
        
        let signing_key = SigningKey::generate(&mut OsRng);
        
        Ok(Self { signing_key })
    }
    
    /// 获取公钥（用于SSH_MSG_KEX_ECDH_REPLY）
    pub fn public_key_bytes(&self) -> Vec<u8> {
        // SSH Ed25519公钥格式（参考OpenSSH sshkey.c）
        let verifying_key = self.signing_key.verifying_key();
        
        // SSH格式：ssh-ed25519 + 公钥bytes
        // 简化：仅返回公钥bytes（32字节）
        verifying_key.as_bytes().to_vec()
    }
    
    /// 签名（参考OpenSSH sshkey.c: sshkey_sign()）
    pub fn sign(&self, data: &[u8]) -> Result<Vec<u8>> {
        // OpenSSH Ed25519签名
        let signature = self.signing_key.sign(data);
        
        // SSH签名格式（参考OpenSSH ssh-sign.c）
        // 简化：仅返回签名bytes（64字节）
        Ok(signature.to_bytes().to_vec())
    }
    
    /// 获取完整SSH公钥格式（参考OpenSSH sshkey.c）
    pub fn ssh_public_key(&self) -> String {
        let public_bytes = self.public_key_bytes();
        
        // SSH公钥格式：ssh-ed25519 <base64-encoded-public-key>
        // 参考OpenSSH ssh-keygen -y
        
        use base64::{Engine as _, engine::general_purpose};
        let encoded = general_purpose::STANDARD.encode(&public_bytes);
        
        format!("ssh-ed25519 {}", encoded)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_curve25519_key_generation() {
        let kex = Curve25519Kex::new();
        assert_eq!(kex.public_key().len(), 32);
    }
    
    #[test]
    fn test_curve25519_shared_secret() {
        let client_kex = Curve25519Kex::new();
        let server_kex = Curve25519Kex::new();
        
        // 客户端计算共享密钥
        let client_secret = client_kex.compute_shared_secret(server_kex.public_key()).unwrap();
        
        // 服务器计算共享密钥
        let server_secret = server_kex.compute_shared_secret(client_kex.public_key()).unwrap();
        
        // 应该相同（Curve25519特性）
        assert_eq!(client_secret, server_secret);
    }
    
    #[test]
    fn test_ed25519_host_key() {
        let host_key = Ed25519HostKey::load_or_generate("test_key").unwrap();
        assert_eq!(host_key.public_key_bytes().len(), 32);
    }
    
    #[test]
    fn test_ed25519_signature() {
        let host_key = Ed25519HostKey::load_or_generate("test_key").unwrap();
        let data = b"test data";
        
        let signature = host_key.sign(data).unwrap();
        assert_eq!(signature.len(), 64);  // Ed25519签名64字节
    }
}
