328 lines
13 KiB
Python
328 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import shutil
|
|
import threading
|
|
import time
|
|
import uuid
|
|
import zipfile
|
|
from collections.abc import Callable
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from app.config import Settings
|
|
from app.crawler import (
|
|
DEFAULT_ASSET_TYPES,
|
|
CaptureCancelled,
|
|
CaptureError,
|
|
CaptureOptions,
|
|
SiteCrawler,
|
|
)
|
|
from app.db import Database
|
|
from app.url_policy import NetworkRules, UrlPolicyError, resolve_and_validate
|
|
|
|
CANCELLATION_POLL_SECONDS = 0.25
|
|
PROGRESS_WRITE_INTERVAL_SECONDS = 0.5
|
|
|
|
|
|
def _remove_tree(path: Path, allowed_root: Path) -> None:
|
|
"""Delete only a child path beneath the configured work root."""
|
|
root = allowed_root.resolve()
|
|
resolved = path.resolve()
|
|
if resolved == root or not resolved.is_relative_to(root):
|
|
raise RuntimeError("Refusing to remove a path outside SiteHarbor's work directory.")
|
|
shutil.rmtree(resolved, ignore_errors=True)
|
|
|
|
|
|
def _create_archive(
|
|
source_dir: Path,
|
|
destination: Path,
|
|
archive_root: str,
|
|
cancelled: Callable[[], bool],
|
|
deadline_exceeded: Callable[[], bool],
|
|
heartbeat: Callable[[], None],
|
|
) -> int:
|
|
partial = destination.with_suffix(".zip.partial")
|
|
partial.unlink(missing_ok=True)
|
|
last_heartbeat = time.monotonic()
|
|
try:
|
|
with zipfile.ZipFile(
|
|
partial,
|
|
mode="w",
|
|
compression=zipfile.ZIP_DEFLATED,
|
|
compresslevel=6,
|
|
strict_timestamps=False,
|
|
) as archive:
|
|
for source_path in sorted(source_dir.rglob("*")):
|
|
if cancelled():
|
|
raise CaptureCancelled("Capture cancelled while packaging the archive")
|
|
if deadline_exceeded():
|
|
raise CaptureError(
|
|
"The configured capture time limit was reached while packaging"
|
|
)
|
|
if time.monotonic() - last_heartbeat >= 5:
|
|
heartbeat()
|
|
last_heartbeat = time.monotonic()
|
|
if not source_path.is_file() or source_path.is_symlink():
|
|
continue
|
|
relative = source_path.relative_to(source_dir).as_posix()
|
|
with source_path.open("rb") as source_file:
|
|
with archive.open(f"{archive_root}/{relative}", mode="w") as archive_file:
|
|
while chunk := source_file.read(1_024 * 1_024):
|
|
if cancelled():
|
|
raise CaptureCancelled(
|
|
"Capture cancelled while packaging the archive"
|
|
)
|
|
if deadline_exceeded():
|
|
raise CaptureError(
|
|
"The configured capture time limit was reached while packaging"
|
|
)
|
|
if time.monotonic() - last_heartbeat >= 5:
|
|
heartbeat()
|
|
last_heartbeat = time.monotonic()
|
|
archive_file.write(chunk)
|
|
os.replace(partial, destination)
|
|
return destination.stat().st_size
|
|
except Exception:
|
|
partial.unlink(missing_ok=True)
|
|
destination.unlink(missing_ok=True)
|
|
raise
|
|
|
|
|
|
class CaptureService:
|
|
"""Runs one leased capture and persists a terminal SQLite job state."""
|
|
|
|
def __init__(self, database: Database, settings: Settings, worker_id: str) -> None:
|
|
self.database = database
|
|
self.settings = settings
|
|
self.worker_id = worker_id
|
|
|
|
def run(self, job: dict[str, Any]) -> None:
|
|
job_id = str(job["id"])
|
|
config = dict(job["config"])
|
|
attempt_id = uuid.uuid4().hex[:12]
|
|
work_dir = self.settings.work_dir / f"{job_id}-{attempt_id}"
|
|
artifact_path = self.settings.artifacts_dir / f"{job_id}-{attempt_id}.zip"
|
|
report_path = self.settings.reports_dir / f"{job_id}-{attempt_id}.json"
|
|
heartbeat_stop = threading.Event()
|
|
lease_lost = threading.Event()
|
|
cancellation_requested = threading.Event()
|
|
cancellation_lock = threading.Lock()
|
|
heartbeat_thread: threading.Thread | None = None
|
|
capture_deadline = time.monotonic() + int(config["max_duration_seconds"])
|
|
last_cancellation_check = 0.0
|
|
last_progress_at = 0.0
|
|
last_progress_phase: str | None = None
|
|
|
|
def capture_cancelled() -> bool:
|
|
nonlocal last_cancellation_check
|
|
if lease_lost.is_set() or cancellation_requested.is_set():
|
|
return True
|
|
now = time.monotonic()
|
|
with cancellation_lock:
|
|
if now - last_cancellation_check < CANCELLATION_POLL_SECONDS:
|
|
return False
|
|
last_cancellation_check = now
|
|
# Fetch workers call this predicate for every response chunk. Poll SQLite sparingly.
|
|
if self.database.worker_should_stop(job_id, self.worker_id):
|
|
cancellation_requested.set()
|
|
return True
|
|
return False
|
|
|
|
def renew_lease() -> None:
|
|
interval = min(20, max(3, self.settings.lease_seconds // 3))
|
|
while not heartbeat_stop.wait(interval):
|
|
try:
|
|
if not self.database.renew_lease(
|
|
job_id, self.worker_id, self.settings.lease_seconds
|
|
):
|
|
lease_lost.set()
|
|
return
|
|
except Exception:
|
|
# Ownership checks stop this attempt if the lease eventually expires.
|
|
continue
|
|
|
|
try:
|
|
_remove_tree_if_present(work_dir, self.settings.work_dir)
|
|
work_dir.mkdir(parents=True, exist_ok=True)
|
|
if capture_cancelled():
|
|
raise CaptureCancelled("Capture cancelled before it started")
|
|
|
|
heartbeat_thread = threading.Thread(
|
|
target=renew_lease,
|
|
name=f"siteharbor-lease-{job_id[-8:]}",
|
|
daemon=True,
|
|
)
|
|
heartbeat_thread.start()
|
|
|
|
rules = NetworkRules(
|
|
allow_private_networks=self.settings.allow_private_networks,
|
|
allow_nonstandard_ports=self.settings.allow_nonstandard_ports,
|
|
)
|
|
resolve_and_validate(str(job["normalized_url"]), rules)
|
|
remaining_seconds = int(capture_deadline - time.monotonic())
|
|
if remaining_seconds <= 0:
|
|
raise CaptureError(
|
|
"The configured capture time limit was reached during validation"
|
|
)
|
|
options = CaptureOptions(
|
|
source_url=str(job["normalized_url"]),
|
|
include_external_assets=bool(config["include_external_assets"]),
|
|
max_pages=int(config["max_pages"]),
|
|
max_depth=int(config["max_depth"]),
|
|
max_bytes=int(config["max_bytes"]),
|
|
max_duration_seconds=remaining_seconds,
|
|
parallel_connections=int(
|
|
config.get("parallel_connections", self.settings.fetch_concurrency)
|
|
),
|
|
download_asset_types=tuple(
|
|
config.get("download_asset_types", DEFAULT_ASSET_TYPES)
|
|
),
|
|
skip_file_types=tuple(config.get("skip_file_types", ())),
|
|
use_proxy=bool(config.get("use_proxy", False)),
|
|
)
|
|
|
|
def update(
|
|
phase: str, message: str, stats: dict[str, int], warnings: list[str]
|
|
) -> None:
|
|
nonlocal last_progress_at, last_progress_phase
|
|
now = time.monotonic()
|
|
if (
|
|
phase == last_progress_phase
|
|
and now - last_progress_at < PROGRESS_WRITE_INTERVAL_SECONDS
|
|
):
|
|
return
|
|
last_progress_at = now
|
|
last_progress_phase = phase
|
|
renewed = self.database.update_progress(
|
|
job_id,
|
|
self.worker_id,
|
|
self.settings.lease_seconds,
|
|
phase=phase,
|
|
message=message,
|
|
stats=stats,
|
|
warnings=warnings,
|
|
)
|
|
if not renewed:
|
|
raise CaptureCancelled("Capture lease was reassigned")
|
|
|
|
if not self.database.update_progress(
|
|
job_id,
|
|
self.worker_id,
|
|
self.settings.lease_seconds,
|
|
phase="Validating",
|
|
message="Checking capture policy and target reachability",
|
|
stats=_job_stats(job),
|
|
warnings=[],
|
|
):
|
|
raise CaptureCancelled("Capture lease was reassigned")
|
|
|
|
crawler = SiteCrawler(
|
|
settings=self.settings,
|
|
options=options,
|
|
cancelled=capture_cancelled,
|
|
on_progress=update,
|
|
)
|
|
result = crawler.run(work_dir)
|
|
|
|
if capture_cancelled():
|
|
raise CaptureCancelled("Capture cancelled before packaging")
|
|
if time.monotonic() >= capture_deadline:
|
|
raise CaptureError("The configured capture time limit was reached before packaging")
|
|
|
|
update("Packaging", "Creating a portable ZIP archive", result.stats, result.warnings)
|
|
archive_size = _create_archive(
|
|
result.site_dir,
|
|
artifact_path,
|
|
archive_root=f"siteharbor-{job_id}",
|
|
cancelled=capture_cancelled,
|
|
deadline_exceeded=lambda: time.monotonic() >= capture_deadline,
|
|
heartbeat=lambda: update(
|
|
"Packaging",
|
|
"Creating a portable ZIP archive",
|
|
result.stats,
|
|
result.warnings,
|
|
),
|
|
)
|
|
|
|
report = {
|
|
**result.report,
|
|
"job_id": job_id,
|
|
"archive": {
|
|
"filename": artifact_path.name,
|
|
"size_bytes": archive_size,
|
|
"archive_root": f"siteharbor-{job_id}",
|
|
},
|
|
}
|
|
if time.monotonic() >= capture_deadline:
|
|
raise CaptureError(
|
|
"The configured capture time limit was reached before completion"
|
|
)
|
|
report_path.write_text(
|
|
json.dumps(report, indent=2, ensure_ascii=True), encoding="utf-8"
|
|
)
|
|
marked_ready = self.database.mark_ready(
|
|
job_id,
|
|
self.worker_id,
|
|
artifact_path=str(artifact_path),
|
|
report_path=str(report_path),
|
|
archive_size=archive_size,
|
|
entry_point=result.entry_point,
|
|
stats=result.stats,
|
|
warnings=result.warnings,
|
|
retention_days=int(config["retention_days"]),
|
|
)
|
|
if not marked_ready:
|
|
artifact_path.unlink(missing_ok=True)
|
|
report_path.unlink(missing_ok=True)
|
|
if self.database.is_cancel_requested(job_id):
|
|
self.database.mark_cancelled(
|
|
job_id,
|
|
self.worker_id,
|
|
"Capture cancelled before the archive became available",
|
|
)
|
|
except CaptureCancelled as error:
|
|
artifact_path.unlink(missing_ok=True)
|
|
report_path.unlink(missing_ok=True)
|
|
if self.database.is_cancel_requested(job_id):
|
|
self.database.mark_cancelled(job_id, self.worker_id, str(error))
|
|
except UrlPolicyError as error:
|
|
artifact_path.unlink(missing_ok=True)
|
|
self.database.mark_failed(job_id, self.worker_id, "target_blocked", str(error))
|
|
except CaptureError as error:
|
|
artifact_path.unlink(missing_ok=True)
|
|
self.database.mark_failed(job_id, self.worker_id, "capture_failed", str(error))
|
|
except Exception as error:
|
|
artifact_path.unlink(missing_ok=True)
|
|
self.database.mark_failed(
|
|
job_id,
|
|
self.worker_id,
|
|
"internal_error",
|
|
f"Capture failed: {error}",
|
|
)
|
|
finally:
|
|
heartbeat_stop.set()
|
|
if heartbeat_thread:
|
|
heartbeat_thread.join(timeout=1)
|
|
try:
|
|
_remove_tree_if_present(work_dir, self.settings.work_dir)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def _job_stats(job: dict[str, Any]) -> dict[str, int]:
|
|
return {
|
|
"pages_found": int(job.get("pages_found", 0)),
|
|
"pages_fetched": int(job.get("pages_fetched", 0)),
|
|
"assets_found": int(job.get("assets_found", 0)),
|
|
"assets_fetched": int(job.get("assets_fetched", 0)),
|
|
"bytes_downloaded": int(job.get("bytes_downloaded", 0)),
|
|
"files_written": int(job.get("files_written", 0)),
|
|
}
|
|
|
|
|
|
def _remove_tree_if_present(path: Path, root: Path) -> None:
|
|
if path.exists():
|
|
_remove_tree(path, root)
|