#!/usr/libexec/platform-python
#
# nicrx_drop  Network RX drop analysis tool
#       For Linux, uses BCC, eBPF.
#
# This tool analyzes network RX drop by monitoring:
# 1. Packet drop detection: monitor /proc/net/dev for rx drop
# 2. Drop stack trace: capture kfree_skb stack traces in RX context
# 3. Memory allocation failure: capture alloc failure stack traces in RX context
# 4. NIC queue check: capture ethtool information

from __future__ import print_function
from bcc import BPF
import argparse
import time
import subprocess
import re
from collections import defaultdict
from datetime import datetime

def read_proc_net_dev():
    """Read /proc/net/dev and return interface names and rx_drop stats"""
    interfaces = []
    stats = {}
    try:
        with open('/proc/net/dev', 'r') as f:
            for line in f:
                line = line.strip()
                if ':' in line:
                    parts = line.split(':')
                    ifname = parts[0].strip()
                    interfaces.append(ifname)
                    data = parts[1].split()
                    if len(data) >= 4:
                        stats[ifname] = int(data[3])
    except:
        pass
    return interfaces, stats

def filter_existing_interfaces(requested_interfaces):
    """Filter out interfaces that don't exist"""
    available, _ = read_proc_net_dev()
    existing = [iface for iface in requested_interfaces if iface in available]
    missing = [iface for iface in requested_interfaces if iface not in available]
    return existing, missing

# arguments
examples = """examples:
    ./netrx_drop.py                    # monitor eth0 and eth1, 1 iteration, 10 seconds
    ./netrx_drop.py -i eth0            # monitor eth0 only
    ./netrx_drop.py -i eth0 -c 3       # 3 iterations
    ./netrx_drop.py -i eth0 -d 30      # 30 seconds per iteration
    ./netrx_drop.py -i eth0 -c 2 -d 20 # 2 iterations, 20 seconds each
"""
parser = argparse.ArgumentParser(
    description="Network RX drop analysis",
    formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog=examples)
parser.add_argument("-i", "--interface", 
                    help="network interface(s), comma separated (default: eth0,eth1)")
parser.add_argument("-c", "--count", type=int, default=1,
                    help="number of iterations (default: 1. max: 20)")
parser.add_argument("-d", "--duration", type=int, default=10,
                    help="duration per iteration in seconds (default: 10 seconds. max: 60 seconds )")
parser.add_argument("--ebpf", action="store_true",
                    help=argparse.SUPPRESS)
args = parser.parse_args()

# Limit duration to maximum 60 seconds
if args.duration > 60:
    print("Warning: Duration exceeds maximum of 60 seconds. Setting to 60 seconds.")
    args.duration = 60

# Limit count to maximum 20 iterations
if args.count > 20:
    print("Warning: Count exceeds maximum of 20 iterations. Setting to 20.")
    args.count = 20

# parse interfaces
if args.interface:
    requested_interfaces = [iface.strip() for iface in args.interface.split(',')]
else:
    requested_interfaces = ['eth0', 'eth1']

# Filter out interfaces that don't exist
interfaces, missing_interfaces = filter_existing_interfaces(requested_interfaces)

if missing_interfaces:
    print("Warning: The following interfaces do not exist and will be skipped: %s" % ', '.join(missing_interfaces))

if not interfaces:
    print("Error: No valid interfaces found. Available interfaces:")
    available, _ = read_proc_net_dev()
    if available:
        print("  %s" % ', '.join(available))
    else:
        print("  (none)")
    exit(1)

