453 lines
11 KiB
Python
453 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import sqlite3
|
|
from datetime import datetime, timezone
|
|
from typing import Any
|
|
|
|
|
|
class JobsRepository:
|
|
def __init__(self, conn: sqlite3.Connection) -> None:
|
|
self._conn = conn
|
|
|
|
@staticmethod
|
|
def _now() -> str:
|
|
return datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace(
|
|
"+00:00",
|
|
"Z",
|
|
)
|
|
|
|
def create_job(
|
|
self,
|
|
job_id: str,
|
|
method_id: str,
|
|
method_version: int,
|
|
status: str,
|
|
hash_file_name: str,
|
|
hash_file_sha256: str,
|
|
hash_count: int,
|
|
step_count: int,
|
|
client_id: str | None = None,
|
|
) -> str:
|
|
created_at = self._now()
|
|
queued_at = created_at if status == "QUEUED" else None
|
|
|
|
self._conn.execute(
|
|
"""
|
|
INSERT INTO crack_jobs (
|
|
id,
|
|
method_id,
|
|
method_version,
|
|
status,
|
|
created_at,
|
|
queued_at,
|
|
started_at,
|
|
completed_at,
|
|
cancelled_at,
|
|
imported_at,
|
|
client_id,
|
|
hash_file_name,
|
|
hash_file_sha256,
|
|
hash_count,
|
|
step_count,
|
|
result_sha256,
|
|
error
|
|
)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
job_id,
|
|
method_id,
|
|
method_version,
|
|
status,
|
|
created_at,
|
|
queued_at,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
client_id,
|
|
hash_file_name,
|
|
hash_file_sha256,
|
|
hash_count,
|
|
step_count,
|
|
None,
|
|
None,
|
|
),
|
|
)
|
|
|
|
return created_at
|
|
|
|
def get_job(self, job_id: str) -> sqlite3.Row | None:
|
|
return self._conn.execute(
|
|
"""
|
|
SELECT
|
|
id,
|
|
method_id,
|
|
method_version,
|
|
status,
|
|
created_at,
|
|
queued_at,
|
|
started_at,
|
|
completed_at,
|
|
cancelled_at,
|
|
imported_at,
|
|
client_id,
|
|
hash_file_name,
|
|
hash_file_sha256,
|
|
hash_count,
|
|
step_count,
|
|
result_sha256,
|
|
error
|
|
FROM crack_jobs
|
|
WHERE id = ?
|
|
""",
|
|
(job_id,),
|
|
).fetchone()
|
|
|
|
def list_jobs(
|
|
self,
|
|
status: str | None = None,
|
|
method_id: str | None = None,
|
|
client_id: str | None = None,
|
|
) -> list[sqlite3.Row]:
|
|
query = """
|
|
SELECT
|
|
id,
|
|
method_id,
|
|
method_version,
|
|
status,
|
|
created_at,
|
|
queued_at,
|
|
started_at,
|
|
completed_at,
|
|
cancelled_at,
|
|
imported_at,
|
|
client_id,
|
|
hash_file_name,
|
|
hash_file_sha256,
|
|
hash_count,
|
|
step_count,
|
|
result_sha256,
|
|
error
|
|
FROM crack_jobs
|
|
WHERE 1 = 1
|
|
"""
|
|
params: list[Any] = []
|
|
|
|
if status is not None:
|
|
query += "\n AND status = ?"
|
|
params.append(status)
|
|
|
|
if method_id is not None:
|
|
query += "\n AND method_id = ?"
|
|
params.append(method_id)
|
|
|
|
if client_id is not None:
|
|
query += "\n AND client_id = ?"
|
|
params.append(client_id)
|
|
|
|
query += "\n ORDER BY created_at DESC, id DESC"
|
|
|
|
return self._conn.execute(query, params).fetchall()
|
|
|
|
def update_job_status(
|
|
self,
|
|
job_id: str,
|
|
status: str,
|
|
) -> None:
|
|
now = self._now()
|
|
|
|
timestamp_column = {
|
|
"QUEUED": "queued_at",
|
|
"RUNNING": "started_at",
|
|
"COMPLETED": "completed_at",
|
|
"CANCELLED": "cancelled_at",
|
|
}.get(status)
|
|
|
|
if timestamp_column is None:
|
|
self._conn.execute(
|
|
"""
|
|
UPDATE crack_jobs
|
|
SET status = ?
|
|
WHERE id = ?
|
|
""",
|
|
(
|
|
status,
|
|
job_id,
|
|
),
|
|
)
|
|
return
|
|
|
|
self._conn.execute(
|
|
f"""
|
|
UPDATE crack_jobs
|
|
SET status = ?,
|
|
{timestamp_column} = ?
|
|
WHERE id = ?
|
|
""",
|
|
(
|
|
status,
|
|
now,
|
|
job_id,
|
|
),
|
|
)
|
|
|
|
def assign_client(
|
|
self,
|
|
job_id: str,
|
|
client_id: str | None,
|
|
) -> None:
|
|
self._conn.execute(
|
|
"""
|
|
UPDATE crack_jobs
|
|
SET client_id = ?
|
|
WHERE id = ?
|
|
""",
|
|
(
|
|
client_id,
|
|
job_id,
|
|
),
|
|
)
|
|
|
|
def update_job_result(
|
|
self,
|
|
job_id: str,
|
|
result_sha256: str | None = None,
|
|
error: str | None = None,
|
|
imported_at: str | None = None,
|
|
) -> None:
|
|
self._conn.execute(
|
|
"""
|
|
UPDATE crack_jobs
|
|
SET result_sha256 = ?,
|
|
error = ?,
|
|
imported_at = ?
|
|
WHERE id = ?
|
|
""",
|
|
(
|
|
result_sha256,
|
|
error,
|
|
imported_at,
|
|
job_id,
|
|
),
|
|
)
|
|
|
|
def add_job_handshake(
|
|
self,
|
|
job_id: str,
|
|
handshake_id: int,
|
|
hash22000: str,
|
|
access_point_id: int | None = None,
|
|
) -> None:
|
|
self._conn.execute(
|
|
"""
|
|
INSERT INTO crack_job_handshakes (
|
|
job_id,
|
|
handshake_id,
|
|
hash22000,
|
|
access_point_id
|
|
)
|
|
VALUES (?, ?, ?, ?)
|
|
""",
|
|
(
|
|
job_id,
|
|
handshake_id,
|
|
hash22000,
|
|
access_point_id,
|
|
),
|
|
)
|
|
|
|
def get_job_handshake(
|
|
self,
|
|
job_id: str,
|
|
handshake_id: int,
|
|
) -> sqlite3.Row | None:
|
|
return self._conn.execute(
|
|
"""
|
|
SELECT
|
|
job_id,
|
|
handshake_id,
|
|
hash22000,
|
|
access_point_id
|
|
FROM crack_job_handshakes
|
|
WHERE job_id = ?
|
|
AND handshake_id = ?
|
|
""",
|
|
(
|
|
job_id,
|
|
handshake_id,
|
|
),
|
|
).fetchone()
|
|
|
|
def list_job_handshakes(
|
|
self,
|
|
job_id: str,
|
|
) -> list[sqlite3.Row]:
|
|
return self._conn.execute(
|
|
"""
|
|
SELECT
|
|
job_id,
|
|
handshake_id,
|
|
hash22000,
|
|
access_point_id
|
|
FROM crack_job_handshakes
|
|
WHERE job_id = ?
|
|
ORDER BY handshake_id
|
|
""",
|
|
(job_id,),
|
|
).fetchall()
|
|
|
|
def create_job_step(
|
|
self,
|
|
job_id: str,
|
|
step_no: int,
|
|
step_id: str,
|
|
status: str,
|
|
session_name: str,
|
|
definition: dict[str, Any],
|
|
) -> None:
|
|
self._conn.execute(
|
|
"""
|
|
INSERT INTO crack_job_steps (
|
|
job_id,
|
|
step_no,
|
|
step_id,
|
|
status,
|
|
session_name,
|
|
definition_json,
|
|
started_at,
|
|
completed_at,
|
|
exit_code,
|
|
restore_seen,
|
|
error
|
|
)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
job_id,
|
|
step_no,
|
|
step_id,
|
|
status,
|
|
session_name,
|
|
json.dumps(
|
|
definition,
|
|
ensure_ascii=False,
|
|
separators=(",", ":"),
|
|
sort_keys=True,
|
|
),
|
|
None,
|
|
None,
|
|
None,
|
|
0,
|
|
None,
|
|
),
|
|
)
|
|
|
|
def get_job_step(
|
|
self,
|
|
job_id: str,
|
|
step_no: int,
|
|
) -> sqlite3.Row | None:
|
|
return self._conn.execute(
|
|
"""
|
|
SELECT
|
|
job_id,
|
|
step_no,
|
|
step_id,
|
|
status,
|
|
session_name,
|
|
definition_json,
|
|
started_at,
|
|
completed_at,
|
|
exit_code,
|
|
restore_seen,
|
|
error
|
|
FROM crack_job_steps
|
|
WHERE job_id = ?
|
|
AND step_no = ?
|
|
""",
|
|
(
|
|
job_id,
|
|
step_no,
|
|
),
|
|
).fetchone()
|
|
|
|
def list_job_steps(
|
|
self,
|
|
job_id: str,
|
|
) -> list[sqlite3.Row]:
|
|
return self._conn.execute(
|
|
"""
|
|
SELECT
|
|
job_id,
|
|
step_no,
|
|
step_id,
|
|
status,
|
|
session_name,
|
|
definition_json,
|
|
started_at,
|
|
completed_at,
|
|
exit_code,
|
|
restore_seen,
|
|
error
|
|
FROM crack_job_steps
|
|
WHERE job_id = ?
|
|
ORDER BY step_no
|
|
""",
|
|
(job_id,),
|
|
).fetchall()
|
|
|
|
def update_job_step(
|
|
self,
|
|
job_id: str,
|
|
step_no: int,
|
|
status: str | None = None,
|
|
started_at: str | None = None,
|
|
completed_at: str | None = None,
|
|
exit_code: int | None = None,
|
|
restore_seen: bool | None = None,
|
|
error: str | None = None,
|
|
) -> None:
|
|
fields: list[str] = []
|
|
params: list[Any] = []
|
|
|
|
if status is not None:
|
|
fields.append("status = ?")
|
|
params.append(status)
|
|
|
|
if started_at is not None:
|
|
fields.append("started_at = ?")
|
|
params.append(started_at)
|
|
|
|
if completed_at is not None:
|
|
fields.append("completed_at = ?")
|
|
params.append(completed_at)
|
|
|
|
if exit_code is not None:
|
|
fields.append("exit_code = ?")
|
|
params.append(exit_code)
|
|
|
|
if restore_seen is not None:
|
|
fields.append("restore_seen = ?")
|
|
params.append(int(restore_seen))
|
|
|
|
if error is not None:
|
|
fields.append("error = ?")
|
|
params.append(error)
|
|
|
|
if not fields:
|
|
return
|
|
|
|
params.extend((job_id, step_no))
|
|
|
|
self._conn.execute(
|
|
f"""
|
|
UPDATE crack_job_steps
|
|
SET {", ".join(fields)}
|
|
WHERE job_id = ?
|
|
AND step_no = ?
|
|
""",
|
|
params,
|
|
)
|