"""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 # Scope check: only crawl URLs under the base_url path prefix # This prevents BFS from wandering the entire domain base_path = urlparse(self.base_url).path.rstrip('/') url_path = urlparse(url).path.rstrip('/') if base_path and not url_path.startswith(base_path) and url != self.base_url: # Allow the base URL itself, but skip other paths outside our scope # unless they match a URL pattern (e.g., linked documents) if not self._matches_url_pattern(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() # Only collect URLs that match our patterns or are the right content type if self._matches_url_pattern(url): 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