#!/usr/bin/env python3
"""
generate_corpus.py - Build a tiny synthetic/defanged PCAP for the Zeek+Suricata lab.

Requires: scapy  (pip install scapy)

The resulting corpus.pcap is deliberately BENIGN and SYNTHETIC:
  - All IP addresses are RFC 5737 documentation addresses (192.0.2.x, 198.51.100.x)
    or RFC 1918 private ranges; they resolve to nothing real.
  - All domain names are in the .invalid / .test / .example TLD space.
  - No real malware payload. The "beacon" is a repeated HTTP GET with a
    suspicious fixed-interval cadence and an unusual user-agent string;
    the "C2 domain" is lab-c2.example.invalid.
  - Files transferred are EICAR-format test strings or empty payloads.

This script writes corpus.pcap into the same directory it lives in.
Run once; the file is then mounted into the network-tools container.
"""

import os
import struct
import time

PCAP_GLOBAL_HEADER = struct.pack(
    '<IHHiIII',
    0xa1b2c3d4,  # magic
    2, 4,         # major, minor version
    0,            # timezone offset
    0,            # accuracy of timestamps
    65535,        # snaplen
    1,            # link-layer type: LINKTYPE_ETHERNET
)

def pcap_record(ts_sec, ts_usec, data):
    return struct.pack('<IIII', ts_sec, ts_usec, len(data), len(data)) + data

def eth_ip_tcp(src_mac, dst_mac, src_ip, dst_ip, src_port, dst_port, payload):
    """Build a minimal Ethernet/IPv4/TCP frame (no real checksum)."""
    # Ethernet header
    def mac_bytes(m):
        return bytes(int(x, 16) for x in m.split(':'))

    eth = mac_bytes(dst_mac) + mac_bytes(src_mac) + b'\x08\x00'

    # IPv4 header (20 bytes, no options)
    tcp_payload = payload.encode() if isinstance(payload, str) else payload
    tcp_len = 20 + len(tcp_payload)  # TCP header + data
    ip_total = 20 + tcp_len
    ip_hdr = struct.pack(
        '>BBHHHBBH4s4s',
        0x45, 0,          # version+IHL, DSCP+ECN
        ip_total,         # total length
        0, 0,             # identification, flags+fragment offset
        64, 6,            # TTL, protocol=TCP
        0,                # checksum (0 = not computed)
        bytes(int(b) for b in src_ip.split('.')),
        bytes(int(b) for b in dst_ip.split('.')),
    )
    # TCP header (20 bytes)
    tcp_hdr = struct.pack(
        '>HHIIBBHHH',
        src_port, dst_port,
        1000, 2000,        # seq, ack
        0x50, 0x18,        # data offset=5, flags=ACK+PSH
        65535, 0, 0        # window, checksum, urgent
    )
    return eth + ip_hdr + tcp_hdr + tcp_payload

def eth_ip_udp(src_mac, dst_mac, src_ip, dst_ip, src_port, dst_port, payload):
    """Build a minimal Ethernet/IPv4/UDP frame."""
    def mac_bytes(m):
        return bytes(int(x, 16) for x in m.split(':'))
    eth = mac_bytes(dst_mac) + mac_bytes(src_mac) + b'\x08\x00'
    udp_payload = payload.encode() if isinstance(payload, str) else payload
    udp_len = 8 + len(udp_payload)
    ip_total = 20 + udp_len
    ip_hdr = struct.pack(
        '>BBHHHBBH4s4s',
        0x45, 0, ip_total, 0, 0, 64, 17, 0,
        bytes(int(b) for b in src_ip.split('.')),
        bytes(int(b) for b in dst_ip.split('.')),
    )
    udp_hdr = struct.pack('>HHHH', src_port, dst_port, udp_len, 0)
    return eth + ip_hdr + udp_hdr + udp_payload

# Addresses - all documentation/RFC ranges, not routable
VICTIM    = '192.0.2.10'
GW        = '192.0.2.1'
C2        = '198.51.100.42'       # "attacker" destination
BENIGN    = '198.51.100.100'      # normal external host
DNS_SRV   = '192.0.2.53'

MAC_VICTIM  = 'aa:bb:cc:dd:ee:01'
MAC_GW      = 'aa:bb:cc:dd:ee:02'

# Build a synthetic DNS query for lab-c2.example.invalid
# (raw DNS wire format, minimal)
def dns_query_packet(qname):
    parts = qname.encode().split(b'.')
    labels = b''.join(struct.pack('B', len(p)) + p for p in parts) + b'\x00'
    header = struct.pack('>HHHHHHH', 0x1234, 0x0100, 1, 0, 0, 0, 0)  # ID, flags, 1 question, 0 answers, 0 auth, 0 add
    question = labels + struct.pack('>HH', 1, 1)  # QTYPE=A, QCLASS=IN
    return header + question

BASE_TIME = 1_700_000_000  # fixed epoch so pcap is deterministic

out_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'corpus.pcap')

records = []

# 1. DNS query for C2 domain (repeated 10 times, 60s apart = beacon pattern)
for i in range(10):
    ts = BASE_TIME + i * 60
    dns_pkt = dns_query_packet('lab-c2.example.invalid')
    frame = eth_ip_udp(MAC_VICTIM, MAC_GW, VICTIM, DNS_SRV, 54321, 53, dns_pkt)
    records.append(pcap_record(ts, 0, frame))

# 2. HTTP beacon to C2 (10 times, 60s apart, ~same byte size each time)
beacon_http = (
    'GET /check HTTP/1.1\r\n'
    'Host: lab-c2.example.invalid\r\n'
    'User-Agent: EvilBot/1.0\r\n'
    'Accept: */*\r\n'
    'Connection: keep-alive\r\n'
    '\r\n'
)
for i in range(10):
    ts = BASE_TIME + i * 60 + 1  # 1s after DNS
    frame = eth_ip_tcp(MAC_VICTIM, MAC_GW, VICTIM, C2, 49152 + i, 80, beacon_http)
    records.append(pcap_record(ts, 0, frame))

# 3. HTTP response with EICAR-format test string (not real malware)
eicar_str = (
    'HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream\r\n'
    'Content-Disposition: attachment; filename="test.exe"\r\n'
    'Content-Length: 68\r\n\r\n'
    'X5O!P%@AP[4\\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*'
)
ts = BASE_TIME + 5
frame = eth_ip_tcp(MAC_GW, MAC_VICTIM, C2, VICTIM, 80, 49152, eicar_str)
records.append(pcap_record(ts, 0, frame))

# 4. Normal HTTPS to benign host (should NOT trigger C2 alerts)
normal_tls = b'\x16\x03\x01' + b'\x00' * 20  # minimal TLS client hello stub
frame = eth_ip_tcp(MAC_VICTIM, MAC_GW, VICTIM, BENIGN, 55000, 443, normal_tls)
records.append(pcap_record(BASE_TIME + 3600, 0, frame))

with open(out_path, 'wb') as f:
    f.write(PCAP_GLOBAL_HEADER)
    for r in sorted(records, key=lambda x: struct.unpack('<I', x[:4])[0]):
        f.write(r)

print(f'Written: {out_path}  ({os.path.getsize(out_path)} bytes)')
print('This is a SYNTHETIC/DEFANGED corpus. No real malware. IPs are RFC 5737/1918.')
