shadow_protocols/

packet_shaper.rs

//! Packet Shaping for Traffic Analysis Resistance
//! 
//! Shapes outgoing traffic to match target patterns and resist statistical analysis.

use std::collections::VecDeque;
use std::time::{Duration, Instant};
use bytes::Bytes;
use rand::Rng;

/// Traffic shaping strategy
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ShapingStrategy {
    /// No shaping (send immediately)
    None,
    /// Constant bitrate (add padding to match rate)
    ConstantRate { bytes_per_sec: u64 },
    /// Packet size distribution matching
    SizeDistribution,
    /// Constant packet rate with padding
    ConstantPacketRate { packets_per_sec: u64 },
    /// Mimic specific protocol
    ProtocolMimic,
}

/// Packet shaper that modifies traffic patterns
pub struct PacketShaper {
    /// Shaping strategy
    strategy: ShapingStrategy,
    /// Pending packets queue
    queue: VecDeque<QueuedPacket>,
    /// Dummy packet generation enabled
    dummy_enabled: bool,
    /// Last send time
    last_send: Instant,
    /// Statistics
    stats: ShapingStats,
}

/// Queued packet with metadata
#[derive(Debug, Clone)]
struct QueuedPacket {
    data: Bytes,
    enqueue_time: Instant,
    priority: u8,
    is_dummy: bool,
}

/// Shaping statistics
#[derive(Debug, Clone, Default)]
pub struct ShapingStats {
    pub packets_shaped: u64,
    pub dummy_packets_sent: u64,
    pub bytes_padded: u64,
    pub avg_queue_delay_ms: f64,
}

impl PacketShaper {
    /// Create new packet shaper
    pub fn new(strategy: ShapingStrategy) -> Self {
        Self {
            strategy,
            queue: VecDeque::new(),
            dummy_enabled: false,
            last_send: Instant::now(),
            stats: ShapingStats::default(),
        }
    }

    /// Enqueue packet for shaping
    pub fn enqueue(&mut self, data: Bytes, priority: u8) {
        let packet = QueuedPacket {
            data,
            enqueue_time: Instant::now(),
            priority,
            is_dummy: false,
        };
        
        self.queue.push_back(packet);
    }

    /// Get next packet to send (with shaping applied)
    pub async fn next_packet(&mut self) -> Option<Bytes> {
        match self.strategy {
            ShapingStrategy::None => {
                self.queue.pop_front().map(|p| p.data)
            }
            
            ShapingStrategy::ConstantRate { bytes_per_sec } => {
                self.constant_rate_shaping(bytes_per_sec).await
            }
            
            ShapingStrategy::ConstantPacketRate { packets_per_sec } => {
                self.constant_packet_rate_shaping(packets_per_sec).await
            }
            
            ShapingStrategy::SizeDistribution => {
                self.size_distribution_shaping().await
            }
            
            ShapingStrategy::ProtocolMimic => {
                self.protocol_mimic_shaping().await
            }
        }
    }

    /// Constant bitrate shaping (adds padding if needed)
    async fn constant_rate_shaping(&mut self, bytes_per_sec: u64) -> Option<Bytes> {
        let target_interval = Duration::from_secs_f64(1.0 / (bytes_per_sec as f64 / 1200.0));
        let elapsed = self.last_send.elapsed();
        
        if elapsed < target_interval {
            tokio::time::sleep(target_interval - elapsed).await;
        }
        
        self.last_send = Instant::now();
        
        // Get next packet or generate dummy
        if let Some(mut packet) = self.queue.pop_front() {
            let delay = packet.enqueue_time.elapsed();
            self.stats.avg_queue_delay_ms = (self.stats.avg_queue_delay_ms * 0.9) + (delay.as_millis() as f64 * 0.1);
            self.stats.packets_shaped += 1;
            
            // Add padding to match target size
            let target_size = 1200;
            if packet.data.len() < target_size {
                let padding_size = target_size - packet.data.len();
                self.stats.bytes_padded += padding_size as u64;
                
                let mut padded = packet.data.to_vec();
                padded.resize(target_size, 0);
                Some(Bytes::from(padded))
            } else {
                Some(packet.data)
            }
        } else if self.dummy_enabled {
            // Generate dummy packet
            self.stats.dummy_packets_sent += 1;
            Some(self.generate_dummy_packet(1200))
        } else {
            None
        }
    }

    /// Constant packet rate shaping
    async fn constant_packet_rate_shaping(&mut self, packets_per_sec: u64) -> Option<Bytes> {
        let target_interval = Duration::from_secs_f64(1.0 / packets_per_sec as f64);
        let elapsed = self.last_send.elapsed();
        
        if elapsed < target_interval {
            tokio::time::sleep(target_interval - elapsed).await;
        }
        
        self.last_send = Instant::now();
        
        // Get next packet or dummy
        if let Some(packet) = self.queue.pop_front() {
            self.stats.packets_shaped += 1;
            Some(packet.data)
        } else if self.dummy_enabled {
            self.stats.dummy_packets_sent += 1;
            Some(self.generate_dummy_packet(self.random_packet_size()))
        } else {
            None
        }
    }

    /// Size distribution matching shaping
    async fn size_distribution_shaping(&mut self) -> Option<Bytes> {
        if let Some(mut packet) = self.queue.pop_front() {
            self.stats.packets_shaped += 1;
            
            // Round size to common distribution bucket
            let target_size = self.round_to_distribution(packet.data.len());
            
            if packet.data.len() < target_size {
                let mut padded = packet.data.to_vec();
                padded.resize(target_size, 0);
                self.stats.bytes_padded += (target_size - packet.data.len()) as u64;
                Some(Bytes::from(padded))
            } else {
                Some(packet.data)
            }
        } else {
            None
        }
    }

