n0mad1k
228 lines
8.7 KiB
Python
228 lines
8.7 KiB
Python
"""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
|