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

# ============================================================
# CONFIGURATION
# ============================================================
TOTAL_RECORDS = 200          # إجمالي عدد السجلات المُرسَلة
ATTACK_RATIO = 0.10          # نسبة الهجوم (10%)
KAFKA_BROKER = '10.0.0.15:9093'
DB_HOST = "10.0.0.18"
TIMEOUT_SECONDS = 120         # مدة انتظار الكشف بعد إرسال آخر سجل

# ============================================================
# TEMPLATES (from real CIC-IDS2018 records)
# ============================================================
BENIGN_TEMPLATE = {
    "Dst Port": 80.0, "Protocol": 6.0, "Flow Duration": 10085.0,
    "Tot Fwd Pkts": 3.0, "Tot Bwd Pkts": 4.0,
    "TotLen Fwd Pkts": 316.0, "TotLen Bwd Pkts": 935.0,
    "Fwd Pkt Len Max": 316.0, "Fwd Pkt Len Min": 0.0,
    "Bwd Pkt Len Max": 935.0, "Bwd Pkt Len Min": 0.0,
    "Flow Pkts/s": 694.1001487, "Flow Byts/s": 124045.6123,
    "Fwd Header Len": 72.0, "Bwd Header Len": 92.0,
    "Fwd Pkts/s": 297.4714923, "Bwd Pkts/s": 396.6286564,
    "Pkt Len Min": 0.0, "Pkt Len Max": 935.0,
    "SYN Flag Cnt": 0.0, "RST Flag Cnt": 1.0,
    "PSH Flag Cnt": 1.0, "ACK Flag Cnt": 0.0,
}

ATTACK_TEMPLATE = {
    "Dst Port": 880.0, "Protocol": 6.0, "Flow Duration": 1.0,
    "Tot Fwd Pkts": 3.0, "Tot Bwd Pkts": 4.0,
    "TotLen Fwd Pkts": 316.0, "TotLen Bwd Pkts": 935.0,
    "Fwd Pkt Len Max": 316.0, "Fwd Pkt Len Min": 0.0,
    "Bwd Pkt Len Max": 935.0, "Bwd Pkt Len Min": 0.0,
    "Flow Pkts/s": 6.1001487, "Flow Byts/s": 15.6123,
    "Fwd Header Len": 7.0, "Bwd Header Len": 9.0,
    "Fwd Pkts/s": 97.4714923, "Bwd Pkts/s": 36.6286564,
    "Pkt Len Min": 8.0, "Pkt Len Max": 95.0,
    "SYN Flag Cnt": 7.0, "RST Flag Cnt": 3.0,
    "PSH Flag Cnt": 2.0, "ACK Flag Cnt": 1.0,
}

# ============================================================
# GENERATE UNIQUE IP POOLS (test-only ranges, won't collide)
# ============================================================
def gen_attack_ip(i):
    return f"203.0.113.{i % 254 + 1}"

def gen_benign_ip(i):
    return f"198.51.100.{i % 254 + 1}"

# ============================================================
# BUILD MIXED RECORD SET
# ============================================================
n_attacks = int(TOTAL_RECORDS * ATTACK_RATIO)
n_benign = TOTAL_RECORDS - n_attacks

records = []

for i in range(n_benign):
    r = dict(BENIGN_TEMPLATE)
    r["src_ip"] = gen_benign_ip(i)
    r["label"] = "Benign"
    r["_expected"] = "Benign"
    records.append(r)

for i in range(n_attacks):
    r = dict(ATTACK_TEMPLATE)
    r["src_ip"] = gen_attack_ip(i)
    r["label"] = "DDOS attack-HOIC"
    r["_expected"] = "Attack"
    records.append(r)

random.shuffle(records)  # نخلطهم عشوائياً قبل الإرسال

print(f"Prepared {len(records)} records total: {n_benign} Benign + {n_attacks} Attack ({ATTACK_RATIO*100:.0f}%)")

# ============================================================
# SEND BURST TO KAFKA
# ============================================================
producer = KafkaProducer(
    bootstrap_servers=[KAFKA_BROKER],
    value_serializer=lambda x: json.dumps(x).encode('utf-8')
)

attack_ips_sent = set()
sent_timestamps = {}

print(f"\n[T=0] Sending burst of {len(records)} records to Kafka...")
burst_start_time = time.time()

for r in records:
    r["timestamp"] = datetime.now().isoformat()
    expected = r.pop("_expected")
    if expected == "Attack":
        attack_ips_sent.add(r["src_ip"])
        sent_timestamps[r["src_ip"]] = time.time()
    producer.send('network-logs', value=r)

producer.flush()
burst_send_duration = time.time() - burst_start_time
print(f"Burst sent in {burst_send_duration:.2f} seconds")
print(f"Attack IPs sent ({len(attack_ips_sent)}): {sorted(attack_ips_sent)}")

# ============================================================
# MONITOR POSTGRESQL
# ============================================================
conn = psycopg2.connect(host=DB_HOST, dbname="ids_db",
                         user="ids_user", password="ids_pass123")
cur = conn.cursor()

detected = {}
false_positives = []

print(f"\nMonitoring PostgreSQL for up to {TIMEOUT_SECONDS}s...")
monitor_start = time.time()

while (time.time() - monitor_start) < TIMEOUT_SECONDS:
    # تحقق من التنبيهات المرتبطة بـ IPs الهجوم
    for ip in attack_ips_sent:
        if ip in detected:
            continue
        cur.execute("""
            SELECT attack_type, created_at FROM alerts
            WHERE src_ip = %s ORDER BY created_at DESC LIMIT 1
        """, (ip,))
        result = cur.fetchone()
        if result:
            latency = time.time() - sent_timestamps[ip]
            detected[ip] = {"attack_type": result[0], "latency": latency}

    # تحقق من false positives (IPs طبيعية ظهرت كتنبيه)
    cur.execute("""
        SELECT DISTINCT src_ip, attack_type FROM alerts
        WHERE src_ip LIKE '198.51.100.%'
    """)
    fps = cur.fetchall()
    for ip, atype in fps:
        if ip not in [f["ip"] for f in false_positives]:
            false_positives.append({"ip": ip, "attack_type": atype})

    if len(detected) == len(attack_ips_sent):
        break
    time.sleep(0.5)

cur.close()
conn.close()

# ============================================================
# RESULTS
# ============================================================
print("\n" + "="*60)
print("RESULTS")
print("="*60)
print(f"Total records sent: {len(records)} ({n_benign} Benign + {n_attacks} Attack)")
print(f"Attacks detected: {len(detected)} / {len(attack_ips_sent)}")
print(f"Detection Rate: {len(detected)/len(attack_ips_sent)*100:.1f}%")

if detected:
    latencies = [d["latency"] for d in detected.values()]
    print(f"\nDetection Latency (seconds):")
    print(f"  Min:    {min(latencies):.2f}")
    print(f"  Max:    {max(latencies):.2f}")
    print(f"  Avg:    {sum(latencies)/len(latencies):.2f}")

print(f"\nFalse Positives (Benign flagged as Attack): {len(false_positives)}")
for fp in false_positives:
    print(f"  {fp['ip']} -> {fp['attack_type']}")

undetected = attack_ips_sent - set(detected.keys())
if undetected:
    print(f"\nUndetected attacks (missed): {len(undetected)}")
    for ip in undetected:
        print(f"  {ip}")
