405 lines
16 KiB
Python
405 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
# Side quest 13/21: A picture is worth a thousand words. We keep millions.
|
|
"""
|
|
Screenshot capture module for neopig.
|
|
|
|
Wraps uri2png for async-compatible page screenshots.
|
|
Screenshots are stored in vault with MD5 hash like other media.
|
|
|
|
Supported engines (in order of preference):
|
|
- wkhtmltoimage: Fast, lightweight, uses Qt WebKit. Install: apt install wkhtmltopdf
|
|
- cutycapt: Fast, lightweight, uses Qt WebKit. Install: apt install cutycapt
|
|
- playwright-webkit: WebKit via Playwright (lighter than Chromium)
|
|
- playwright-firefox: Firefox via Playwright
|
|
- playwright-chromium: Chromium via Playwright (heaviest, but most compatible)
|
|
- selenium-*: Various Selenium drivers
|
|
|
|
The module auto-detects available engines and picks the lightest one,
|
|
or you can specify an engine explicitly.
|
|
"""
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import io
|
|
import logging
|
|
import shutil
|
|
import warnings
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Optional, List, Dict, Any
|
|
import tempfile
|
|
|
|
from PIL import Image
|
|
|
|
# Suppress PIL decompression bomb warnings for large screenshots
|
|
warnings.filterwarnings('ignore', category=Image.DecompressionBombWarning)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Engine preference order - lightest/fastest first
|
|
ENGINE_PREFERENCE = [
|
|
'cutycapt', # Native Qt WebKit - fast, supports full-page screenshots
|
|
'wkhtmltoimage', # Native Qt WebKit - fast, viewport-only
|
|
'playwright-webkit', # WebKit via Playwright - lighter than Chromium
|
|
'playwright-firefox',
|
|
'playwright-chromium',
|
|
'playwright', # Default Playwright (Chromium)
|
|
'selenium-chrome',
|
|
'selenium-firefox',
|
|
'selenium',
|
|
]
|
|
|
|
# Native tools that don't require browser downloads
|
|
NATIVE_ENGINES = {'wkhtmltoimage', 'cutycapt'}
|
|
|
|
|
|
@dataclass
|
|
class ScreenshotConfig:
|
|
"""Screenshot capture configuration."""
|
|
enabled: bool = True
|
|
width: int = 1024
|
|
height: int = 768
|
|
delay: int = 1000 # ms after DOM load
|
|
timeout: int = 30000 # ms total timeout
|
|
user_agent: Optional[str] = None
|
|
engine: Optional[str] = None # None = auto-detect best available
|
|
full_page: bool = True
|
|
format: str = 'jpeg' # 'jpeg' or 'png'
|
|
quality: int = 93 # JPEG quality (1-100)
|
|
|
|
|
|
class ScreenshotCapture:
|
|
"""
|
|
Async screenshot capture using uri2png.
|
|
|
|
Supports multiple backends with automatic selection of the lightest
|
|
available engine. Native tools (wkhtmltoimage, cutycapt) are preferred
|
|
over browser-based solutions.
|
|
|
|
Usage:
|
|
config = ScreenshotConfig(enabled=True, engine='wkhtmltoimage')
|
|
capture = ScreenshotCapture(config)
|
|
result = await capture.capture('https://example.com')
|
|
"""
|
|
|
|
def __init__(self, config: ScreenshotConfig = None):
|
|
self.config = config or ScreenshotConfig()
|
|
self._engine = None
|
|
self._engine_name = None
|
|
self._available_engines: Optional[List[Dict[str, str]]] = None
|
|
self._initialized = False
|
|
|
|
async def _get_available_engines(self) -> List[Dict[str, str]]:
|
|
"""Get list of available screenshot engines."""
|
|
if self._available_engines is not None:
|
|
return self._available_engines
|
|
|
|
def _check():
|
|
try:
|
|
from uri2png import get_available_engines
|
|
return get_available_engines()
|
|
except ImportError:
|
|
return []
|
|
|
|
self._available_engines = await asyncio.to_thread(_check)
|
|
return self._available_engines
|
|
|
|
async def _select_engine(self) -> Optional[str]:
|
|
"""Select the best available engine based on preference order."""
|
|
available = await self._get_available_engines()
|
|
available_names = {e['name'] for e in available}
|
|
|
|
# If user specified an engine, try to use it
|
|
if self.config.engine:
|
|
if self.config.engine in available_names:
|
|
return self.config.engine
|
|
else:
|
|
logger.warning(f"Requested engine '{self.config.engine}' not available")
|
|
logger.info(f"Available engines: {', '.join(available_names)}")
|
|
|
|
# Check native tools first (they're fast and don't need browser downloads)
|
|
for engine in ENGINE_PREFERENCE:
|
|
if engine in available_names:
|
|
# For native engines, verify the binary exists
|
|
if engine in NATIVE_ENGINES:
|
|
binary = 'wkhtmltoimage' if engine == 'wkhtmltoimage' else 'cutycapt'
|
|
if shutil.which(binary):
|
|
return engine
|
|
else:
|
|
logger.debug(f"Engine {engine} listed but binary not found")
|
|
continue
|
|
return engine
|
|
|
|
return None
|
|
|
|
async def initialize(self) -> bool:
|
|
"""Initialize the screenshot engine."""
|
|
if self._initialized:
|
|
return self._engine is not None
|
|
|
|
engine_name = await self._select_engine()
|
|
if not engine_name:
|
|
logger.warning("No screenshot engine available")
|
|
logger.info("Install one of: wkhtmltopdf, cutycapt, or playwright")
|
|
self._initialized = True
|
|
return False
|
|
|
|
def _create_engine():
|
|
try:
|
|
from uri2png import create_engine
|
|
# Pass options as kwargs to create_engine
|
|
return create_engine(
|
|
engine_name,
|
|
width=self.config.width,
|
|
height=self.config.height,
|
|
delay=self.config.delay,
|
|
timeout=self.config.timeout,
|
|
full_page=self.config.full_page,
|
|
user_agent=self.config.user_agent,
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to create engine '{engine_name}': {e}")
|
|
return None
|
|
|
|
self._engine = await asyncio.to_thread(_create_engine)
|
|
self._engine_name = engine_name
|
|
self._initialized = True
|
|
|
|
if self._engine:
|
|
logger.info(f"Screenshot engine: {engine_name}")
|
|
return True
|
|
return False
|
|
|
|
async def is_available(self) -> bool:
|
|
"""Check if screenshots are available."""
|
|
if not self._initialized:
|
|
await self.initialize()
|
|
return self._engine is not None
|
|
|
|
def get_engine_name(self) -> Optional[str]:
|
|
"""Get the name of the active engine."""
|
|
return self._engine_name
|
|
|
|
async def list_engines(self) -> List[Dict[str, str]]:
|
|
"""List all available screenshot engines."""
|
|
return await self._get_available_engines()
|
|
|
|
@staticmethod
|
|
def calculate_delay_for_content(content_length: int) -> int:
|
|
"""
|
|
Calculate optimal screenshot delay based on content size.
|
|
|
|
Longer pages need more time for JavaScript to render all content.
|
|
|
|
Args:
|
|
content_length: Length of raw HTML content in bytes
|
|
|
|
Returns:
|
|
Delay in milliseconds (min 3000ms, max 60000ms)
|
|
"""
|
|
# Base delay of 3 seconds
|
|
base_delay = 3000
|
|
|
|
# Add 3 seconds per 50KB of content, up to a max of 57 seconds
|
|
additional_delay = min((content_length // 50000) * 3000, 57000)
|
|
|
|
return base_delay + additional_delay
|
|
|
|
async def capture(self, uri: str, content_length: int = 0) -> Optional[dict]:
|
|
"""
|
|
Capture screenshot of a URI.
|
|
|
|
Args:
|
|
uri: URL to capture
|
|
content_length: Optional content length for dynamic delay calculation.
|
|
If provided and engine supports it, delay will be adjusted.
|
|
|
|
Returns:
|
|
Dict with 'data', 'md5_hash', 'mime_type', 'engine' or None on failure
|
|
"""
|
|
if not self.config.enabled:
|
|
return None
|
|
|
|
if not await self.is_available():
|
|
return None
|
|
|
|
# Calculate dynamic delay if content_length provided and we're using a native engine
|
|
effective_engine = self._engine
|
|
if content_length > 0 and self._engine_name in NATIVE_ENGINES:
|
|
dynamic_delay = self.calculate_delay_for_content(content_length)
|
|
if dynamic_delay > self.config.delay:
|
|
logger.debug(f"Using dynamic delay {dynamic_delay}ms for {content_length} bytes")
|
|
# Create a one-off engine with the adjusted delay
|
|
try:
|
|
from uri2png import create_engine
|
|
effective_engine = await asyncio.to_thread(
|
|
lambda: create_engine(
|
|
self._engine_name,
|
|
width=self.config.width,
|
|
height=self.config.height,
|
|
delay=dynamic_delay,
|
|
timeout=self.config.timeout,
|
|
full_page=self.config.full_page,
|
|
user_agent=self.config.user_agent,
|
|
)
|
|
)
|
|
except Exception as e:
|
|
logger.debug(f"Could not create dynamic engine: {e}")
|
|
effective_engine = self._engine
|
|
|
|
try:
|
|
# Create temp file for output
|
|
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
|
|
output_path = f.name
|
|
|
|
try:
|
|
# Capture - API is capture(url, output_path)
|
|
capture_coro = effective_engine.capture(uri, output_path)
|
|
|
|
# All uri2png engines return coroutines
|
|
result = await asyncio.wait_for(
|
|
capture_coro,
|
|
timeout=self.config.timeout / 1000 + 5
|
|
)
|
|
|
|
if not result.success:
|
|
# Extract meaningful error from wkhtmltoimage output
|
|
raw_error = result.error or 'unknown error'
|
|
# Filter out wkhtmltoimage progress output (progress bars, loading messages)
|
|
lines = raw_error.split('\n')
|
|
# Skip lines that are progress bars or loading messages
|
|
meaningful = [l.strip() for l in lines if l.strip()
|
|
and not l.strip().startswith('[') # Progress bars like [>
|
|
and not l.strip().startswith('Loading')
|
|
and '%' not in l] # Percentage indicators
|
|
# Look for actual error/fail lines first
|
|
error_lines = [l for l in meaningful if 'error' in l.lower() or 'fail' in l.lower()]
|
|
if error_lines:
|
|
error_msg = error_lines[0][:60]
|
|
elif meaningful:
|
|
error_msg = meaningful[-1][:60]
|
|
else:
|
|
error_msg = 'capture failed'
|
|
logger.warning(f"Screenshot failed: {uri} - {error_msg}")
|
|
return None
|
|
|
|
# Read bytes from output file
|
|
png_data = Path(output_path).read_bytes()
|
|
if not png_data:
|
|
logger.warning(f"Screenshot empty: {uri}")
|
|
return None
|
|
|
|
# Convert to JPEG if configured
|
|
if self.config.format == 'jpeg':
|
|
img = Image.open(io.BytesIO(png_data))
|
|
|
|
# Convert RGBA to RGB (JPEG doesn't support alpha)
|
|
if img.mode in ('RGBA', 'LA', 'P'):
|
|
img = img.convert('RGB')
|
|
|
|
# Split oversized images (JPEG max is 65500px)
|
|
max_dim = 65000
|
|
if img.height > max_dim:
|
|
logger.debug(f"Splitting oversized image: {img.height}px into chunks")
|
|
results = []
|
|
chunk_idx = 0
|
|
y = 0
|
|
while y < img.height:
|
|
chunk_height = min(max_dim, img.height - y)
|
|
chunk = img.crop((0, y, img.width, y + chunk_height))
|
|
output = io.BytesIO()
|
|
chunk.save(output, format='JPEG', quality=self.config.quality, optimize=True)
|
|
chunk_data = output.getvalue()
|
|
chunk_hash = hashlib.md5(chunk_data).hexdigest()
|
|
results.append({
|
|
'data': chunk_data,
|
|
'md5_hash': chunk_hash,
|
|
'mime_type': 'image/jpeg',
|
|
'size': len(chunk_data),
|
|
'source_uri': uri,
|
|
'engine': self._engine_name,
|
|
'format': 'jpg',
|
|
'chunk': chunk_idx,
|
|
'total_chunks': (img.height + max_dim - 1) // max_dim,
|
|
})
|
|
y += max_dim
|
|
chunk_idx += 1
|
|
return results # Return list for oversized images
|
|
|
|
output = io.BytesIO()
|
|
img.save(output, format='JPEG', quality=self.config.quality, optimize=True)
|
|
data = output.getvalue()
|
|
mime_type = 'image/jpeg'
|
|
ext = 'jpg'
|
|
else:
|
|
data = png_data
|
|
mime_type = 'image/png'
|
|
ext = 'png'
|
|
|
|
md5_hash = hashlib.md5(data).hexdigest()
|
|
logger.debug(f"Screenshot captured ({self._engine_name}, {ext}): {uri} -> {md5_hash} ({len(data)//1024}KB)")
|
|
|
|
return {
|
|
'data': data,
|
|
'md5_hash': md5_hash,
|
|
'mime_type': mime_type,
|
|
'size': len(data),
|
|
'source_uri': uri,
|
|
'engine': self._engine_name,
|
|
'format': ext,
|
|
}
|
|
finally:
|
|
# Clean up temp file
|
|
try:
|
|
Path(output_path).unlink(missing_ok=True)
|
|
except Exception:
|
|
pass
|
|
|
|
except asyncio.TimeoutError:
|
|
logger.warning(f"Screenshot timeout: {uri}")
|
|
return None
|
|
except Exception as e:
|
|
logger.warning(f"Screenshot error for {uri}: {e}")
|
|
return None
|
|
|
|
async def capture_to_file(self, uri: str, output_path: str) -> bool:
|
|
"""
|
|
Capture screenshot directly to a file.
|
|
|
|
Returns:
|
|
True on success, False on failure
|
|
"""
|
|
result = await self.capture(uri)
|
|
if result is None:
|
|
return False
|
|
|
|
try:
|
|
await asyncio.to_thread(Path(output_path).write_bytes, result['data'])
|
|
return True
|
|
except Exception as e:
|
|
logger.warning(f"Failed to write screenshot: {e}")
|
|
return False
|
|
|
|
async def cleanup(self):
|
|
"""Cleanup engine resources."""
|
|
if self._engine:
|
|
try:
|
|
def _cleanup():
|
|
if hasattr(self._engine, 'cleanup'):
|
|
self._engine.cleanup()
|
|
await asyncio.to_thread(_cleanup)
|
|
except Exception as e:
|
|
logger.debug(f"Engine cleanup error: {e}")
|
|
|
|
|
|
async def list_available_engines() -> List[Dict[str, str]]:
|
|
"""List all available screenshot engines (utility function)."""
|
|
capture = ScreenshotCapture()
|
|
return await capture.list_engines()
|
|
|
|
|
|
async def get_best_engine() -> Optional[str]:
|
|
"""Get the best available engine name (utility function)."""
|
|
capture = ScreenshotCapture(ScreenshotConfig(enabled=True))
|
|
await capture.initialize()
|
|
return capture.get_engine_name()
|