import os
import time
import subprocess

INPUT_DIR      = "/input"
OUTPUT_DIR     = "/source_attacks"
PROCESSED_LOG  = "/input/.processed.txt"
POLL_INTERVAL  = 5  # seconds

os.makedirs(OUTPUT_DIR, exist_ok=True)
if not os.path.exists(PROCESSED_LOG):
    open(PROCESSED_LOG, "w").close()


def get_processed():
    with open(PROCESSED_LOG) as f:
        return set(line.strip() for line in f)


def mark_processed(name):
    with open(PROCESSED_LOG, "a") as f:
        f.write(name + "\n")


def run_cicflowmeter(pcap_path, csv_path):
    print(f"[CICFlowMeter] Converting {pcap_path} -> {csv_path}", flush=True)
    result = subprocess.run(
        ["cicflowmeter", "-f", pcap_path, "-c", csv_path],
        capture_output=True, text=True
    )
    if result.returncode != 0:
        print(f"[ERROR] cicflowmeter failed for {pcap_path}: {result.stderr}", flush=True)
        return False
    print(f"[OK] Converted -> {csv_path}", flush=True)
    return True


def watch_loop():
    print(f"Watching {INPUT_DIR} for new .pcap files...", flush=True)
    while True:
        processed = get_processed()
        for fname in sorted(os.listdir(INPUT_DIR)):
            if not fname.endswith(".pcap") or fname in processed:
                continue

            pcap_path = os.path.join(INPUT_DIR, fname)
            csv_name  = fname.replace(".pcap", ".csv")
            csv_path  = os.path.join(OUTPUT_DIR, csv_name)

            if run_cicflowmeter(pcap_path, csv_path):
                mark_processed(fname)
            else:
                print(f"[SKIP] Will retry {fname} next cycle", flush=True)

        time.sleep(POLL_INTERVAL)


if __name__ == "__main__":
    watch_loop()
