Add collector engine: base, custom source, profile loader, wikileaks helpers
This commit is contained in:
n0mad1k
@@ -0,0 +1,292 @@
|
||||
"""Abstract base collector for Mosaic document acquisition."""
|
||||
import os
|
||||
import tempfile
|
||||
import hashlib
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import httpx
|
||||
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
|
||||
|
||||
from storage.db import MosaicDB
|
||||
from storage.models import Document
|
||||
from utils.http_limiter import HTTPRateLimiter
|
||||
from utils.cache import DownloadCache
|
||||
from utils.sanitize import sanitize_filename, validate_content_type
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BaseCollector(ABC):
|
||||
"""Abstract base for all document collectors."""
|
||||
|
||||
def __init__(self, db: MosaicDB, cache: DownloadCache,
|
||||
rate_limiter: Optional[HTTPRateLimiter] = None,
|
||||
config: Optional[dict] = None):
|
||||
self.db = db
|
||||
self.cache = cache
|
||||
self.rate_limiter = rate_limiter or HTTPRateLimiter()
|
||||
self.config = config or {}
|
||||
self.user_agent = self.config.get('user_agent', 'Mozilla/5.0 (compatible; research-crawler/1.0)')
|
||||
self.timeout = int(self.config.get('timeout', 30))
|
||||
self.max_retries = int(self.config.get('max_retries', 3))
|
||||
|
||||
@abstractmethod
|
||||
def discover(self, since: Optional[str] = None,
|
||||
until: Optional[str] = None,
|
||||
limit: Optional[int] = None) -> list[dict]:
|
||||
"""Discover document URLs from the source.
|
||||
|
||||
Returns list of dicts: {'url': str, 'title': str, 'date': str|None, 'doc_type': str}
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_name(self) -> str:
|
||||
"""Return the source name."""
|
||||
pass
|
||||
|
||||
def expected_count(self) -> int:
|
||||
"""Minimum expected document count for validation."""
|
||||
return 0
|
||||
|
||||
def fetch(self, url: str, collection: str) -> Optional[Document]:
|
||||
"""Download a URL with atomic writes and hash verification.
|
||||
|
||||
Returns Document if successful, None if failed.
|
||||
"""
|
||||
# Check if already cached
|
||||
existing = self.db.get_document_by_url(url)
|
||||
if existing and existing.status == 'CACHED':
|
||||
cached_path = self.cache.get_cached_path(
|
||||
existing.content_hash, collection,
|
||||
self._url_extension(url)
|
||||
)
|
||||
if cached_path:
|
||||
logger.debug("Already cached: %s", url)
|
||||
return existing
|
||||
|
||||
# Rate limit
|
||||
self.rate_limiter.wait()
|
||||
|
||||
# Download to temp file
|
||||
try:
|
||||
with httpx.Client(timeout=self.timeout, follow_redirects=True) as client:
|
||||
response = client.get(url, headers={'User-Agent': self.user_agent})
|
||||
response.raise_for_status()
|
||||
|
||||
content = response.content
|
||||
content_hash = hashlib.sha256(content).hexdigest()
|
||||
content_length = len(content)
|
||||
etag = response.headers.get('etag')
|
||||
last_modified = response.headers.get('last-modified')
|
||||
content_type = response.headers.get('content-type', '').split(';')[0].strip()
|
||||
|
||||
# Write to temp file
|
||||
ext = self._url_extension(url) or self._mime_to_ext(content_type)
|
||||
tmp_fd, tmp_path = tempfile.mkstemp(suffix=f'.{ext}.tmp', dir=str(self.cache.base_dir))
|
||||
try:
|
||||
os.write(tmp_fd, content)
|
||||
os.close(tmp_fd)
|
||||
|
||||
# Validate content type by magic bytes
|
||||
try:
|
||||
detected_type = validate_content_type(tmp_path)
|
||||
except ValueError:
|
||||
detected_type = content_type
|
||||
|
||||
# Atomic store
|
||||
final_path = self.cache.store_atomic(
|
||||
tmp_path, content_hash, collection, ext, content_length
|
||||
)
|
||||
except Exception:
|
||||
# Clean up temp file on failure
|
||||
os.close(tmp_fd) if not os.path.exists(tmp_path) else None
|
||||
if os.path.exists(tmp_path):
|
||||
os.unlink(tmp_path)
|
||||
raise
|
||||
|
||||
# Build document
|
||||
doc_type = self._detect_doc_type(detected_type)
|
||||
doc = Document(
|
||||
id=content_hash,
|
||||
source=collection,
|
||||
source_url=url,
|
||||
doc_type=doc_type,
|
||||
title=self._extract_title(url),
|
||||
raw_path=final_path,
|
||||
content_hash=content_hash,
|
||||
fetch_date=datetime.utcnow().isoformat(),
|
||||
etag=etag,
|
||||
last_modified=last_modified,
|
||||
char_count=len(content.decode('utf-8', errors='replace')),
|
||||
status='CACHED',
|
||||
parse_status='PENDING',
|
||||
)
|
||||
|
||||
self.db.insert_document(doc)
|
||||
return doc
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
logger.error("HTTP %d for %s: %s", e.response.status_code, url, e)
|
||||
except httpx.RequestError as e:
|
||||
logger.error("Request error for %s: %s", url, e)
|
||||
except Exception as e:
|
||||
logger.error("Download failed for %s: %s", url, e)
|
||||
|
||||
# Record failure
|
||||
if existing:
|
||||
self.db.update_document_status(existing.id, status='DOWNLOAD_FAILED')
|
||||
return None
|
||||
|
||||
def check_update(self, url: str) -> bool:
|
||||
"""Check if a cached URL has been updated (HTTP HEAD with ETag/Last-Modified).
|
||||
|
||||
Returns True if the document has changed and should be re-downloaded.
|
||||
"""
|
||||
existing = self.db.get_document_by_url(url)
|
||||
if not existing:
|
||||
return True
|
||||
|
||||
try:
|
||||
self.rate_limiter.wait()
|
||||
with httpx.Client(timeout=self.timeout, follow_redirects=True) as client:
|
||||
response = client.head(url, headers={'User-Agent': self.user_agent})
|
||||
response.raise_for_status()
|
||||
|
||||
new_etag = response.headers.get('etag')
|
||||
new_modified = response.headers.get('last-modified')
|
||||
|
||||
if existing.etag and new_etag:
|
||||
return existing.etag != new_etag
|
||||
if existing.last_modified and new_modified:
|
||||
return existing.last_modified != new_modified
|
||||
|
||||
# Can't determine — assume changed
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.warning("HEAD check failed for %s: %s", url, e)
|
||||
return False
|
||||
|
||||
def collect(self, since: Optional[str] = None, until: Optional[str] = None,
|
||||
limit: Optional[int] = None, update: bool = False,
|
||||
retry_failed: bool = False) -> dict:
|
||||
"""Main collection entry point with progress tracking.
|
||||
|
||||
Returns dict with counts: {'discovered', 'downloaded', 'cached', 'failed', 'updated'}
|
||||
"""
|
||||
name = self.get_name()
|
||||
stats = {'discovered': 0, 'downloaded': 0, 'cached': 0, 'failed': 0, 'updated': 0}
|
||||
|
||||
# Check disk space
|
||||
if not self.cache.check_disk_space(required_mb=100):
|
||||
logger.error("Insufficient disk space for collection")
|
||||
return stats
|
||||
|
||||
# Discover URLs
|
||||
urls = self.discover(since=since, until=until, limit=limit)
|
||||
stats['discovered'] = len(urls)
|
||||
|
||||
# Validate against expected count
|
||||
expected = self.expected_count()
|
||||
if expected > 0 and len(urls) < expected * 0.1:
|
||||
logger.warning(
|
||||
"Discovered only %d URLs (expected >%d). "
|
||||
"Site structure may have changed.",
|
||||
len(urls), int(expected * 0.1)
|
||||
)
|
||||
|
||||
# Filter for update mode
|
||||
if update:
|
||||
urls_to_process = []
|
||||
for item in urls:
|
||||
if self.check_update(item['url']):
|
||||
urls_to_process.append(item)
|
||||
stats['updated'] += 1
|
||||
urls = urls_to_process
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
transient=True,
|
||||
) as progress:
|
||||
task = progress.add_task(f"Collecting {name}", total=len(urls))
|
||||
|
||||
for item in urls:
|
||||
url = item['url']
|
||||
|
||||
# Skip already cached (unless update mode)
|
||||
existing = self.db.get_document_by_url(url)
|
||||
if existing and existing.status == 'CACHED' and not update:
|
||||
stats['cached'] += 1
|
||||
progress.advance(task)
|
||||
continue
|
||||
|
||||
# Skip failed unless retry
|
||||
if existing and existing.status == 'DOWNLOAD_FAILED' and not retry_failed:
|
||||
stats['failed'] += 1
|
||||
progress.advance(task)
|
||||
continue
|
||||
|
||||
doc = self.fetch(url, name)
|
||||
if doc:
|
||||
# Update with discovered metadata
|
||||
if item.get('title'):
|
||||
doc.title = item['title']
|
||||
if item.get('date'):
|
||||
doc.date = item['date']
|
||||
self.db.insert_document(doc)
|
||||
stats['downloaded'] += 1
|
||||
else:
|
||||
stats['failed'] += 1
|
||||
|
||||
progress.advance(task)
|
||||
|
||||
return stats
|
||||
|
||||
def _url_extension(self, url: str) -> str:
|
||||
"""Extract file extension from URL."""
|
||||
from urllib.parse import urlparse
|
||||
path = urlparse(url).path
|
||||
ext = Path(path).suffix.lstrip('.')
|
||||
if ext and len(ext) <= 10:
|
||||
return sanitize_filename(ext).lower()
|
||||
return 'html'
|
||||
|
||||
def _mime_to_ext(self, mime: str) -> str:
|
||||
"""Convert MIME type to file extension."""
|
||||
mapping = {
|
||||
'text/html': 'html',
|
||||
'application/pdf': 'pdf',
|
||||
'text/plain': 'txt',
|
||||
'message/rfc822': 'eml',
|
||||
'application/json': 'json',
|
||||
}
|
||||
return mapping.get(mime, 'bin')
|
||||
|
||||
def _detect_doc_type(self, content_type: str) -> str:
|
||||
"""Map content type to our doc_type enum."""
|
||||
if 'html' in content_type:
|
||||
return 'html'
|
||||
elif 'pdf' in content_type:
|
||||
return 'pdf'
|
||||
elif 'email' in content_type or 'rfc822' in content_type:
|
||||
return 'email'
|
||||
elif 'text' in content_type:
|
||||
return 'text'
|
||||
return 'html' # default for web pages
|
||||
|
||||
def _extract_title(self, url: str) -> str:
|
||||
"""Extract a title from URL path."""
|
||||
from urllib.parse import urlparse, unquote
|
||||
path = urlparse(url).path
|
||||
parts = [p for p in path.split('/') if p]
|
||||
if parts:
|
||||
return unquote(parts[-1]).replace('-', ' ').replace('_', ' ')
|
||||
return url
|
||||
Reference in New Issue
Block a user