aria2-core 0.2.3

//! MSE (Message Stream Encryption) encrypted handshake module
//!
//! Implements BitTorrent BEP 10 encryption layer protocol, overlaying optional encrypted handshake on BT protocol handshake.
//! Contains three-phase handshake: Method Selection, PAD/DH key exchange, SKEY/SVC verification.

use rc4::{KeyInit, Rc4 as Rc4Cipher, StreamCipher};
use ring::agreement::{self, EphemeralPrivateKey, UnparsedPublicKey};
use ring::rand::SystemRandom;
use sha1::{Digest, Sha1};
use std::sync::Mutex;

use crate::error::{Aria2Error, Result};

/// MSE encryption method
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CryptoMethod {
    /// Plaintext transmission, no encryption
    Plain = 0x0001,
    /// RC4 stream cipher
    Rc4 = 0x0002,
    /// AES-128-CBC (optional)
    Aes128Cbc = 0x0003,
}

impl CryptoMethod {
    /// Create CryptoMethod from u16 value
    pub fn from_u16(value: u16) -> Option<Self> {
        match value {
            0x0001 => Some(CryptoMethod::Plain),
            0x0002 => Some(CryptoMethod::Rc4),
            0x0003 => Some(CryptoMethod::Aes128Cbc),
            _ => None,
        }
    }

    /// Convert to u16
    pub fn to_u16(self) -> u16 {
        self as u16
    }
}

/// MSE handshake state machine
#[derive(Debug, Clone, PartialEq)]
#[allow(clippy::large_enum_variant)]
pub enum MseState {
    /// Idle state
    Idle,
    /// Method selection sent, waiting for remote response
    MethodSelectionSent,
    /// DH parameter exchange in progress
    KeyExchangeInProgress,
    /// Waiting for SKEY verification
    VerificationPending,
    /// Handshake complete
    Established(MseCryptoContext),
    /// Handshake failed
    Failed(String),
}

/// MSE encryption context (used for encrypting/decrypting BT messages after handshake complete)
pub struct MseCryptoContext {
    send_key: Vec<u8>,
    recv_key: Vec<u8>,
    crypto_method: CryptoMethod,
    rc4_send: Option<Rc4Cipher>,
    rc4_recv: Option<Rc4Cipher>,
}

impl std::fmt::Debug for MseCryptoContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MseCryptoContext")
            .field("send_key", &self.send_key)
            .field("recv_key", &self.recv_key)
            .field("crypto_method", &self.crypto_method)
            .field("rc4_send", &self.rc4_send.is_some())
            .field("rc4_recv", &self.rc4_recv.is_some())
            .finish()
    }
}

impl Clone for MseCryptoContext {
    fn clone(&self) -> Self {
        // Note: RC4 state cannot truly be cloned (will lose state), here creating new instance
        match self.crypto_method {
            CryptoMethod::Rc4 => {
                // Attempt to sync state (imperfect but usable)
                Self::new(&self.send_key, &self.recv_key, self.crypto_method)
            }
            _ => Self {
                send_key: self.send_key.clone(),
                recv_key: self.recv_key.clone(),
                crypto_method: self.crypto_method,
                rc4_send: None,
                rc4_recv: None,
            },
        }
    }
}

impl PartialEq for MseCryptoContext {
    fn eq(&self, other: &Self) -> bool {
        self.send_key == other.send_key
            && self.recv_key == other.recv_key
            && self.crypto_method == other.crypto_method
    }
}

