collect_data.py

import csv
import time
from datetime import datetime
from pathlib import Path
from typing import Any

import requests


ATTACKER_URL = "http://192.168.193.133:9000"
WATCHER_URL = "http://192.168.193.133:9001"

NUM_ROUNDS = 10000
ACTION_SLEEP_SECONDS = 30
EXTRA_WAIT_SECONDS = 5

OUTPUT_DIR = Path("./data/emulation")


def get_json(base_url: str, path: str) -> dict[str, Any]:
    response = requests.get(f"{base_url}{path}", timeout=10)
    response.raise_for_status()
    return response.json()


def post_json(
    base_url: str,
    path: str,
    payload: dict[str, Any] | None = None,
) -> dict[str, Any]:
    response = requests.post(
        f"{base_url}{path}",
        json=payload or {},
        timeout=10,
    )

    if response.status_code >= 400:
        print("error:", response.status_code, response.text)

    response.raise_for_status()
    return response.json()


def flatten_stats(prefix: str, data: dict[str, Any]) -> dict[str, Any]:
    flat: dict[str, Any] = {}

    def walk(path: str, value: Any) -> None:
        if isinstance(value, dict):
            for key, child in value.items():
                next_path = f"{path}_{key}" if path else str(key)
                walk(next_path, child)
        else:
            flat[f"{prefix}_{path}"] = value

    walk("", data)
    return flat


def attacker_is_done(status: dict[str, Any]) -> bool:
    return status.get("status") not in {"running", "aborting"}


def stop_attacker_if_needed() -> dict[str, Any]:
    status = get_json(ATTACKER_URL, "/api/status")

    if not attacker_is_done(status):
        return post_json(ATTACKER_URL, "/api/stop")

    return status


def flatten_action(action_data: dict[str, Any]) -> str:
    action_name = action_data.get("action", "unknown")

    payload = action_data.get("payload", {})
    if not isinstance(payload, dict):
        payload = {}

    host_id = payload.get("host_id", "")
    plc_id = payload.get("plc_id", "")
    process_id = payload.get("process_id", "")

    if action_name == "scan_network":
        host_id = "all"
        plc_id = ""
    elif action_name in {
        "idle",
        "scan_host",
        "exploit_host",
        "inspect_host",
    }:
        plc_id = ""
    elif action_name == "tamper_process":
        if isinstance(process_id, str) and process_id.startswith("process-"):
            plc_id = process_id
        elif isinstance(plc_id, str) and plc_id.isdigit():
            plc_id = f"process-{plc_id}"

    return f"{action_name}:{host_id}:{plc_id}"


def build_action_columns(
    action_response: dict[str, Any],
    attacker_status: dict[str, Any],
) -> dict[str, Any]:
    latest_action = attacker_status.get("latest_action") or action_response

    if not isinstance(latest_action, dict):
        latest_action = {}

    latest_payload = latest_action.get("payload", {})
    if not isinstance(latest_payload, dict):
        latest_payload = {}

    effective_stealth_level = latest_payload.get("stealth_level")
    mode_stealth_level = attacker_status.get("stealth_level")

    if effective_stealth_level is None:
        effective_stealth_level = mode_stealth_level

    flattened_action = flatten_action(latest_action)

    return {
        "action": flattened_action,
        "attacker_status": (
            "completed"
            if flattened_action == "idle::"
            else attacker_status.get("status")
        ),
        "stealth_level_effective": effective_stealth_level,
        "stealth_level_mode": mode_stealth_level,
        "action_duration": attacker_status.get("duration_seconds"),
    }


def wait_for_attacker_after_initial_sleep() -> dict[str, Any]:
    attacker_status = get_json(ATTACKER_URL, "/api/status")

    if attacker_is_done(attacker_status):
        return attacker_status

    print("action still running; waiting another 5 seconds")
    time.sleep(EXTRA_WAIT_SECONDS)

    attacker_status = get_json(ATTACKER_URL, "/api/status")

    if attacker_is_done(attacker_status):
        return attacker_status

    print("action still not complete; aborting")
    return post_json(ATTACKER_URL, "/api/stop")


def collect_round(round_id: int) -> dict[str, Any]:
    print(f"\n== ROUND {round_id} ==")

    stop_attacker_if_needed()

    watcher_start = post_json(WATCHER_URL, "/start-collecting")
    print("watcher start:", watcher_start)

    action_response = post_json(ATTACKER_URL, "/api/act", {})
    print("action:", action_response)

    time.sleep(ACTION_SLEEP_SECONDS)

    attacker_status = wait_for_attacker_after_initial_sleep()
    print("attacker status:", attacker_status.get("status"))

    watcher_stop = post_json(WATCHER_URL, "/stop-collecting")
    observation = watcher_stop.get("observation", {})

    main_statistics = observation.get("main_statistics", {})
    extra_statistics = observation.get("extra_statistics", {})

    row: dict[str, Any] = {
        "round_id": round_id,
        "duration": watcher_stop.get("duration_s"),
    }

    row.update(
        build_action_columns(
            action_response=action_response,
            attacker_status=attacker_status,
        )
    )

    row.update(flatten_stats("main", main_statistics))
    row.update(flatten_stats("extra", extra_statistics))

    return row


def append_row(path: Path, row: dict[str, Any], header_written: bool) -> bool:
    columns = list(row.keys())

    with path.open("a", newline="") as file:
        writer = csv.DictWriter(file, fieldnames=columns)

        if not header_written:
            writer.writeheader()
            header_written = True

            print("\n=== CSV COLUMNS ===")
            for column in columns:
                print(column)

        writer.writerow(row)
        file.flush()

    return header_written


def main() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    started_at = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_path = OUTPUT_DIR / f"watcher_collection_{started_at}.csv"

    print("attacker health:", get_json(ATTACKER_URL, "/api/health"))
    print("watcher status:", get_json(WATCHER_URL, "/status"))
    print(f"output file: {output_path}")

    header_written = False

    try:
        for round_id in range(1, NUM_ROUNDS + 1):
            row = collect_round(round_id)

            header_written = append_row(
                path=output_path,
                row=row,
                header_written=header_written,
            )

            print(f"saved round {round_id}: {output_path}")

    except KeyboardInterrupt:
        print("\ncollection interrupted by user")

    finally:
        stop_attacker_if_needed()
        print(f"data saved to: {output_path}")


if __name__ == "__main__":
    main()