Add collector engine: base, custom source, profile loader, wikileaks helpers
This commit is contained in:
n0mad1k
@@ -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
|
||||
Reference in New Issue
Block a user