wifi-gps-mapper-public/cracklab/client/result_package.py

from __future__ import annotations

import hashlib
import json
import os
import tempfile
import zipfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


class ResultPackageError(RuntimeError):
    pass


def utc_now() -> str:
    return (
        datetime.now(timezone.utc)
        .replace(microsecond=0)
        .isoformat()
        .replace("+00:00", "Z")
    )


def _sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def _sha256_file(path: Path) -> str:
    digest = hashlib.sha256()

    with path.open("rb") as handle:
        while True:
            chunk = handle.read(1024 * 1024)

            if not chunk:
                break

            digest.update(chunk)

    return digest.hexdigest()


def _json_bytes(data: dict[str, Any]) -> bytes:
    return (
        json.dumps(
            data,
            ensure_ascii=False,
            indent=2,
        )
        + "\n"
    ).encode("utf-8")


def _atomic_write(
    path: Path,
    data: bytes,
) -> None:
    path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    fd, temporary = tempfile.mkstemp(
        prefix=f".{path.name}.",
        suffix=".tmp",
        dir=str(path.parent),
    )

    temporary_path = Path(temporary)

    try:
        with os.fdopen(fd, "wb") as handle:
            handle.write(data)
            handle.flush()
            os.fsync(handle.fileno())

        temporary_path.replace(path)
    except BaseException:
        temporary_path.unlink(
            missing_ok=True
        )
        raise


def _validate_result_steps(
    job: dict[str, Any],
    step_states: list[dict[str, Any]],
) -> None:
    job_steps = job.get("steps")

    if not isinstance(job_steps, list):
        raise ResultPackageError(
            "Job has no valid steps list."
        )

    if len(job_steps) != len(step_states):
        raise ResultPackageError(
            "Result step count does not match Job step count."
        )

    for job_step, result_step in zip(
        job_steps,
        step_states,
    ):
        if (
            result_step.get("step_no")
            != job_step.get("step_no")
        ):
            raise ResultPackageError(
                "Result step_no does not match Job."
            )

        if (
            result_step.get("step_id")
            != job_step.get("step_id")
        ):
            raise ResultPackageError(
                "Result step_id does not match Job."
            )

        if (
            result_step.get("session_name")
            != job_step.get("session_name")
        ):
            raise ResultPackageError(
                "Result session_name does not match Job."
            )

        if result_step.get("status") != "COMPLETED":
            raise ResultPackageError(
                f"Step {job_step.get('step_no')} "
                "is not COMPLETED."
            )

        if result_step.get("exit_code") not in (0, 1):
            raise ResultPackageError(
                f"Step {job_step.get('step_no')} "
                "has unsupported exit_code."
            )

        if result_step.get("restore_seen"):
            raise ResultPackageError(
                f"Step {job_step.get('step_no')} "
                "still has an active restore state."
            )


def build_result(
    job: dict[str, Any],
    step_states: list[dict[str, Any]],
    *,
    completed_at: str | None = None,
) -> tuple[dict[str, Any], bytes]:
    job_id = job.get("job_id")

    if not isinstance(job_id, str) or not job_id:
        raise ResultPackageError(
            "Job has no valid job_id."
        )

    method_id = job.get("method_id")
    method_version = job.get("method_version")

    if not isinstance(method_id, str) or not method_id:
        raise ResultPackageError(
            "Job has no valid method_id."
        )

    if (
        isinstance(method_version, bool)
        or not isinstance(method_version, int)
    ):
        raise ResultPackageError(
            "Job has no valid method_version."
        )

    _validate_result_steps(
        job,
        step_states,
    )

    result = {
        "schema_version": 1,
        "job_id": job_id,
        "status": "COMPLETED",
        "completed_at": completed_at or utc_now(),
        "hash_file_sha256": job["hash_file_sha256"],
        "hash_count": job["hash_count"],
        "method": {
            "id": method_id,
            "version": method_version,
        },
        "steps": step_states,
        "show_file": "hashcat-show.txt",
    }

    unsigned_bytes = _json_bytes(result)

    result["result_sha256"] = _sha256_bytes(
        unsigned_bytes
    )

    return result, _json_bytes(result)


def create_result_package(
    job_dir: Path,
    job: dict[str, Any],
    step_states: list[dict[str, Any]],
    outbox: Path,
    *,
    logger=None,
) -> Path:
    results_dir = job_dir / "results"

    show_path = results_dir / "hashcat-show.txt"

    if not show_path.is_file():
        raise ResultPackageError(
            f"Missing Hashcat show output: {show_path}"
        )

    result, result_bytes = build_result(
        job,
        step_states,
    )

    result_path = results_dir / "result.json"

    _atomic_write(
        result_path,
        result_bytes,
    )

    result_sha256 = _sha256_file(
        result_path
    )

    if result_sha256 != _sha256_bytes(
        result_bytes
    ):
        raise ResultPackageError(
            "Written result.json SHA-256 mismatch."
        )

    outbox.mkdir(
        parents=True,
        exist_ok=True,
    )

    job_id = job["job_id"]

    destination = (
        outbox / f"RESULT-{job_dir.name}.zip"
    )

    if destination.exists():
        raise ResultPackageError(
            f"Result package already exists: "
            f"{destination}"
        )

    fd, temporary = tempfile.mkstemp(
        prefix=f".RESULT-{job_dir.name}.",
        suffix=".zip.tmp",
        dir=str(outbox),
    )

    temporary_path = Path(temporary)

    try:
        os.close(fd)

        with zipfile.ZipFile(
            temporary_path,
            mode="w",
            compression=zipfile.ZIP_DEFLATED,
        ) as archive:
            archive.write(
                result_path,
                "result.json",
            )
            archive.write(
                show_path,
                "hashcat-show.txt",
            )

        temporary_path.replace(
            destination
        )

    except BaseException:
        temporary_path.unlink(
            missing_ok=True
        )
        raise

    if logger is not None:
        logger.info(
            "Result package created: %s "
            "result_sha256=%s "
            "show_sha256=%s",
            destination.name,
            result["result_sha256"],
            _sha256_file(show_path),
        )

    return destination