impl MseCryptoContext {
    /// Create new encryption context using derived keys
    ///
    /// For RC4 method, initializes two RC4 instances and discards first 1024 bytes keystream
    /// (MSE spec section 5.2 to prevent keystream attack)
    pub fn new(send_key: &[u8], recv_key: &[u8], method: CryptoMethod) -> Self {
        match method {
            CryptoMethod::Rc4 => {
                // Initialize send direction RC4 and discard 1024 bytes keystream
                let mut rc4_send = Rc4Cipher::new_from_slice(send_key).unwrap();
                let mut discard = vec![0u8; 1024];
                rc4_send.apply_keystream(&mut discard);

                // Initialize receive direction RC4 and discard 1024 bytes keystream
                let mut rc4_recv = Rc4Cipher::new_from_slice(recv_key).unwrap();
                let mut discard = vec![0u8; 1024];
                rc4_recv.apply_keystream(&mut discard);

                Self {
                    send_key: send_key.to_vec(),
                    recv_key: recv_key.to_vec(),
                    crypto_method: method,
                    rc4_send: Some(rc4_send),
                    rc4_recv: Some(rc4_recv),
                }
            }
            _ => Self {
                send_key: send_key.to_vec(),
                recv_key: recv_key.to_vec(),
                crypto_method: method,
                rc4_send: None,
                rc4_recv: None,
            },
        }
    }

    /// Encrypt data
    pub fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>> {
        match self.crypto_method {
            CryptoMethod::Rc4 => {
                if let Some(ref mut rc4) = self.rc4_send {
                    let mut data = plaintext.to_vec();
                    rc4.apply_keystream(&mut data);
                    Ok(data)
                } else {
                    Err(Aria2Error::Fatal(crate::error::FatalError::Config(
                        "RC4 cipher not initialized".to_string(),
                    )))
                }
            }
            CryptoMethod::Plain | CryptoMethod::Aes128Cbc => {
                // Plain and AES modes currently not supported or return plaintext directly
                Ok(plaintext.to_vec())
            }
        }
    }

    /// Decrypt data
    pub fn decrypt(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>> {
        match self.crypto_method {
            CryptoMethod::Rc4 => {
                if let Some(ref mut rc4) = self.rc4_recv {
                    let mut data = ciphertext.to_vec();
                    rc4.apply_keystream(&mut data);
                    Ok(data)
                } else {
                    Err(Aria2Error::Fatal(crate::error::FatalError::Config(
                        "RC4 cipher not initialized".to_string(),
                    )))
                }
            }
            CryptoMethod::Plain | CryptoMethod::Aes128Cbc => {
                // Plain and AES modes currently not supported or return plaintext directly
                Ok(ciphertext.to_vec())
            }
        }
    }

    /// Get current encryption method
    pub fn crypto_method(&self) -> CryptoMethod {
        self.crypto_method
    }

    /// Whether using encryption
    pub fn is_encrypted(&self) -> bool {
        self.crypto_method != CryptoMethod::Plain
    }
}

/// Plaintext fallback: create unencrypted context
impl Default for MseCryptoContext {
    fn default() -> Self {
        Self {
            send_key: vec![],
            recv_key: vec![],
            crypto_method: CryptoMethod::Plain,
            rc4_send: None,
            rc4_recv: None,
        }
    }
}

/// MSE handshake manager
pub struct MseHandshakeManager {
    /// Handshake state machine.
    /// Uses std::sync::Mutex because the lock is only held for short synchronous
    /// reads/writes and never across .await points.
    state: Mutex<MseState>,
    local_dh_private_key: Option<EphemeralPrivateKey>,
    local_dh_pubkey: Vec<u8>,
    remote_dh_pubkey: Option<Vec<u8>>,
    pub(crate) shared_secret: Option<Vec<u8>>,
    info_hash: [u8; 20],
    pad_length: u16,
}

impl MseHandshakeManager {
    /// Create new MSE handshake instance
    pub fn new(info_hash: [u8; 20]) -> Result<Self> {
        let rng = SystemRandom::new();

        // Generate X25519 key pair
        let private_key = EphemeralPrivateKey::generate(&agreement::X25519, &rng).map_err(|e| {
            Aria2Error::Fatal(crate::error::FatalError::Config(format!(
                "Failed to generate DH key: {}",
                e
            )))
        })?;

        // Get public key (X25519 public key is fixed at 32 bytes)
        let public_key = private_key.compute_public_key().map_err(|e| {
            Aria2Error::Fatal(crate::error::FatalError::Config(format!(
                "Failed to compute public key: {}",
                e
            )))
        })?;

        let pubkey_slice = public_key.as_ref();
        let mut local_dh_pubkey = vec![0u8; pubkey_slice.len()];
        local_dh_pubkey.copy_from_slice(pubkey_slice);

        // Random PAD length (0-512 bytes)
        use rand::RngCore;
        let mut rng_core = rand::thread_rng();
        let pad_length = rng_core.next_u32() as u16 % 513;

        Ok(Self {
            state: Mutex::new(MseState::Idle),
            local_dh_private_key: Some(private_key),
            local_dh_pubkey,
            remote_dh_pubkey: None,
            shared_secret: None,
            info_hash,
            pad_length,
        })
    }

