Update public source
This commit is contained in:
@@ -0,0 +1,310 @@
|
||||
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
|
||||
Reference in New Issue
Block a user