# BPF program for kfree_skb stack trace
bpf_text_kfree_skb = """
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>
#include <linux/netdevice.h>
#include <linux/skbuff.h>

#define IFNAMSIZ 16

struct kfree_skb_event {
    u64 timestamp;
    u32 cpu;
    u32 reason;
    char ifname[IFNAMSIZ];
    u32 stack_id;
    u32 pid;
    char comm[TASK_COMM_LEN];
};

BPF_STACK_TRACE(stack_traces_kfree, 2048);
BPF_PERF_OUTPUT(events_kfree);

TRACEPOINT_PROBE(skb, kfree_skb) {
    u32 reason = (u32)args->reason;
    
    // Skip normal consumption
    if (reason == 0) {
        return 0;
    }
    
    struct sk_buff *skb = (struct sk_buff *)args->skbaddr;
    struct kfree_skb_event event = {};
    
    event.timestamp = bpf_ktime_get_ns();
    event.cpu = bpf_get_smp_processor_id();
    event.reason = reason;
    event.stack_id = stack_traces_kfree.get_stackid(args, BPF_F_REUSE_STACKID);
    event.pid = bpf_get_current_pid_tgid() >> 32;
    bpf_get_current_comm(&event.comm, sizeof(event.comm));
    
    // Get interface name from skb
    if (skb) {
        struct net_device *dev = NULL;
        bpf_probe_read_kernel(&dev, sizeof(dev), &skb->dev);
        if (dev) {
            bpf_probe_read_kernel_str(event.ifname, sizeof(event.ifname), dev->name);
        }
    }
    
    events_kfree.perf_submit(args, &event, sizeof(event));
    return 0;
}
"""

# BPF program for memory allocation failure
bpf_text_alloc = """
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>

struct alloc_fail_event {
    u64 timestamp;
    u32 cpu;
    u32 stack_id;
    u32 pid;
    char comm[TASK_COMM_LEN];
    u32 gfp_mask;
};

BPF_STACK_TRACE(stack_traces_alloc, 2048);
BPF_PERF_OUTPUT(events_alloc);

int trace_warn_alloc(struct pt_regs *ctx) {
    struct alloc_fail_event event = {};
    
    event.timestamp = bpf_ktime_get_ns();
    event.cpu = bpf_get_smp_processor_id();
    event.stack_id = stack_traces_alloc.get_stackid(ctx, BPF_F_REUSE_STACKID);
    event.pid = bpf_get_current_pid_tgid() >> 32;
    bpf_get_current_comm(&event.comm, sizeof(event.comm));
    event.gfp_mask = (u32)PT_REGS_PARM1(ctx);
    
    events_alloc.perf_submit(ctx, &event, sizeof(event));
    return 0;
}
"""

# Network RX path indicators
NET_RX_INDICATORS = [
    'netif_receive_skb',
    '__netif_receive_skb',
    '__netif_receive_skb_core',
    'net_rx_action',
    'napi_gro_receive',
    'virtnet_poll',
    'try_fill_recv',
    'igb_clean_rx_irq',
    'ixgbe_clean_rx_irq',
]

def is_net_rx_path(bpf_obj, stack_traces, stack_id):
    """Check if stack trace is in network RX path"""
    try:
        stack = stack_traces.walk(stack_id)
        for addr in stack:
            sym = bpf_obj.ksym(addr)
            if sym:
                for indicator in NET_RX_INDICATORS:
                    if indicator in sym:
                        return True
    except:
        pass
    return False

class DropDetector:
    """Module 1: Drop detection module"""
    def __init__(self, interfaces):
        self.interfaces = interfaces
        self.initial_stats = {}
        self.final_stats = {}
        self.drops_detected = False
        
    def start(self):
        """Start detection - read initial stats"""
        _, self.initial_stats = read_proc_net_dev()
        
    def stop(self):
        """Stop detection - read final stats and check for drops"""
        _, self.final_stats = read_proc_net_dev()
        self.drops_detected = False
        
        for ifname in self.interfaces:
            initial = self.initial_stats.get(ifname, 0)
            final = self.final_stats.get(ifname, 0)
            if final > initial:
                self.drops_detected = True
                break
                
        return self.drops_detected
    
    def get_drop_count(self, ifname):
        """Get drop count for interface"""
        initial = self.initial_stats.get(ifname, 0)
        final = self.final_stats.get(ifname, 0)
        return final - initial

