scripts/enrich_attack_cases.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path

from prompt_injection_framework.io.data_loader import read_jsonl_records, validate_core_fields
from prompt_injection_framework.mapping.map_cases import map_case


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Enrich normalized attack cases with taxonomy metadata")
    parser.add_argument("--input", type=Path, default=Path("data/normalized/attack_cases.jsonl"))
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("data/normalized/attack_cases.enriched.jsonl"),
    )
    parser.add_argument("--limit", type=int, default=0)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    args.output.parent.mkdir(parents=True, exist_ok=True)

    total = 0
    with args.output.open("w", encoding="utf-8") as output_handle:
        for case in read_jsonl_records(args.input):
            validate_core_fields(case)
            output_handle.write(json.dumps(map_case(case), ensure_ascii=False) + "\n")
            total += 1
            if args.limit and total >= args.limit:
                break

    print(f"Wrote {total} enriched records to {args.output}")


if __name__ == "__main__":
    main()