"""
Traffic Dripper -- simulates a live, continuous traffic feed from a
static CICFlowMeter CSV (captured from a pre-made attack PCAP).

Reads the full source CSV once, then periodically writes small chunks
of rows into /output as new, uniquely-named files -- each with the
timestamp column overwritten to "now" so downstream Kafka/HDFS/
PostgreSQL joins use realistic, current timestamps instead of the
original PCAP capture time.

Filebeat is already configured to watch /output/*.csv and pick up new
files automatically -- no change needed there.
"""

import csv
import os
import time
import shutil
from datetime import datetime

SOURCE_DIR   = "/home/ubuntu/ids-project/source_attacks"   # raw CICFlowMeter CSVs go here
OUTPUT_DIR   = "/home/ubuntu/ids-project/output"            # Filebeat watches this
ROWS_PER_DRIP = 30          # how many flow-rows per drip file
DRIP_INTERVAL = 10          # seconds between drips
TIMESTAMP_COL_INDEX = 5     # same index used in the JS processor (p[5])
LOOP_WHEN_DONE = True       # replay source files continuously for demo purposes

os.makedirs(OUTPUT_DIR, exist_ok=True)

def get_source_files():
    return sorted([
        f for f in os.listdir(SOURCE_DIR)
        if f.endswith(".csv")
    ])

def drip_file(source_path, header, rows, batch_num):
    """Write a chunk of rows to a new timestamped CSV in OUTPUT_DIR,
    with the timestamp column overwritten to the current time."""
    now_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    out_rows = []
    for row in rows:
        row = row.copy()
        if len(row) > TIMESTAMP_COL_INDEX:
            row[TIMESTAMP_COL_INDEX] = now_str
        out_rows.append(row)

    fname = f"drip_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}.csv"
    out_path = os.path.join(OUTPUT_DIR, fname)

    with open(out_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(header)
        writer.writerows(out_rows)

    print(f"[{now_str}] Dripped {len(out_rows)} rows -> {fname}")

def process_source_file(path):
    with open(path, "r") as f:
        reader = csv.reader(f)
        header = next(reader)
        rows = list(reader)

    print(f"Loaded {len(rows)} rows from {os.path.basename(path)}")

    batch_num = 0
    for i in range(0, len(rows), ROWS_PER_DRIP):
        chunk = rows[i:i + ROWS_PER_DRIP]
        drip_file(path, header, chunk, batch_num)
        batch_num += 1
        time.sleep(DRIP_INTERVAL)

def main():
    print("Traffic Dripper started.")
    while True:
        source_files = get_source_files()
        if not source_files:
            print(f"No CSV files found in {SOURCE_DIR}. Waiting...")
            time.sleep(30)
            continue

        for fname in source_files:
            full_path = os.path.join(SOURCE_DIR, fname)
            process_source_file(full_path)

        if not LOOP_WHEN_DONE:
            print("All source files processed. Exiting (LOOP_WHEN_DONE=False).")
            break

        print("Finished one full pass over all source files. Looping...")

if __name__ == "__main__":
    main()