class StackCollector:
    """Module 2: Drop stack trace collector"""
    def __init__(self, interfaces, bpf_obj):
        self.interfaces = interfaces
        self.bpf = bpf_obj
        self.stack_traces = bpf_obj.get_table("stack_traces_kfree")
        self.stacks = defaultdict(int)
        self.total_count = 0
        
    def process_event(self, cpu, data, size):
        """Process kfree_skb event"""
        event = self.bpf["events_kfree"].event(data)
        
        # Check if in RX path
        if not is_net_rx_path(self.bpf, self.stack_traces, event.stack_id):
            return
            
        # Check interface
        ifname = event.ifname.decode('utf-8', 'replace').strip('\x00')
        if ifname not in self.interfaces:
            return
            
        # Record stack trace
        stack_key = str(event.stack_id)
        self.stacks[stack_key] += 1
        self.total_count += 1
        
    def get_top_stacks(self, n=3):
        """Get top N stack traces"""
        sorted_stacks = sorted(self.stacks.items(), key=lambda x: x[1], reverse=True)
        return sorted_stacks[:n]
    
    def print_results(self):
        """Print results"""
        if self.total_count == 0:
            return
            
        print("\n=== Drop Stack Trace Module ===")
        print("Total drop events in RX context: %d" % self.total_count)
        print("\nTop 3 stack traces:")
        
        top_stacks = self.get_top_stacks(3)
        for idx, (stack_key, count) in enumerate(top_stacks, 1):
            print("\n%d. [%d times]" % (idx, count))
            try:
                stack_id = int(stack_key)
                stack = self.stack_traces.walk(stack_id)
                for addr in stack:
                    sym = self.bpf.ksym(addr)
                    if sym:
                        print("    %s" % sym)
            except:
                print("    [Unable to resolve stack]")

class MemoryDetector:
    """Module 3: Memory allocation failure detector"""
    def __init__(self, interfaces, bpf_obj):
        self.interfaces = interfaces
        self.bpf = bpf_obj
        self.stack_traces = bpf_obj.get_table("stack_traces_alloc")
        self.stacks = defaultdict(int)
        self.total_count = 0
        
    def process_event(self, cpu, data, size):
        """Process memory allocation failure event"""
        event = self.bpf["events_alloc"].event(data)
        
        # Check if in RX path
        if not is_net_rx_path(self.bpf, self.stack_traces, event.stack_id):
            return
            
        # Record stack trace
        stack_key = str(event.stack_id)
        self.stacks[stack_key] += 1
        self.total_count += 1
        
    def get_top_stacks(self, n=3):
        """Get top N stack traces"""
        sorted_stacks = sorted(self.stacks.items(), key=lambda x: x[1], reverse=True)
        return sorted_stacks[:n]
    
    def print_results(self):
        """Print results"""
        if self.total_count == 0:
            return
            
        print("\n=== Memory Allocation Failure Module ===")
        print("Total alloc failures in RX context: %d" % self.total_count)
        print("\nTop 3 stack traces:")
        
        top_stacks = self.get_top_stacks(3)
        for idx, (stack_key, count) in enumerate(top_stacks, 1):
            print("\n%d. [%d times]" % (idx, count))
            try:
                stack_id = int(stack_key)
                stack = self.stack_traces.walk(stack_id)
                for addr in stack:
                    sym = self.bpf.ksym(addr)
                    if sym:
                        print("    %s" % sym)
            except:
                print("    [Unable to resolve stack]")

