n0mad1k
130 lines
3.6 KiB
Python
130 lines
3.6 KiB
Python
"""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
|