349 lines
9.5 KiB
Python
349 lines
9.5 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import tempfile
|
|
import zipfile
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from zoneinfo import ZoneInfo
|
|
|
|
import yaml
|
|
|
|
from cracklab.server.repositories.jobs import JobsRepository
|
|
from cracklab.server.repositories.methods import MethodsRepository
|
|
from cracklab.server.services.method_definition import (
|
|
validate_method_definition,
|
|
)
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[3]
|
|
CONFIG_FILE = PROJECT_ROOT / "config" / "config.yaml"
|
|
|
|
|
|
class JobExportError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class JobExportService:
|
|
def __init__(
|
|
self,
|
|
conn,
|
|
timezone_name: str | None = None,
|
|
) -> None:
|
|
self._conn = conn
|
|
self._jobs = JobsRepository(conn)
|
|
self._methods = MethodsRepository(conn)
|
|
self._timezone_name = (
|
|
timezone_name
|
|
if timezone_name is not None
|
|
else self._load_timezone()
|
|
)
|
|
|
|
try:
|
|
self._timezone = ZoneInfo(self._timezone_name)
|
|
except Exception as exc:
|
|
raise JobExportError(
|
|
f"Invalid timezone: {self._timezone_name}"
|
|
) from exc
|
|
|
|
@staticmethod
|
|
def _load_timezone() -> str:
|
|
if not CONFIG_FILE.is_file():
|
|
raise JobExportError(
|
|
f"Configuration file not found: {CONFIG_FILE}"
|
|
)
|
|
|
|
try:
|
|
with CONFIG_FILE.open(
|
|
"r",
|
|
encoding="utf-8",
|
|
) as fp:
|
|
config = yaml.safe_load(fp) or {}
|
|
except Exception as exc:
|
|
raise JobExportError(
|
|
f"Failed to load configuration: {CONFIG_FILE}"
|
|
) from exc
|
|
|
|
timezone_name = config.get(
|
|
"timezone",
|
|
"Europe/Moscow",
|
|
)
|
|
|
|
if not isinstance(timezone_name, str) or not timezone_name:
|
|
raise JobExportError(
|
|
"Configuration timezone must be a non-empty string"
|
|
)
|
|
|
|
return timezone_name
|
|
|
|
@staticmethod
|
|
def _parse_utc(text: str) -> datetime:
|
|
try:
|
|
value = datetime.fromisoformat(text)
|
|
except ValueError as exc:
|
|
raise JobExportError(
|
|
f"Invalid job created_at: {text}"
|
|
) from exc
|
|
|
|
if value.tzinfo is None:
|
|
raise JobExportError(
|
|
f"Job timestamp is not timezone-aware: {text}"
|
|
)
|
|
|
|
value = value.astimezone(timezone.utc)
|
|
|
|
return value
|
|
|
|
@staticmethod
|
|
def _unique_hash_content(rows) -> tuple[str, int]:
|
|
unique_hashes: list[str] = []
|
|
seen: set[str] = set()
|
|
|
|
for row in rows:
|
|
hash22000 = str(row["hash22000"]).strip()
|
|
|
|
if not hash22000:
|
|
raise JobExportError(
|
|
"Job contains an empty hash22000 value"
|
|
)
|
|
|
|
if hash22000 in seen:
|
|
continue
|
|
|
|
seen.add(hash22000)
|
|
unique_hashes.append(hash22000)
|
|
|
|
content = "\n".join(unique_hashes) + "\n"
|
|
|
|
return content, len(unique_hashes)
|
|
|
|
@staticmethod
|
|
def _sha256_text(content: str) -> str:
|
|
return hashlib.sha256(
|
|
content.encode("utf-8")
|
|
).hexdigest()
|
|
|
|
@staticmethod
|
|
def _json_bytes(data: dict) -> bytes:
|
|
return (
|
|
json.dumps(
|
|
data,
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
+ "\n"
|
|
).encode("utf-8")
|
|
|
|
def export_job(
|
|
self,
|
|
job_id: str,
|
|
output_dir: str | Path,
|
|
) -> Path:
|
|
job = self._jobs.get_job(job_id)
|
|
|
|
if job is None:
|
|
raise JobExportError(
|
|
f"Job not found: {job_id}"
|
|
)
|
|
|
|
method_version = self._methods.get_method_version(
|
|
str(job["method_id"]),
|
|
int(job["method_version"]),
|
|
)
|
|
|
|
if method_version is None:
|
|
raise JobExportError(
|
|
"Method version not found: "
|
|
f"{job['method_id']} v{job['method_version']}"
|
|
)
|
|
|
|
try:
|
|
method_definition = json.loads(
|
|
method_version["definition_json"]
|
|
)
|
|
validate_method_definition(
|
|
method_definition
|
|
)
|
|
except (json.JSONDecodeError, TypeError, ValueError) as exc:
|
|
raise JobExportError(
|
|
"Invalid method version definition: "
|
|
f"{job['method_id']} v{job['method_version']}"
|
|
) from exc
|
|
|
|
hash_mode = int(
|
|
method_definition["hash_mode"]
|
|
)
|
|
|
|
handshake_rows = self._jobs.list_job_handshakes(
|
|
job_id
|
|
)
|
|
step_rows = self._jobs.list_job_steps(
|
|
job_id
|
|
)
|
|
|
|
hash_content, hash_count = self._unique_hash_content(
|
|
handshake_rows
|
|
)
|
|
|
|
hash_sha256 = self._sha256_text(
|
|
hash_content
|
|
)
|
|
|
|
if hash_sha256 != job["hash_file_sha256"]:
|
|
raise JobExportError(
|
|
"Job hash snapshot SHA-256 mismatch: "
|
|
f"expected {job['hash_file_sha256']}, "
|
|
f"calculated {hash_sha256}"
|
|
)
|
|
|
|
if hash_count != int(job["hash_count"]):
|
|
raise JobExportError(
|
|
"Job hash count mismatch: "
|
|
f"expected {job['hash_count']}, "
|
|
f"calculated {hash_count}"
|
|
)
|
|
|
|
if len(handshake_rows) == 0:
|
|
raise JobExportError(
|
|
f"Job contains no handshake snapshot: {job_id}"
|
|
)
|
|
|
|
if len(step_rows) != int(job["step_count"]):
|
|
raise JobExportError(
|
|
"Job step count mismatch: "
|
|
f"expected {job['step_count']}, "
|
|
f"calculated {len(step_rows)}"
|
|
)
|
|
|
|
created_at = self._parse_utc(
|
|
str(job["created_at"])
|
|
)
|
|
|
|
local_created_at = created_at.astimezone(
|
|
self._timezone
|
|
)
|
|
|
|
zip_filename = (
|
|
"JOB-"
|
|
+ local_created_at.strftime(
|
|
"%d-%m-%Y_%H%M%S"
|
|
)
|
|
+ "-"
|
|
+ str(job["id"]).removeprefix("JOB-")
|
|
+ ".zip"
|
|
)
|
|
|
|
job_json = {
|
|
"job_id": str(job["id"]),
|
|
"created_at": str(job["created_at"]),
|
|
"method_id": str(job["method_id"]),
|
|
"method_version": int(job["method_version"]),
|
|
"hash_mode": hash_mode,
|
|
"hash_file_name": str(job["hash_file_name"]),
|
|
"hash_file_sha256": str(job["hash_file_sha256"]),
|
|
"hash_count": int(job["hash_count"]),
|
|
"handshake_count": len(handshake_rows),
|
|
"step_count": int(job["step_count"]),
|
|
"handshakes": [
|
|
{
|
|
"handshake_id": int(row["handshake_id"]),
|
|
"hash22000": str(row["hash22000"]),
|
|
"access_point_id": (
|
|
int(row["access_point_id"])
|
|
if row["access_point_id"] is not None
|
|
else None
|
|
),
|
|
}
|
|
for row in handshake_rows
|
|
],
|
|
"steps": [
|
|
{
|
|
"step_no": int(row["step_no"]),
|
|
"step_id": str(row["step_id"]),
|
|
"session_name": str(row["session_name"]),
|
|
"definition": json.loads(
|
|
row["definition_json"]
|
|
),
|
|
}
|
|
for row in step_rows
|
|
],
|
|
}
|
|
|
|
job_bytes = self._json_bytes(job_json)
|
|
hash_bytes = hash_content.encode("utf-8")
|
|
|
|
manifest_files = [
|
|
{
|
|
"path": "job.json",
|
|
"size": len(job_bytes),
|
|
"sha256": hashlib.sha256(
|
|
job_bytes
|
|
).hexdigest(),
|
|
},
|
|
{
|
|
"path": f"hashes/{job['hash_file_name']}",
|
|
"size": len(hash_bytes),
|
|
"sha256": hashlib.sha256(
|
|
hash_bytes
|
|
).hexdigest(),
|
|
},
|
|
]
|
|
|
|
manifest = {
|
|
"format_version": 1,
|
|
"job_id": str(job["id"]),
|
|
"created_at": str(job["created_at"]),
|
|
"zip_filename": zip_filename,
|
|
"files": manifest_files,
|
|
}
|
|
|
|
manifest_bytes = self._json_bytes(
|
|
manifest
|
|
)
|
|
|
|
output_path = Path(output_dir)
|
|
output_path.mkdir(
|
|
parents=True,
|
|
exist_ok=True,
|
|
)
|
|
|
|
final_path = output_path / zip_filename
|
|
|
|
with tempfile.NamedTemporaryFile(
|
|
mode="wb",
|
|
suffix=".zip",
|
|
dir=output_path,
|
|
delete=False,
|
|
) as tmp:
|
|
temporary_path = Path(tmp.name)
|
|
|
|
try:
|
|
with zipfile.ZipFile(
|
|
temporary_path,
|
|
mode="w",
|
|
compression=zipfile.ZIP_DEFLATED,
|
|
) as archive:
|
|
archive.writestr(
|
|
"job.json",
|
|
job_bytes,
|
|
)
|
|
archive.writestr(
|
|
f"hashes/{job['hash_file_name']}",
|
|
hash_bytes,
|
|
)
|
|
archive.writestr(
|
|
"manifest.json",
|
|
manifest_bytes,
|
|
)
|
|
|
|
temporary_path.replace(final_path)
|
|
|
|
except Exception:
|
|
temporary_path.unlink(
|
|
missing_ok=True
|
|
)
|
|
raise
|
|
|
|
return final_path
|