class NICQueueChecker:
    """Module 4: NIC queue checker"""
    def __init__(self, interfaces):
        self.interfaces = interfaces
        self.ethtool_g_info = {}
        self.ethtool_s_stats = {}
        
        # Key variables to extract from ethtool -S
        self.key_variables = [
            'rx_dropped',
            'rx_errors',
            'rx_missed_errors',
            'rx_no_buffer',
            'rx_fifo_errors',
            'rx_over_errors',
        ]
        
    def collect_info(self):
        """Collect ethtool information"""
        # Find ethtool command path
        ethtool_cmd = 'ethtool'
        try:
            result = subprocess.run(['which', 'ethtool'], 
                                  stdout=subprocess.PIPE, stderr=subprocess.PIPE, 
                                  universal_newlines=True, timeout=2)
            if result.returncode == 0 and result.stdout.strip():
                ethtool_cmd = result.stdout.strip()
        except:
            pass  # Use default 'ethtool'
        
        for ifname in self.interfaces:
            # ethtool -g
            try:
                result = subprocess.run([ethtool_cmd, '-g', ifname],
                                      stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                      universal_newlines=True, timeout=5)
                if result.returncode == 0:
                    self.ethtool_g_info[ifname] = result.stdout
                else:
                    error_msg = result.stderr.strip() if result.stderr else "Return code: %d" % result.returncode
                    self.ethtool_g_info[ifname] = "Failed: %s" % error_msg
            except subprocess.TimeoutExpired:
                self.ethtool_g_info[ifname] = "Failed: Command timeout"
            except FileNotFoundError:
                self.ethtool_g_info[ifname] = "Failed: ethtool command not found (tried: %s)" % ethtool_cmd
            except Exception as e:
                self.ethtool_g_info[ifname] = "Failed: %s" % str(e)
                
            # ethtool -S - extract key variables
            try:
                result = subprocess.run([ethtool_cmd, '-S', ifname],
                                      stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                      universal_newlines=True, timeout=5)
                if result.returncode == 0:
                    stats = {}
                    for line in result.stdout.split('\n'):
                        line = line.strip()
                        if ':' in line:
                            parts = line.split(':', 1)
                            if len(parts) == 2:
                                key = parts[0].strip()
                                value = parts[1].strip()
                                # Check if this is one of our key variables
                                for var in self.key_variables:
                                    if var.lower() in key.lower():
                                        stats[key] = value
                                        break
                    self.ethtool_s_stats[ifname] = stats
                else:
                    error_msg = result.stderr.strip() if result.stderr else "Return code: %d" % result.returncode
                    self.ethtool_s_stats[ifname] = {"_error": error_msg}
            except subprocess.TimeoutExpired:
                self.ethtool_s_stats[ifname] = {"_error": "Command timeout"}
            except FileNotFoundError:
                self.ethtool_s_stats[ifname] = {"_error": "ethtool command not found (tried: %s)" % ethtool_cmd}
            except Exception as e:
                self.ethtool_s_stats[ifname] = {"_error": str(e)}
    
    def print_results(self):
        """Print results"""
        print("\n=== NIC Queue Check Module ===")
        for ifname in self.interfaces:
            print("\nInterface: %s" % ifname)
            print("\nRing buffer info (ethtool -g):")
            print(self.ethtool_g_info.get(ifname, "N/A"))
            
            print("\nKey statistics (ethtool -S - current values):")
            stats = self.ethtool_s_stats.get(ifname, {})
            if stats:
                if "_error" in stats:
                    print("  Error: %s" % stats["_error"])
                else:
                    for key in sorted(stats.keys()):
                        print("  %s: %s" % (key, stats[key]))
            else:
                print("  No key statistics found")