    /// Phase 1: Build Method Selection payload
    ///
    /// Returns `\x13MSegadd` (supports encryption) or `\x00` (plaintext only)
    pub fn build_method_selection(&self) -> Vec<u8> {
        b"\x13MSegadd".to_vec()
    }

    /// Parse remote Method Selection
    ///
    /// data: Method selection bytes received from remote
    pub fn parse_remote_method_selection(data: &[u8]) -> Result<CryptoMethod> {
        if data.is_empty() {
            return Err(Aria2Error::Parse("Empty method selection".to_string()));
        }

        // Check if MSegadd (\x13MSegadd)
        if data == b"\x13MSegadd" {
            return Ok(CryptoMethod::Rc4); // Supports encryption
        }

        // Check if plaintext mode (\x00)
        if data == b"\x00" || data[0] == 0x00 {
            return Ok(CryptoMethod::Plain);
        }

        Err(Aria2Error::Parse(format!(
            "Invalid method selection: {:?}",
            data
        )))
    }

    /// Phase 2: Build PAD + DH public key + CryptoProvisions payload
    ///
    /// Format:
    /// ```text
    /// ┌────────┬──────────┬─────────────┬────────────┐
    /// │ PAD_D  │ PAD_LEN  │ Crypto_Pro  │ ICB/IV     │
    /// │ (2B BE)│ (2B BE)  │ (2B BE)     │ (optional 16B) │
    /// ├────────┼──────────┼─────────────┼────────────┤
    /// │          DH Public Key (X25519 = 32 bytes)                │
    /// └────────┴──────────┴─────────────┴────────────┘
    /// ```
    pub fn build_key_exchange_payload(&self, crypto_methods: &[CryptoMethod]) -> Result<Vec<u8>> {
        let mut payload = Vec::new();

        // PAD_D (2 bytes big-endian): Random padding length
        payload.extend_from_slice(&self.pad_length.to_be_bytes());

        // PAD_LEN (2 bytes big-endian): Same as PAD_D
        payload.extend_from_slice(&self.pad_length.to_be_bytes());

        // Crypto_Provisions (2 bytes big-endian): Supported encryption method bitmask
        let mut crypto_provisions: u16 = 0;
        for method in crypto_methods {
            crypto_provisions |= method.to_u16();
        }
        payload.extend_from_slice(&crypto_provisions.to_be_bytes());

        // Random PAD data
        use rand::RngCore;
        let mut rng_core = rand::thread_rng();
        let mut pad_data = vec![0u8; self.pad_length as usize];
        rng_core.fill_bytes(&mut pad_data);
        payload.extend_from_slice(&pad_data);

        // DH Public Key (32 bytes for X25519)
        payload.extend_from_slice(&self.local_dh_pubkey);

        Ok(payload)
    }

