import json
import time
import random
from datetime import datetime
from kafka import KafkaProducer

# Configure Kafka Producer
producer = KafkaProducer(
    bootstrap_servers=['localhost:9092'],
    value_serializer=lambda x: json.dumps(x).encode('utf-8')
)

# Fake IP addresses
NORMAL_IPS  = [f"192.168.1.{i}" for i in range(1, 20)]
ATTACK_IPS  = ["10.0.0.99", "185.220.101.5", "45.33.32.156"]
PORTS       = [22, 80, 443, 3306, 8080]
PROTOCOLS   = ["TCP", "UDP"]

def normal_log():
    return {
        "timestamp":   datetime.now().isoformat(),
        "src_ip":      random.choice(NORMAL_IPS),
        "dst_ip":      "192.168.1.100",
        "src_port":    random.randint(1024, 65535),
        "dst_port":    random.choice(PORTS),
        "protocol":    random.choice(PROTOCOLS),
        "bytes":       random.randint(64, 1500),
        "packets":     random.randint(1, 10),
        "label":       "NORMAL"
    }

def brute_force_log():
    return {
        "timestamp":   datetime.now().isoformat(),
        "src_ip":      random.choice(ATTACK_IPS),
        "dst_ip":      "192.168.1.100",
        "src_port":    random.randint(1024, 65535),
        "dst_port":    22,   # SSH port
        "protocol":    "TCP",
        "bytes":       random.randint(40, 100),
        "packets":     1,
        "label":       "BRUTE_FORCE"
    }

def ddos_log():
    return {
        "timestamp":   datetime.now().isoformat(),
        "src_ip":      f"{random.randint(1,254)}.{random.randint(1,254)}.{random.randint(1,254)}.{random.randint(1,254)}",
        "dst_ip":      "192.168.1.100",
        "src_port":    random.randint(1024, 65535),
        "dst_port":    80,
        "protocol":    "UDP",
        "bytes":       random.randint(1000, 1500),
        "packets":     random.randint(50, 200),
        "label":       "DDOS"
    }

print("Starting log generation...")
count = 0

while True:
    # Traffic distribution: 70% normal, 20% Brute Force, 10% DDoS
    rand = random.random()
    if rand < 0.70:
        log = normal_log()
    elif rand < 0.90:
        log = brute_force_log()
    else:
        log = ddos_log()

    # Send log to Kafka topic
    producer.send('network-logs', value=log)
    count += 1

    # Print progress every 10 logs
    if count % 10 == 0:
        print(f"[+] Sent {count} logs | Last type: {log['label']}")

    # Send one log every 0.5 seconds
    time.sleep(0.5)