def run_iteration(iteration_num, duration, interfaces):
    """Run one iteration"""
    print("\n" + "="*60)
    print("Iteration %d/%d (Duration: %d seconds)" % (iteration_num, args.count, duration))
    print("="*60)
    
    # Initialize modules
    drop_detector = DropDetector(interfaces)
    stack_collector = None
    memory_detector = None
    nic_checker = NICQueueChecker(interfaces)
    bpf_kfree = None
    bpf_alloc = None
    
    # Load BPF programs (Module 2 and 3 start first)
    try:
        bpf_kfree = BPF(text=bpf_text_kfree_skb)
        bpf_alloc = BPF(text=bpf_text_alloc)
        
        # TRACEPOINT_PROBE(skb, kfree_skb) automatically attaches when BPF program loads
        # The tracepoint "skb:kfree_skb" is attached by BCC's TRACEPOINT_PROBE macro
        
        # Attach kprobe for warn_alloc
        try:
            bpf_alloc.attach_kprobe(event="warn_alloc", fn_name="trace_warn_alloc")
            print("warn_alloc kprobe attached successfully")
        except Exception as e:
            print("Warning: Failed to attach kprobe for warn_alloc: %s" % str(e))
            print("Memory allocation failure detection may not work")
        
        stack_collector = StackCollector(interfaces, bpf_kfree)
        memory_detector = MemoryDetector(interfaces, bpf_alloc)
        
        # Open perf buffers
        bpf_kfree["events_kfree"].open_perf_buffer(stack_collector.process_event)
        bpf_alloc["events_alloc"].open_perf_buffer(memory_detector.process_event)
        
        print("BPF programs loaded successfully")
    except Exception as e:
        print("Warning: Failed to load BPF programs: %s" % str(e))
        print("Continuing with drop detection only...")
    
    # Start monitoring
    start_time = time.time()
    end_time = start_time + duration
    
    print("\nMonitoring interfaces: %s" % ', '.join(interfaces))
    print("Monitoring started at: %s" % datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
    
    # Start drop detection module last (Module 1: last to start)
    # Wait a bit to let other modules start first
    time.sleep(0.5)
    drop_detector.start()
    
    # Poll BPF events and monitor
    drop_detection_end_time = end_time - 1.0  # Stop drop detection 1 second earlier
    
    try:
        while time.time() < end_time:
            current_time = time.time()
            
            # Poll BPF events (Module 2 and 3: concurrent via async events)
            if bpf_kfree:
                bpf_kfree.perf_buffer_poll(timeout=100)
            if bpf_alloc:
                bpf_alloc.perf_buffer_poll(timeout=100)
            
            # Stop drop detection first (Module 1: first to stop)
            # Note: /proc/net/dev rx_dropped is cumulative, so we only need to read at start and end
            if current_time >= drop_detection_end_time and not hasattr(drop_detector, '_stopped'):
                drops_detected = drop_detector.stop()
                drop_detector._stopped = True
                if drops_detected:
                    print("Drops detected! Other modules continue to capture data...")
                    # Let other modules continue for a bit more to ensure all drops are captured
                    time.sleep(0.5)
            
            time.sleep(0.1)
    except KeyboardInterrupt:
        print("\nInterrupted by user")
        if not hasattr(drop_detector, '_stopped'):
            drop_detector.stop()
    
    # Ensure drop detection is stopped
    if not hasattr(drop_detector, '_stopped'):
        drops_detected = drop_detector.stop()
    else:
        drops_detected = drop_detector.drops_detected
    
    print("\nMonitoring ended at: %s" % datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
    
    # Print drop statistics
    print("\n=== Drop Detection Module ===")
    for ifname in interfaces:
        initial = drop_detector.initial_stats.get(ifname, 0)
        final = drop_detector.final_stats.get(ifname, 0)
        drop_count = final - initial
        print("Interface %s:" % ifname)
        print("  Previous value: %d" % initial)
        print("  Current value:  %d" % final)
        print("  Difference:     %d" % drop_count)
    
    # Print other modules if drops were detected
    if drops_detected:
        if stack_collector:
            stack_collector.print_results()
        if memory_detector:
            memory_detector.print_results()
    
    # Module 4: Always print ethtool info (even if no drops detected)
    nic_checker.collect_info()
    nic_checker.print_results()

# Main execution
print("=== Network RX Drop Analysis Tool ===")
print("Interfaces: %s" % ', '.join(interfaces))
print("Iterations: %d" % args.count)
print("Duration per iteration: %d seconds" % args.duration)

for i in range(1, args.count + 1):
    run_iteration(i, args.duration, interfaces)
    if i < args.count:
        time.sleep(1)  # Small delay between iterations

print("\n" + "="*60)
print("Analysis completed")
print("="*60)
