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

import json
import re
import sys
import zipfile
from pathlib import Path
from typing import Any


USER_CODE_RE = re.compile(r"\b[A-Z0-9]{4}-[A-Z0-9]{4}\b")


def read_text(zf: zipfile.ZipFile, name: str) -> str:
    return zf.read(name).decode("utf-8")


def read_jsonl(zf: zipfile.ZipFile, name: str) -> list[dict[str, Any]]:
    rows = []
    for line in read_text(zf, name).splitlines():
        line = line.strip()
        if line:
            rows.append(json.loads(line))
    return rows


def normalize_code(value: str) -> str:
    return re.sub(r"[^A-Za-z0-9]", "", value).lower()


def nested_modified_property(row: dict[str, Any], display_name: str) -> str | None:
    for resource in row.get("targetResources", []):
        for prop in resource.get("modifiedProperties", []):
            if prop.get("displayName") == display_name:
                return prop.get("newValue")
    return None


def recover_components(bundle_path: str | Path) -> dict[str, str]:
    with zipfile.ZipFile(bundle_path) as zf:
        transcript = read_text(zf, "endpoint/powershell_transcript.txt")
        code_match = USER_CODE_RE.search(transcript)
        if not code_match:
            raise ValueError("no device user code found in endpoint transcript")
        user_code = code_match.group(0)

        defender_rows = read_jsonl(zf, "endpoint/defender_device_events.jsonl")
        if not any(row.get("AdditionalFields", {}).get("UserCodeObserved") == user_code for row in defender_rows):
            raise ValueError("endpoint telemetry does not confirm the same user code")

        signin_rows = read_jsonl(zf, "entra/signin_logs.jsonl")
        signin = next(
            row
            for row in signin_rows
            if row.get("authenticationProtocol") == "DeviceCode"
            and row.get("userCodeEvidence") == user_code
            and row.get("status", {}).get("errorCode") == 0
        )
        correlation_id = signin["correlationId"]
        flow = signin.get("flow_key", "dc")

        audit_rows = read_jsonl(zf, "entra/audit_logs.jsonl")
        consent = next(
            row
            for row in audit_rows
            if row.get("activityDisplayName") == "Consent to application"
            and row.get("correlationId") == correlation_id
            and row.get("result") == "success"
        )
        grant_id = nested_modified_property(consent, "OAuth2PermissionGrant.Id")
        if not grant_id:
            raise ValueError("consent event does not include OAuth2PermissionGrant.Id")
        grant_tail = grant_id.rsplit("-", 1)[-1]

        m365_rows = read_jsonl(zf, "m365/unified_audit_log.jsonl")
        export = next(
            row
            for row in m365_rows
            if row.get("Operation") == "GraphExportJobCreated"
            and row.get("CorrelationId") == correlation_id
            and row.get("OAuth2PermissionGrantId") == grant_id
            and row.get("ResultStatus") == "Succeeded"
        )
        export_tail = export["ExportId"].rsplit("-", 1)[-1]

    return {
        "flow": flow,
        "user_code": normalize_code(user_code),
        "grant_tail": grant_tail,
        "export_tail": export_tail,
    }


def solve_bundle(bundle_path: str | Path) -> str:
    components = recover_components(bundle_path)
    return "caplag{" + "_".join(
        [
            components["flow"],
            components["user_code"],
            components["grant_tail"],
            components["export_tail"],
        ]
    ) + "}"


def main(argv: list[str]) -> int:
    if len(argv) != 2:
        print(f"usage: {Path(argv[0]).name} public/grant_residue_case.zip", file=sys.stderr)
        return 2
    print(solve_bundle(argv[1]))
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))