diff --git a/collectors/__init__.py b/collectors/__init__.py index e69de29..ec5685d 100644 --- a/collectors/__init__.py +++ b/collectors/__init__.py @@ -0,0 +1,3 @@ +from .base import BaseCollector +from .custom import CustomSource +from .profile_loader import load_profile, list_profiles diff --git a/collectors/base.py b/collectors/base.py new file mode 100644 index 0000000..becd7c0 --- /dev/null +++ b/collectors/base.py @@ -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 diff --git a/collectors/custom.py b/collectors/custom.py new file mode 100644 index 0000000..58ae20e --- /dev/null +++ b/collectors/custom.py @@ -0,0 +1,217 @@ +"""Universal source collector configured by YAML profiles.""" +import re +import logging +from urllib.parse import urlparse, urljoin, unquote +from typing import Optional +from datetime import datetime + +import httpx +from bs4 import BeautifulSoup + +from .base import BaseCollector +from storage.db import MosaicDB +from utils.http_limiter import HTTPRateLimiter +from utils.cache import DownloadCache + +logger = logging.getLogger(__name__) + + +class CustomSource(BaseCollector): + """Universal collector that crawls any source based on profile configuration.""" + + def __init__(self, profile: dict, db: MosaicDB, cache: DownloadCache, + rate_limiter: Optional[HTTPRateLimiter] = None, + config: Optional[dict] = None): + # Set rate from profile + rate = profile.get('rate_limit', 1.0) + if rate_limiter is None: + rate_limiter = HTTPRateLimiter(rate=rate, jitter=0.5) + + super().__init__(db, cache, rate_limiter, config) + self.profile = profile + self.name = profile['name'] + self.base_url = profile['base_url'] + self.domain_allowlist = set(profile.get('domain_allowlist', [])) + self.crawl_rules = profile.get('crawl_rules', {}) + self.content_types = set(profile.get('content_types', ['text/html', 'application/pdf'])) + self.mirrors = profile.get('mirrors', []) + + # Compile URL patterns + self.url_patterns = [ + re.compile(self._glob_to_regex(p)) + for p in self.crawl_rules.get('url_patterns', []) + ] + self.exclude_patterns = [ + re.compile(self._glob_to_regex(p)) + for p in self.crawl_rules.get('exclude_patterns', []) + ] + + # Import wikileaks helpers if applicable + self._wikileaks_helpers = None + if any(d.endswith('wikileaks.org') for d in self.domain_allowlist): + try: + from . import wikileaks + self._wikileaks_helpers = wikileaks + logger.debug("WikiLeaks helpers loaded for %s", self.name) + except ImportError: + pass + + def get_name(self) -> str: + return self.name + + def expected_count(self) -> int: + return self.profile.get('expected_min_count', 0) + + def discover(self, since: Optional[str] = None, + until: Optional[str] = None, + limit: Optional[int] = None) -> list[dict]: + """Crawl source to discover document URLs.""" + max_depth = self.crawl_rules.get('max_depth', 2) + max_docs = limit or self.crawl_rules.get('max_documents', 1000) + + discovered = [] + visited = set() + queue = [(self.base_url, 0)] + + while queue and len(discovered) < max_docs: + url, depth = queue.pop(0) + + if url in visited: + continue + visited.add(url) + + if depth > max_depth: + continue + + # Domain check + if not self._is_allowed_domain(url): + continue + + # Exclude check + if self._is_excluded(url): + continue + + logger.debug("Crawling: %s (depth %d)", url, depth) + self.rate_limiter.wait() + + 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_type = response.headers.get('content-type', '').split(';')[0].strip() + + # If it's a document (not just a page to crawl), add it + if self._matches_url_pattern(url) or self._is_document_type(content_type): + item = { + 'url': url, + 'title': self._extract_title_from_url(url), + 'date': None, + 'doc_type': self._content_type_to_doc_type(content_type), + } + + # Date filtering + if since and item.get('date') and item['date'] < since: + continue + if until and item.get('date') and item['date'] > until: + continue + + discovered.append(item) + + # Extract links if HTML and should follow links + if 'html' in content_type and self.crawl_rules.get('follow_links', True): + links = self._extract_links(response.text, url) + + # Use wikileaks helpers if available + if self._wikileaks_helpers: + extra = self._wikileaks_helpers.extract_extra_links( + response.text, url, self.name + ) + links.extend(extra) + + for link in links: + if link not in visited and len(discovered) < max_docs: + queue.append((link, depth + 1)) + + except httpx.HTTPStatusError as e: + logger.warning("HTTP %d for %s", e.response.status_code, url) + except httpx.RequestError as e: + logger.warning("Request error for %s: %s", url, e) + except Exception as e: + logger.warning("Error crawling %s: %s", url, e) + + logger.info("Discovered %d documents from %s", len(discovered), self.name) + return discovered[:max_docs] + + def _is_allowed_domain(self, url: str) -> bool: + """Check if URL domain is in the allowlist.""" + if not self.domain_allowlist: + return True + parsed = urlparse(url) + return parsed.netloc in self.domain_allowlist + + def _is_excluded(self, url: str) -> bool: + """Check if URL matches any exclude pattern.""" + for pattern in self.exclude_patterns: + if pattern.search(url): + return True + return False + + def _matches_url_pattern(self, url: str) -> bool: + """Check if URL matches any include pattern.""" + if not self.url_patterns: + return True + for pattern in self.url_patterns: + if pattern.search(url): + return True + return False + + def _is_document_type(self, content_type: str) -> bool: + """Check if content type is one we want to collect.""" + return content_type in self.content_types + + def _extract_links(self, html: str, base_url: str) -> list[str]: + """Extract absolute URLs from HTML page.""" + links = [] + try: + soup = BeautifulSoup(html, 'html.parser') + for tag in soup.find_all('a', href=True): + href = tag['href'] + # Skip fragments, javascript, mailto + if href.startswith(('#', 'javascript:', 'mailto:')): + continue + absolute = urljoin(base_url, href) + # Strip fragments + absolute = absolute.split('#')[0] + if absolute and self._is_allowed_domain(absolute): + links.append(absolute) + except Exception as e: + logger.warning("Link extraction failed for %s: %s", base_url, e) + return links + + def _extract_title_from_url(self, url: str) -> str: + """Extract a title from URL path.""" + path = urlparse(url).path + parts = [p for p in path.split('/') if p] + if parts: + return unquote(parts[-1]).replace('-', ' ').replace('_', ' ') + return url + + def _content_type_to_doc_type(self, content_type: str) -> str: + """Map content type to doc_type.""" + if 'html' in content_type: + return 'html' + elif 'pdf' in content_type: + return 'pdf' + elif 'text' in content_type: + return 'text' + return 'html' + + @staticmethod + def _glob_to_regex(pattern: str) -> str: + """Convert glob pattern to regex.""" + # Escape regex special chars except * and ? + pattern = re.escape(pattern) + # Convert glob wildcards + pattern = pattern.replace(r'\*', '.*') + pattern = pattern.replace(r'\?', '.') + return pattern diff --git a/collectors/profile_loader.py b/collectors/profile_loader.py new file mode 100644 index 0000000..660b17f --- /dev/null +++ b/collectors/profile_loader.py @@ -0,0 +1,129 @@ +"""Load YAML source profiles for Mosaic collectors.""" +import logging +from pathlib import Path +from typing import Optional + +import yaml + +logger = logging.getLogger(__name__) + +PROFILES_DIR = Path(__file__).parent / 'profiles' + +REQUIRED_FIELDS = ['name', 'base_url', 'domain_allowlist'] +DEFAULT_CRAWL_RULES = { + 'max_depth': 2, + 'max_documents': 1000, + 'follow_links': True, + 'url_patterns': [], + 'exclude_patterns': ['*.css', '*.js', '*.ico', '*.png', '*.jpg', '*.gif'], +} + + +def load_profile(name_or_path: str) -> dict: + """Load a source profile by name or file path. + + Args: + name_or_path: Profile name (e.g., 'vault7') or path to YAML file + + Returns: + Profile dict with all fields populated (defaults filled in) + + Raises: + FileNotFoundError: If profile not found + ValueError: If profile missing required fields + """ + path = Path(name_or_path) + + # Check if it's a direct path + if path.exists() and path.suffix in ('.yaml', '.yml'): + profile_path = path + else: + # Look in profiles directory + profile_path = PROFILES_DIR / f"{name_or_path}.yaml" + if not profile_path.exists(): + profile_path = PROFILES_DIR / f"{name_or_path}.yml" + + if not profile_path.exists(): + raise FileNotFoundError(f"Profile not found: {name_or_path}") + + with open(profile_path, 'r') as f: + profile = yaml.safe_load(f) + + if not profile: + raise ValueError(f"Empty profile: {profile_path}") + + # Validate required fields + for field in REQUIRED_FIELDS: + if field not in profile: + raise ValueError(f"Profile missing required field: {field}") + + # Fill defaults + profile.setdefault('description', '') + profile.setdefault('content_types', ['text/html', 'application/pdf']) + profile.setdefault('expected_min_count', 0) + profile.setdefault('rate_limit', 1.0) + profile.setdefault('mirrors', []) + + crawl_rules = profile.get('crawl_rules', {}) + for key, default in DEFAULT_CRAWL_RULES.items(): + crawl_rules.setdefault(key, default) + profile['crawl_rules'] = crawl_rules + + return profile + + +def create_adhoc_profile(url: str, name: str) -> dict: + """Create an ad-hoc profile for a URL. + + Args: + url: Base URL to crawl + name: Name for this source + + Returns: + Profile dict + """ + from urllib.parse import urlparse + parsed = urlparse(url) + domain = parsed.netloc + + return { + 'name': name, + 'description': f"Ad-hoc source: {url}", + 'base_url': url, + 'domain_allowlist': [domain], + 'crawl_rules': dict(DEFAULT_CRAWL_RULES), + 'content_types': ['text/html', 'application/pdf'], + 'expected_min_count': 0, + 'rate_limit': 1.0, + 'mirrors': [], + } + + +def list_profiles() -> list[dict]: + """List all available source profiles. + + Returns: + List of dicts with 'name', 'description', 'path' for each profile + """ + profiles = [] + + if not PROFILES_DIR.exists(): + return profiles + + for path in sorted(PROFILES_DIR.glob('*.yaml')): + if path.stem == 'example': + continue + try: + with open(path, 'r') as f: + data = yaml.safe_load(f) + if data: + profiles.append({ + 'name': data.get('name', path.stem), + 'description': data.get('description', ''), + 'path': str(path), + 'expected_count': data.get('expected_min_count', 0), + }) + except Exception as e: + logger.warning("Failed to load profile %s: %s", path, e) + + return profiles diff --git a/collectors/wikileaks.py b/collectors/wikileaks.py new file mode 100644 index 0000000..9f7a349 --- /dev/null +++ b/collectors/wikileaks.py @@ -0,0 +1,142 @@ +"""WikiLeaks-specific crawl helpers for Mosaic collectors.""" +import re +import logging +from urllib.parse import urljoin +from typing import Optional + +logger = logging.getLogger(__name__) + + +def extract_extra_links(html: str, base_url: str, source_name: str) -> list[str]: + """Extract additional links using WikiLeaks-specific patterns. + + Different WikiLeaks collections have different page structures. + This function handles the quirks of each. + """ + links = [] + + if source_name == 'vault7': + links.extend(_vault7_links(html, base_url)) + elif source_name == 'vault8': + links.extend(_vault8_links(html, base_url)) + elif source_name == 'cablegate': + links.extend(_cablegate_links(html, base_url)) + elif source_name == 'spyfiles': + links.extend(_spyfiles_links(html, base_url)) + elif source_name == 'emails': + links.extend(_email_links(html, base_url)) + + return links + + +def _vault7_links(html: str, base_url: str) -> list[str]: + """Extract Vault 7 project page links. + + Vault 7 has a main index with project names that link to sub-pages. + Each project page has document links. + """ + links = [] + # Match project links like /vault7/#Dark Matter + for match in re.finditer(r'href=["\'](/vault7/[^"\']+)["\']', html): + link = urljoin(base_url, match.group(1)) + links.append(link) + + # Match document/file links + for match in re.finditer(r'href=["\']([^"\']*(?:\.pdf|\.zip|\.tar\.gz)[^"\']*)["\']', html, re.IGNORECASE): + link = urljoin(base_url, match.group(1)) + links.append(link) + + return links + + +def _vault8_links(html: str, base_url: str) -> list[str]: + """Extract Vault 8 source code repository links.""" + links = [] + for match in re.finditer(r'href=["\'](/vault8/[^"\']+)["\']', html): + link = urljoin(base_url, match.group(1)) + links.append(link) + return links + + +def _cablegate_links(html: str, base_url: str) -> list[str]: + """Extract Cablegate cable links with pagination support. + + Cablegate uses /plusd/cables/ with date-based navigation. + Individual cables are at /plusd/cables/CABLE_ID + """ + links = [] + + # Individual cable links + for match in re.finditer(r'href=["\'](/plusd/cables/[0-9A-Z]+)["\']', html): + link = urljoin(base_url, match.group(1)) + links.append(link) + + # Pagination links + for match in re.finditer(r'href=["\'](/plusd/[^"\']*(?:page|offset|start)=[^"\']+)["\']', html): + link = urljoin(base_url, match.group(1)) + links.append(link) + + # Date-based navigation + for match in re.finditer(r'href=["\'](/plusd/cables/\?[^"\']+)["\']', html): + link = urljoin(base_url, match.group(1)) + links.append(link) + + return links + + +def _spyfiles_links(html: str, base_url: str) -> list[str]: + """Extract Spy Files document links.""" + links = [] + for match in re.finditer(r'href=["\'](/spyfiles/[^"\']+)["\']', html): + link = urljoin(base_url, match.group(1)) + links.append(link) + return links + + +def _email_links(html: str, base_url: str) -> list[str]: + """Extract email dump links with ID-based navigation.""" + links = [] + + # Email ID links + for match in re.finditer(r'href=["\']([^"\']*emailid/\d+)["\']', html): + link = urljoin(base_url, match.group(1)) + links.append(link) + + # Pagination + for match in re.finditer(r'href=["\']([^"\']*(?:page|offset)=\d+)["\']', html): + link = urljoin(base_url, match.group(1)) + links.append(link) + + return links + + +def parse_cable_date(cable_id: str) -> Optional[str]: + """Extract date from cable ID format. + + Cable IDs often start with a two-digit year and month: + e.g., 07PARIS1234 → 2007 (approximate) + """ + match = re.match(r'^(\d{2})[A-Z]', cable_id) + if match: + year = int(match.group(1)) + # WikiLeaks cables span roughly 1966-2010 + if year >= 66: + return f"19{year:02d}" + else: + return f"20{year:02d}" + return None + + +def detect_wikileaks_source(url: str) -> Optional[str]: + """Detect which WikiLeaks collection a URL belongs to.""" + if '/vault7/' in url: + return 'vault7' + elif '/vault8/' in url: + return 'vault8' + elif '/plusd/' in url or '/cablegate/' in url: + return 'cablegate' + elif '/spyfiles/' in url: + return 'spyfiles' + elif 'emailid/' in url or '/dnc-emails/' in url or '/podesta-emails/' in url: + return 'emails' + return None