#!/usr/bin/env python3
"""Parse Volatility 3 JSON output to highlight likely process-injection candidates.

Volatility 3 can render plugin output as JSON (``vol -r json ...``). This tool consumes a
``windows.malfind`` JSON document (from a file or stdin) and ranks regions by injection
signal: RWX protection and the presence of an MZ header or executable prologue bytes.

It analyzes plugin output only; it does not parse raw memory itself.
"""
from __future__ import annotations

import argparse
import json
import sys


def _rows(doc) -> list[dict]:
    # Volatility 3 JSON renderer emits a list of row dicts; tolerate a wrapped form.
    if isinstance(doc, list):
        return doc
    if isinstance(doc, dict):
        for key in ("rows", "data"):
            if isinstance(doc.get(key), list):
                return doc[key]
    return []


def rank_malfind(doc) -> list[dict]:
    out = []
    for row in _rows(doc):
        prot = str(row.get("Protection", row.get("protection", ""))).upper()
        disasm = str(row.get("Disasm", row.get("disasm", "")))
        hexdump = str(row.get("Hexdump", row.get("hexdump", "")))
        score = 0
        if "EXECUTE_READWRITE" in prot or "RWX" in prot:
            score += 3
        if "4D 5A" in hexdump or hexdump.lstrip().startswith("MZ"):
            score += 3  # PE header in an injected region
        if any(op in disasm for op in ("push", "mov", "call", "jmp")):
            score += 1
        out.append({
            "pid": row.get("PID", row.get("pid")),
            "process": row.get("Process", row.get("process")),
            "address": row.get("Start VPN", row.get("Address", row.get("address"))),
            "protection": prot,
            "injection_score": score,
        })
    return sorted(out, key=lambda r: r["injection_score"], reverse=True)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Rank Volatility malfind output")
    sub = parser.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("malfind")
    p.add_argument("json", help="path to malfind JSON, or - for stdin")
    args = parser.parse_args(argv)

    text = sys.stdin.read() if args.json == "-" else open(args.json, encoding="utf-8").read()
    doc = json.loads(text)
    print(json.dumps(rank_malfind(doc), indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())