Archived
287 lines
10 KiB
Python
287 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from dataclasses import dataclass
|
|
from datetime import UTC, datetime
|
|
|
|
from ..models import UpdateResult
|
|
from .base import StoreConnectionMixin
|
|
from .types import RecordRow, UpdateEventRow, record_row, update_event_row
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class UpdateEventInput:
|
|
hostname: str
|
|
ipv4: str | None
|
|
ipv6: str | None
|
|
status: str
|
|
detail: str
|
|
source_ip: str
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class RecordUpdateInput:
|
|
hostname: str
|
|
ipv4: str | None
|
|
ipv6: str | None
|
|
update_ipv4: bool = True
|
|
update_ipv6: bool = True
|
|
|
|
|
|
def _update_event_input(event: UpdateEventInput | str, legacy_values: tuple[object, ...]) -> UpdateEventInput:
|
|
if isinstance(event, UpdateEventInput):
|
|
return event
|
|
if len(legacy_values) != 5:
|
|
raise TypeError("legacy update event values must be ipv4, ipv6, status, detail, source_ip")
|
|
ipv4, ipv6, status, detail, source_ip = legacy_values
|
|
return UpdateEventInput(
|
|
hostname=event,
|
|
ipv4=str(ipv4) if ipv4 is not None else None,
|
|
ipv6=str(ipv6) if ipv6 is not None else None,
|
|
status=str(status),
|
|
detail=str(detail),
|
|
source_ip=str(source_ip),
|
|
)
|
|
|
|
|
|
def _planned_record_update(
|
|
update: RecordUpdateInput,
|
|
*,
|
|
current_ipv4: str | None,
|
|
current_ipv6: str | None,
|
|
existing: bool,
|
|
) -> UpdateResult:
|
|
effective_ipv4 = update.ipv4 if update.update_ipv4 else current_ipv4
|
|
effective_ipv6 = update.ipv6 if update.update_ipv6 else current_ipv6
|
|
return UpdateResult(
|
|
hostname=update.hostname,
|
|
ipv4=effective_ipv4,
|
|
ipv6=effective_ipv6,
|
|
changed=not existing or current_ipv4 != effective_ipv4 or current_ipv6 != effective_ipv6,
|
|
update_ipv4=update.update_ipv4,
|
|
update_ipv6=update.update_ipv6,
|
|
)
|
|
|
|
|
|
class RecordsEventsMixin(StoreConnectionMixin):
|
|
def upsert(
|
|
self,
|
|
hostname: str,
|
|
ipv4: str | None,
|
|
ipv6: str | None,
|
|
update_ipv4: bool = True,
|
|
update_ipv6: bool = True,
|
|
) -> UpdateResult:
|
|
now = datetime.now(UTC).isoformat()
|
|
update = RecordUpdateInput(hostname, ipv4, ipv6, update_ipv4, update_ipv6)
|
|
with self.connect() as conn:
|
|
existing = conn.execute("SELECT ipv4, ipv6 FROM records WHERE hostname = ?", (hostname,)).fetchone()
|
|
current_ipv4 = existing["ipv4"] if existing else None
|
|
current_ipv6 = existing["ipv6"] if existing else None
|
|
result = _planned_record_update(
|
|
update,
|
|
current_ipv4=current_ipv4,
|
|
current_ipv6=current_ipv6,
|
|
existing=existing is not None,
|
|
)
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO records(hostname, ipv4, ipv6, updated_at)
|
|
VALUES (?, ?, ?, ?)
|
|
ON CONFLICT(hostname) DO UPDATE SET
|
|
ipv4=excluded.ipv4,
|
|
ipv6=excluded.ipv6,
|
|
updated_at=excluded.updated_at
|
|
""",
|
|
(result.hostname, result.ipv4, result.ipv6, now),
|
|
)
|
|
return result
|
|
|
|
def commit_update_result(self, result: UpdateResult) -> UpdateResult:
|
|
now = datetime.now(UTC).isoformat()
|
|
with self.connect() as conn:
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO records(hostname, ipv4, ipv6, updated_at)
|
|
VALUES (?, ?, ?, ?)
|
|
ON CONFLICT(hostname) DO UPDATE SET
|
|
ipv4=excluded.ipv4,
|
|
ipv6=excluded.ipv6,
|
|
updated_at=excluded.updated_at
|
|
""",
|
|
(result.hostname, result.ipv4, result.ipv6, now),
|
|
)
|
|
return result
|
|
|
|
def preview_update(
|
|
self,
|
|
hostname: str,
|
|
ipv4: str | None,
|
|
ipv6: str | None,
|
|
update_ipv4: bool = True,
|
|
update_ipv6: bool = True,
|
|
) -> UpdateResult:
|
|
update = RecordUpdateInput(hostname, ipv4, ipv6, update_ipv4, update_ipv6)
|
|
with self.connect() as conn:
|
|
existing = conn.execute("SELECT ipv4, ipv6 FROM records WHERE hostname = ?", (hostname,)).fetchone()
|
|
current_ipv4 = existing["ipv4"] if existing else None
|
|
current_ipv6 = existing["ipv6"] if existing else None
|
|
return _planned_record_update(
|
|
update,
|
|
current_ipv4=current_ipv4,
|
|
current_ipv6=current_ipv6,
|
|
existing=existing is not None,
|
|
)
|
|
|
|
def list_records(self) -> list[RecordRow]:
|
|
return self.list_records_page(limit=10_000, offset=0)
|
|
|
|
def list_records_page(self, limit: int = 100, offset: int = 0) -> list[RecordRow]:
|
|
with self.connect() as conn:
|
|
rows = conn.execute(
|
|
"""
|
|
SELECT hostname, ipv4, ipv6, updated_at
|
|
FROM records
|
|
ORDER BY hostname
|
|
LIMIT ? OFFSET ?
|
|
""",
|
|
(limit, offset),
|
|
).fetchall()
|
|
return [record_row(row) for row in rows]
|
|
|
|
def count_records(self) -> int:
|
|
with self.connect() as conn:
|
|
row = conn.execute("SELECT COUNT(*) FROM records").fetchone()
|
|
return int(row[0])
|
|
|
|
def log_update_event(self, event: UpdateEventInput | str, *legacy_values: object) -> None:
|
|
event = _update_event_input(event, legacy_values)
|
|
with self.connect() as conn:
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO update_events(hostname, ipv4, ipv6, status, detail, source_ip, created_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
event.hostname,
|
|
event.ipv4,
|
|
event.ipv6,
|
|
event.status,
|
|
event.detail[:500],
|
|
event.source_ip,
|
|
datetime.now(UTC).isoformat(),
|
|
),
|
|
)
|
|
|
|
def log_update_rollup(self, event: UpdateEventInput | str, *legacy_values: object) -> None:
|
|
event = _update_event_input(event, legacy_values)
|
|
now = datetime.now(UTC).isoformat()
|
|
with self.connect() as conn:
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO update_event_rollups(
|
|
hostname, ipv4, ipv6, status, detail, source_ip, first_seen_at, last_seen_at, count
|
|
)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, 1)
|
|
ON CONFLICT(hostname, ipv4, ipv6, status, detail, source_ip) DO UPDATE SET
|
|
last_seen_at=excluded.last_seen_at,
|
|
count=count + 1
|
|
""",
|
|
(
|
|
event.hostname,
|
|
event.ipv4 or "",
|
|
event.ipv6 or "",
|
|
event.status,
|
|
event.detail[:500],
|
|
event.source_ip,
|
|
now,
|
|
now,
|
|
),
|
|
)
|
|
|
|
def list_update_events(self, limit: int = 100, offset: int = 0) -> list[UpdateEventRow]:
|
|
return self.list_update_events_page(limit=limit, offset=offset)
|
|
|
|
def list_update_events_page(self, limit: int = 100, offset: int = 0) -> list[UpdateEventRow]:
|
|
with self.connect() as conn:
|
|
rows = conn.execute(
|
|
"""
|
|
SELECT hostname, ipv4, ipv6, status, detail, source_ip, created_at, count
|
|
FROM (
|
|
SELECT hostname, ipv4, ipv6, status, detail, source_ip, created_at, 1 AS count
|
|
FROM update_events
|
|
UNION ALL
|
|
SELECT hostname, ipv4, ipv6, status, detail, source_ip, last_seen_at AS created_at, count
|
|
FROM update_event_rollups
|
|
)
|
|
ORDER BY created_at DESC
|
|
LIMIT ? OFFSET ?
|
|
""",
|
|
(limit, offset),
|
|
).fetchall()
|
|
return [update_event_row(row) for row in rows]
|
|
|
|
def list_update_events_before(self, limit: int = 100, before: str | None = None) -> list[UpdateEventRow]:
|
|
where_clause = "WHERE created_at < ?" if before else ""
|
|
params: tuple[object, ...] = (before, limit) if before else (limit,)
|
|
with self.connect() as conn:
|
|
rows = conn.execute(
|
|
f"""
|
|
SELECT hostname, ipv4, ipv6, status, detail, source_ip, created_at, count
|
|
FROM (
|
|
SELECT hostname, ipv4, ipv6, status, detail, source_ip, created_at, 1 AS count
|
|
FROM update_events
|
|
UNION ALL
|
|
SELECT hostname, ipv4, ipv6, status, detail, source_ip, last_seen_at AS created_at, count
|
|
FROM update_event_rollups
|
|
)
|
|
{where_clause}
|
|
ORDER BY created_at DESC
|
|
LIMIT ?
|
|
""",
|
|
params,
|
|
).fetchall()
|
|
return [update_event_row(row) for row in rows]
|
|
|
|
def count_update_events(self) -> int:
|
|
with self.connect() as conn:
|
|
event_count = conn.execute("SELECT COUNT(*) FROM update_events").fetchone()[0]
|
|
rollup_count = conn.execute("SELECT COUNT(*) FROM update_event_rollups").fetchone()[0]
|
|
return int(event_count) + int(rollup_count)
|
|
|
|
def list_update_events_for_hostname(self, hostname: str, limit: int = 25) -> list[UpdateEventRow]:
|
|
with self.connect() as conn:
|
|
event_rows = conn.execute(
|
|
"""
|
|
SELECT hostname, ipv4, ipv6, status, detail, source_ip, created_at, 1 AS count
|
|
FROM update_events
|
|
WHERE hostname = ?
|
|
ORDER BY created_at DESC
|
|
LIMIT ?
|
|
""",
|
|
(hostname, limit),
|
|
).fetchall()
|
|
rollup_rows = conn.execute(
|
|
"""
|
|
SELECT hostname, ipv4, ipv6, status, detail, source_ip, last_seen_at AS created_at, count
|
|
FROM update_event_rollups
|
|
WHERE hostname = ?
|
|
ORDER BY last_seen_at DESC
|
|
LIMIT ?
|
|
""",
|
|
(hostname, limit),
|
|
).fetchall()
|
|
return _merge_event_rows(event_rows, rollup_rows, limit, 0)
|
|
|
|
|
|
def _merge_event_rows(
|
|
event_rows: list[sqlite3.Row],
|
|
rollup_rows: list[sqlite3.Row],
|
|
limit: int,
|
|
offset: int,
|
|
) -> list[UpdateEventRow]:
|
|
rows = [update_event_row(row) for row in event_rows]
|
|
rows.extend(update_event_row(row) for row in rollup_rows)
|
|
rows.sort(key=lambda row: str(row["created_at"]), reverse=True)
|
|
return rows[offset : offset + limit]
|