    /// Parse remote key exchange payload, compute shared secret
    pub fn process_remote_key_exchange(&mut self, data: &[u8]) -> Result<()> {
        // Minimum payload = 6-byte header (PAD_D + PAD_LEN + Crypto_Pro) + 32-byte
        // DH public key = 38 bytes. PAD data length is variable (0..=512) and is
        // validated separately below against the embedded PAD_LEN field.
        if data.len() < 6 + 32 {
            return Err(Aria2Error::Parse(
                "Key exchange payload too short".to_string(),
            ));
        }

        // Parse fields
        let _pad_d = u16::from_be_bytes([data[0], data[1]]);
        let _pad_len = u16::from_be_bytes([data[2], data[3]]);
        let _crypto_pro = u16::from_be_bytes([data[4], data[5]]);
        // Skip possible ICB/IV (6-22)

        // Calculate PAD end position
        let pad_end = 6 + (_pad_len as usize);
        if data.len() < pad_end + 32 {
            return Err(Aria2Error::Parse(
                "Key exchange payload truncated".to_string(),
            ));
        }

        // Extract remote DH public key (last 32 bytes)
        let remote_pubkey_start = data.len() - 32;
        let remote_pubkey = &data[remote_pubkey_start..];

        // Store remote public key
        self.remote_dh_pubkey = Some(remote_pubkey.to_vec());

        // Compute shared secret
        let peer_public_key = UnparsedPublicKey::new(&agreement::X25519, remote_pubkey);

        if let Some(private_key) = self.local_dh_private_key.take() {
            let shared_secret: Vec<u8> = agreement::agree_ephemeral(
                private_key,
                &peer_public_key,
                |shared_secret: &[u8]| shared_secret.to_vec(),
            )
            .map_err(|e| {
                Aria2Error::Fatal(crate::error::FatalError::Config(format!(
                    "DH key agreement failed: {}",
                    e
                )))
            })?;

            self.shared_secret = Some(shared_secret);
        } else {
            return Err(Aria2Error::Fatal(crate::error::FatalError::Config(
                "Local DH private key not set".to_string(),
            )));
        }

        Ok(())
    }

    /// Phase 3: Build SKEY + VC + CryptoSelect verification payload
    ///
    /// Format:
    /// - SKEY hash (20 bytes): SHA-1(info_hash + shared_secret)
    /// - VC (2 bytes): version cipher (0x0001 for RC4)
    /// - CryptoSelect (2 bytes): Selected encryption method
    /// - len(I) (2 bytes): Optional initial data length (usually 0)
    pub fn build_verification_payload(&self, selected_method: CryptoMethod) -> Result<Vec<u8>> {
        let skey = self.compute_skey()?;
        let mut payload = Vec::new();

        // SKEY hash (20 bytes)
        payload.extend_from_slice(&skey);

        // VC (2 bytes): Version cipher
        payload.extend_from_slice(&0x0001u16.to_be_bytes());

        // CryptoSelect (2 bytes): Selected encryption method
        payload.extend_from_slice(&selected_method.to_u16().to_be_bytes());

        // len(I) (2 bytes): Initial data length
        payload.extend_from_slice(&0x0000u16.to_be_bytes());

        Ok(payload)
    }

    /// Parse remote verification payload, verify SKEY, complete handshake
    pub fn process_remote_verification(&mut self, data: &[u8]) -> Result<MseCryptoContext> {
        if data.len() < 26 {
            return Err(Aria2Error::Parse(
                "Verification payload too short".to_string(),
            ));
        }

        // Parse fields
        let remote_skey = &data[..20];
        let vc = u16::from_be_bytes([data[20], data[21]]);
        let crypto_select = u16::from_be_bytes([data[22], data[23]]);
        // len(I) at [24,25]

        // Validate VC
        if vc != 0x0001 {
            return Err(Aria2Error::Parse(format!("Invalid VC value: {:#06x}", vc)));
        }

        // Validate SKEY
        let expected_skey = self.compute_skey()?;
        if remote_skey != expected_skey.as_slice() {
            return Err(Aria2Error::Checksum("SKEY verification failed".to_string()));
        }

        // Parse selected encryption method
        let selected_method = CryptoMethod::from_u16(crypto_select).ok_or_else(|| {
            Aria2Error::Parse(format!(
                "Unknown crypto method selected: {:#06x}",
                crypto_select
            ))
        })?;

        // Derive encryption keys
        let shared_secret = self.shared_secret.as_ref().ok_or_else(|| {
            Aria2Error::Fatal(crate::error::FatalError::Config(
                "Shared secret not computed".to_string(),
            ))
        })?;

        let (send_key, recv_key) = Self::derive_keys(&expected_skey, shared_secret);

        // Create encryption context
        let ctx = MseCryptoContext::new(&send_key, &recv_key, selected_method);

        Ok(ctx)
    }

