#!/usr/bin/env bash
set -euo pipefail

if [ "$#" -lt 2 ]; then
    echo "Usage: $0 <name> doca-*-vuln.csv [more.csv ...]" >&2
    echo "Example: $0 doca doca-*-vuln.csv" >&2
    exit 1
fi

NAME="$1"
shift

python3 - "$NAME" "$@" <<'PY'
import csv
import os
import re
import sys

FILENAME_RE = re.compile(
    r"^doca-(?P<container>.+)-(?P<version>\d+\.\d+\.\d+)-(?P<host>host|dpu)-vuln\.csv$"
)

META_COLUMNS = ["doca_version", "doca_host", "doca_container"]

OUTPUT_BUCKETS = [
    ("network", "nopriv"),
    ("network", "priv"),
    ("local", "nopriv"),
    ("local", "priv"),
]


def parse_metadata(path):
    base = os.path.basename(path)
    match = FILENAME_RE.match(base)

    if not match:
        raise ValueError(
            f"Filename does not match expected pattern: {base}\n"
            f"Expected: doca-(container)-(version)-(host)-vuln.csv"
        )

    return {
        "doca_version": match.group("version"),
        "doca_host": match.group("host"),
        "doca_container": match.group("container"),
    }


def get_av(row):
    return (
        row.get("nvd_AV")
        or row.get("AV")
        or row.get("cvss_AV")
        or ""
    ).strip().upper()


def get_pr(row):
    return (
        row.get("nvd_PR")
        or row.get("PR")
        or row.get("cvss_PR")
        or ""
    ).strip().upper()


def av_bucket(av):
    # User requested N and A in the same file.
    if av in {"N", "A"}:
        return "network"

    return "local"


def pr_bucket(pr):
    if pr == "N":
        return "nopriv"

    # L, H, empty, unknown all go to privileged bucket.
    return "priv"


def main(name, paths):
    output_files = {
        bucket: open(f"{name}-{bucket[0]}-{bucket[1]}.csv", "w", encoding="utf-8", newline="")
        for bucket in OUTPUT_BUCKETS
    }

    writers = {}
    output_header = None

    try:
        for path in paths:
            metadata = parse_metadata(path)

            with open(path, "r", encoding="utf-8", newline="") as f:
                reader = csv.DictReader(f)

                if reader.fieldnames is None:
                    continue

                if output_header is None:
                    original_header = reader.fieldnames
                    output_header = META_COLUMNS + original_header

                    for bucket, fh in output_files.items():
                        writer = csv.DictWriter(
                            fh,
                            fieldnames=output_header,
                            extrasaction="ignore",
                        )
                        writer.writeheader()
                        writers[bucket] = writer

                for row in reader:
                    av = get_av(row)
                    pr = get_pr(row)

                    bucket = (av_bucket(av), pr_bucket(pr))

                    output_row = {}
                    output_row.update(metadata)
                    output_row.update(row)

                    writers[bucket].writerow(output_row)

    finally:
        for fh in output_files.values():
            fh.close()


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2:])
PY