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