    /// Protocol mimicry shaping (e.g., HTTP, WebRTC patterns)
    async fn protocol_mimic_shaping(&mut self) -> Option<Bytes> {
        // For now, use WebRTC-like distribution
        self.size_distribution_shaping().await
    }

    /// Generate dummy packet
    fn generate_dummy_packet(&self, size: usize) -> Bytes {
        let mut rng = rand::thread_rng();
        let data: Vec<u8> = (0..size).map(|_| rng.gen()).collect();
        Bytes::from(data)
    }

    /// Random packet size following realistic distribution
    fn random_packet_size(&self) -> usize {
        let mut rng = rand::thread_rng();
        let r: f64 = rng.gen();
        
        if r < 0.1 {
            // Small packets (10%)
            rng.gen_range(64..256)
        } else if r < 0.7 {
            // Medium packets (60%)
            rng.gen_range(256..1200)
        } else {
            // Large packets (30%)
            rng.gen_range(1200..1500)
        }
    }

    /// Round size to common distribution bucket
    fn round_to_distribution(&self, size: usize) -> usize {
        // Common packet size buckets (reduces identifiable size variance)
        const BUCKETS: &[usize] = &[64, 128, 256, 512, 1024, 1200, 1400];
        
        for &bucket in BUCKETS {
            if size <= bucket {
                return bucket;
            }
        }
        
        1500 // MTU
    }

    /// Enable/disable dummy traffic
    pub fn set_dummy_traffic(&mut self, enabled: bool) {
        self.dummy_enabled = enabled;
    }

    /// Change shaping strategy
    pub fn set_strategy(&mut self, strategy: ShapingStrategy) {
        self.strategy = strategy;
    }

    /// Get current statistics
    pub fn stats(&self) -> &ShapingStats {
        &self.stats
    }

    /// Clear queue
    pub fn clear_queue(&mut self) {
        self.queue.clear();
    }

    /// Queue length
    pub fn queue_len(&self) -> usize {
        self.queue.len()
    }
}

/// Traffic mixer that combines multiple streams
pub struct TrafficMixer {
    /// Input queues (one per priority level)
    queues: Vec<VecDeque<Bytes>>,
    /// Priority weights
    weights: Vec<u8>,
}

impl TrafficMixer {
    /// Create new traffic mixer with priority levels
    pub fn new(priority_levels: usize) -> Self {
        let mut queues = Vec::new();
        let mut weights = Vec::new();
        
        for i in 0..priority_levels {
            queues.push(VecDeque::new());
            // Higher priority = higher weight
            weights.push((priority_levels - i) as u8);
        }
        
        Self { queues, weights }
    }

    /// Add packet to priority queue
    pub fn enqueue(&mut self, data: Bytes, priority: u8) {
        let idx = priority as usize % self.queues.len();
        self.queues[idx].push_back(data);
    }

    /// Get next packet (weighted round-robin)
    pub fn next_packet(&mut self) -> Option<Bytes> {
        // Weighted round-robin scheduling
        let mut rng = rand::thread_rng();
        let total_weight: u32 = self.weights.iter().map(|&w| w as u32).sum();
        let mut choice: u32 = rng.gen_range(0..total_weight);
        
        for (idx, &weight) in self.weights.iter().enumerate() {
            if choice < weight as u32 {
                return self.queues[idx].pop_front();
            }
            choice -= weight as u32;
        }
        
        None
    }
}

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

    #[test]
    fn test_packet_shaper_enqueue() {
        let mut shaper = PacketShaper::new(ShapingStrategy::None);
        
        shaper.enqueue(Bytes::from(vec![1, 2, 3]), 0);
        shaper.enqueue(Bytes::from(vec![4, 5, 6]), 1);
        
        assert_eq!(shaper.queue_len(), 2);
    }

    #[tokio::test]
    async fn test_no_shaping() {
        let mut shaper = PacketShaper::new(ShapingStrategy::None);
        
        shaper.enqueue(Bytes::from(vec![1, 2, 3]), 0);
        
        let packet = shaper.next_packet().await.unwrap();
        assert_eq!(packet.as_ref(), &[1, 2, 3]);
    }

    #[test]
    fn test_size_rounding() {
        let shaper = PacketShaper::new(ShapingStrategy::SizeDistribution);
        
        assert_eq!(shaper.round_to_distribution(100), 128);
        assert_eq!(shaper.round_to_distribution(500), 512);
        assert_eq!(shaper.round_to_distribution(1300), 1400);
    }

    #[test]
    fn test_traffic_mixer() {
        let mut mixer = TrafficMixer::new(3);
        
        mixer.enqueue(Bytes::from(vec![1]), 0);
        mixer.enqueue(Bytes::from(vec![2]), 1);
        mixer.enqueue(Bytes::from(vec![3]), 2);
        
        // Should get packets, prioritizing higher priority
        assert!(mixer.next_packet().is_some());
    }

    #[tokio::test]
    async fn test_dummy_packet_generation() {
        let mut shaper = PacketShaper::new(ShapingStrategy::ConstantPacketRate { packets_per_sec: 10 });
        shaper.set_dummy_traffic(true);
        
        // Should generate dummy packet when queue is empty
        let packet = shaper.next_packet().await;
        assert!(packet.is_some());
        assert_eq!(shaper.stats().dummy_packets_sent, 1);
    }
}