    /// Internal: Compute SKEY = SHA-1(info_hash || shared_secret)
    pub(crate) fn compute_skey(&self) -> Result<[u8; 20]> {
        let shared_secret = self.shared_secret.as_ref().ok_or_else(|| {
            Aria2Error::Fatal(crate::error::FatalError::Config(
                "Shared secret not computed yet".to_string(),
            ))
        })?;

        let mut hasher = Sha1::new();
        hasher.update(self.info_hash);
        hasher.update(shared_secret);
        let result = hasher.finalize();

        let mut skey = [0u8; 20];
        skey.copy_from_slice(&result);
        Ok(skey)
    }

    /// Internal: Derive encryption keys
    ///
    /// ```text
    /// send_key = SHA-1(SKEY + "keyA" + shared_secret)[:16]
    /// recv_key = SHA-1(SKEY + "keyB" + shared_secret)[:16]
    /// ```
    pub(crate) fn derive_keys(skey: &[u8], shared_secret: &[u8]) -> (Vec<u8>, Vec<u8>) {
        // Send direction key
        let mut hasher_a = Sha1::new();
        hasher_a.update(skey);
        hasher_a.update(b"keyA");
        hasher_a.update(shared_secret);
        let result_a = hasher_a.finalize();

        // Receive direction key
        let mut hasher_b = Sha1::new();
        hasher_b.update(skey);
        hasher_b.update(b"keyB");
        hasher_b.update(shared_secret);
        let result_b = hasher_b.finalize();

        // Take first 16 bytes
        (result_a[..16].to_vec(), result_b[..16].to_vec())
    }

    /// Plaintext fallback: create unencrypted context
    pub fn plaintext_fallback() -> MseCryptoContext {
        MseCryptoContext::default()
    }

    /// Get current state
    pub fn state(&self) -> MseState {
        self.state.lock().unwrap_or_else(|e| e.into_inner()).clone()
    }

    /// Update state
    pub fn set_state(&self, state: MseState) {
        *self.state.lock().unwrap_or_else(|e| e.into_inner()) = state;
    }

    /// Get local DH public key (for debugging)
    pub fn local_dh_pubkey(&self) -> &[u8] {
        &self.local_dh_pubkey
    }

    /// Get shared secret (for debugging)
    pub fn shared_secret(&self) -> Option<&[u8]> {
        self.shared_secret.as_deref()
    }
}

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

    #[test]
    fn test_crypto_method_conversions() {
        assert_eq!(CryptoMethod::from_u16(0x0001), Some(CryptoMethod::Plain));
        assert_eq!(CryptoMethod::from_u16(0x0002), Some(CryptoMethod::Rc4));
        assert_eq!(
            CryptoMethod::from_u16(0x0003),
            Some(CryptoMethod::Aes128Cbc)
        );
        assert_eq!(CryptoMethod::from_u16(0x9999), None);

        assert_eq!(CryptoMethod::Plain.to_u16(), 0x0001);
        assert_eq!(CryptoMethod::Rc4.to_u16(), 0x0002);
        assert_eq!(CryptoMethod::Aes128Cbc.to_u16(), 0x0003);
    }
}