#!/usr/bin/env python3
"""
Generate golden test vectors for Packet Filter testbench.
Implements the firewall rules with strict priority order.
"""

import random

def check_packet(src_ip, protocol, dst_port):
    # Rule 1: Source IP in 10.0.0.0/8 -> DROP
    if (src_ip & 0xFF000000) == 0x0A000000:
        return 0 # DROP
        
    # Rule 2: Source IP in 192.168.0.0/16 -> DROP
    if (src_ip & 0xFFFF0000) == 0xC0A80000:
        return 0 # DROP
        
    # Rule 3: Protocol = ICMP (1) -> ALLOW
    if protocol == 1:
        return 1 # ALLOW
        
    # Rule 4: Protocol = TCP (6) AND dst_port = 80 -> ALLOW
    if protocol == 6 and dst_port == 80:
        return 1 # ALLOW
        
    # Rule 5: Protocol = TCP (6) AND dst_port = 443 -> ALLOW
    if protocol == 6 and dst_port == 443:
        return 1 # ALLOW
        
    # Rule 6: Protocol = TCP (6) AND dst_port = 22 -> ALLOW
    if protocol == 6 and dst_port == 22:
        return 1 # ALLOW
        
    # Rule 7: Protocol = UDP (17) AND dst_port = 53 -> ALLOW
    if protocol == 17 and dst_port == 53:
        return 1 # ALLOW
        
    # Rule 8: (Default) No match -> DROP
    return 0 # DROP

def generate_test_cases():
    test_cases = []
    
    # 1. Rule 1: Block 10.x.x.x (even if protocol is allowed)
    # 10.0.0.1, ICMP (Rule 3 says allow, but Rule 1 says drop) -> Expect DROP
    test_cases.append({
        'src_ip': 0x0A000001, 'protocol': 1, 'dst_port': 0, 
        'expected': 0, 'desc': "Rule 1 Block 10.x with ICMP"
    })
    # 10.1.2.3, TCP port 80 -> Expect DROP
    test_cases.append({
        'src_ip': 0x0A010203, 'protocol': 6, 'dst_port': 80, 
        'expected': 0, 'desc': "Rule 1 Block 10.x with HTTP"
    })
    
    # 2. Rule 2: Block 192.168.x.x
    # 192.168.1.1, UDP port 53 (DNS) -> Expect DROP
    test_cases.append({
        'src_ip': 0xC0A80101, 'protocol': 17, 'dst_port': 53, 
        'expected': 0, 'desc': "Rule 2 Block 192.168.x with DNS"
    })
    
    # 3. Rule 3: Allow ICMP (from public IP)
    # 8.8.8.8, ICMP -> Expect ALLOW
    test_cases.append({
        'src_ip': 0x08080808, 'protocol': 1, 'dst_port': 0, 
        'expected': 1, 'desc': "Rule 3 Allow ICMP public"
    })
    
    # 4. Rule 4: Allow HTTP
    test_cases.append({
        'src_ip': 0x01020304, 'protocol': 6, 'dst_port': 80, 
        'expected': 1, 'desc': "Rule 4 Allow HTTP"
    })
    
    # 5. Rule 5: Allow HTTPS
    test_cases.append({
        'src_ip': 0x01020304, 'protocol': 6, 'dst_port': 443, 
        'expected': 1, 'desc': "Rule 5 Allow HTTPS"
    })
    
    # 6. Rule 6: Allow SSH
    test_cases.append({
        'src_ip': 0x01020304, 'protocol': 6, 'dst_port': 22, 
        'expected': 1, 'desc': "Rule 6 Allow SSH"
    })
    
    # 7. Rule 7: Allow DNS
    test_cases.append({
        'src_ip': 0x01020304, 'protocol': 17, 'dst_port': 53, 
        'expected': 1, 'desc': "Rule 7 Allow DNS"
    })
    
    # 8. Rule 8: Default Drop
    # Public IP, TCP random port -> DROP
    test_cases.append({
        'src_ip': 0x01020304, 'protocol': 6, 'dst_port': 12345, 
        'expected': 0, 'desc': "Rule 8 Default Drop TCP"
    })
    # Public IP, UDP random port -> DROP
    test_cases.append({
        'src_ip': 0x01020304, 'protocol': 17, 'dst_port': 12345, 
        'expected': 0, 'desc': "Rule 8 Default Drop UDP"
    })
    # Public IP, Random Protocol -> DROP
    test_cases.append({
        'src_ip': 0x01020304, 'protocol': 99, 'dst_port': 80, 
        'expected': 0, 'desc': "Rule 8 Default Drop Proto"
    })
    
    # 9. Edge Cases / Boundaries
    # 10.255.255.255 -> DROP
    test_cases.append({
        'src_ip': 0x0AFFFFFF, 'protocol': 6, 'dst_port': 80, 
        'expected': 0, 'desc': "Boundary 10.255.255.255"
    })
    # 11.0.0.0 (Public) -> ALLOW if HTTP
    test_cases.append({
        'src_ip': 0x0B000000, 'protocol': 6, 'dst_port': 80, 
        'expected': 1, 'desc': "Boundary 11.0.0.0"
    })
    
    # 10. Random cases
    random.seed(42)
    for i in range(20):
        src_ip = random.randint(0, 0xFFFFFFFF)
        protocol = random.choice([1, 6, 17, 0, 255])
        dst_port = random.choice([80, 443, 22, 53, 0, 65535, random.randint(0, 65535)])
        
        expected = check_packet(src_ip, protocol, dst_port)
        test_cases.append({
            'src_ip': src_ip, 'protocol': protocol, 'dst_port': dst_port,
            'expected': expected, 'desc': f"Random {i}"
        })
        
    return test_cases

def main():
    test_cases = generate_test_cases()
    
    print("// ==============================================")
    print("// Packet Filter Golden Test Vectors")
    print("// Generated by: python3 generate_golden.py")
    print("// ==============================================")
    print()
    print(f"localparam NUM_TESTS = {len(test_cases)};")
    print()
    print("reg [31:0] tb_src_ip [0:{len(test_cases)-1}];")
    print("reg [15:0] tb_dst_port [0:{len(test_cases)-1}];")
    print("reg [7:0]  tb_protocol [0:{len(test_cases)-1}];")
    print("reg        exp_allow [0:{len(test_cases)-1}];")
    print()
    print("initial begin")
    
    for i, tc in enumerate(test_cases):
        print(f"    // Test {i}: {tc['desc']}")
        print(f"    tb_src_ip[{i}]   = 32'h{tc['src_ip']:08x};")
        print(f"    tb_dst_port[{i}] = 16'd{tc['dst_port']};")
        print(f"    tb_protocol[{i}] = 8'd{tc['protocol']};")
        print(f"    exp_allow[{i}]   = 1'b{tc['expected']};")
        print()
        
    print("end")

if __name__ == "__main__":
    main()