stats.py

"""Statistics tracking for DDoS testing"""

import threading
import time
from datetime import datetime

class StatsTracker:
    """Thread-safe statistics tracker"""
    
    def __init__(self):
        self.lock = threading.Lock()
        self.packets_sent = 0
        self.bytes_sent = 0
        self.start_time = time.time()
        self.last_display = time.time()
        self.last_packets = 0
        self.last_bytes = 0
    
    def increment_packets(self):
        """Increment packet counter"""
        with self.lock:
            self.packets_sent += 1
    
    def add_bytes(self, bytes_count):
        """Add to bytes counter"""
        with self.lock:
            self.bytes_sent += bytes_count
    
    def display_stats(self):
        """Display real-time statistics"""
        with self.lock:
            current_time = time.time()
            elapsed = current_time - self.start_time
            
            # Calculate rates
            packets_per_sec = self.packets_sent / elapsed if elapsed > 0 else 0
            bytes_per_sec = self.bytes_sent / elapsed if elapsed > 0 else 0
            
            # Calculate bandwidth in MB/s
            mbps = (bytes_per_sec * 8) / (1024 * 1024)
            
            # Current rates (last second)
            time_diff = current_time - self.last_display
            if time_diff >= 1:
                current_pps = (self.packets_sent - self.last_packets) / time_diff
                current_bps = (self.bytes_sent - self.last_bytes) / time_diff
                current_mbps = (current_bps * 8) / (1024 * 1024)
                
                self.last_packets = self.packets_sent
                self.last_bytes = self.bytes_sent
                self.last_display = current_time
            else:
                current_pps = packets_per_sec
                current_mbps = mbps
            
            # Display stats
            print(f"\r[{datetime.now().strftime('%H:%M:%S')}] "
                  f"Packets: {self.packets_sent:,} | "
                  f"PPS: {current_pps:.1f} | "
                  f"Bandwidth: {current_mbps:.2f} Mbps | "
                  f"Elapsed: {elapsed:.1f}s", end="", flush=True)
    
    def get_total_stats(self):
        """Get total statistics"""
        with self.lock:
            return {
                'packets': self.packets_sent,
                'bytes': self.bytes_sent,
                'duration': time.time() - self.start_time
            }