pig.py/screenshot.py

#!/usr/bin/env python3
"""
Screenshot capture module for neopig.

Wraps uri2png for async-compatible page screenshots.
Screenshots are stored in vault with MD5 hash like other media.
"""

import asyncio
import hashlib
import logging
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)


@dataclass
class ScreenshotConfig:
    """Screenshot capture configuration."""
    enabled: bool = False
    width: int = 1280
    height: int = 1024
    delay: int = 1000  # ms after DOM load
    user_agent: Optional[str] = None


class ScreenshotCapture:
    """
    Async screenshot capture using uri2png.

    Since uri2png uses GTK main loop, we run it in a subprocess
    to avoid blocking the async event loop.
    """

    def __init__(self, config: ScreenshotConfig = None):
        self.config = config or ScreenshotConfig()
        self._uri2png_available = None

    async def is_available(self) -> bool:
        """Check if uri2png is installed and available."""
        if self._uri2png_available is not None:
            return self._uri2png_available

        try:
            proc = await asyncio.create_subprocess_exec(
                'python', '-c', 'from uri2png import Uri2Png',
                stdout=asyncio.subprocess.DEVNULL,
                stderr=asyncio.subprocess.DEVNULL,
            )
            await proc.wait()
            self._uri2png_available = proc.returncode == 0
        except Exception:
            self._uri2png_available = False

        if not self._uri2png_available:
            logger.warning("uri2png not available - screenshots disabled")

        return self._uri2png_available

    async def capture(self, uri: str) -> Optional[dict]:
        """
        Capture screenshot of a URI.

        Returns:
            Dict with 'data', 'md5_hash', 'mime_type' or None on failure
        """
        if not self.config.enabled:
            return None

        if not await self.is_available():
            return None

        # Create temp file for screenshot
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
            tmp_path = tmp.name

        try:
            # Build uri2png command
            cmd = [
                'python', '-c',
                f'''
from uri2png import Uri2Png
Uri2Png(
    uri="{uri}",
    filepath="{tmp_path}",
    width={self.config.width},
    height={self.config.height},
    delay={self.config.delay},
    user_agent={repr(self.config.user_agent)},
).capture()
'''
            ]

            # Run with timeout (uri2png can hang on bad URLs)
            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )

            try:
                stdout, stderr = await asyncio.wait_for(
                    proc.communicate(),
                    timeout=30.0  # 30 second timeout
                )
            except asyncio.TimeoutError:
                proc.kill()
                await proc.wait()
                logger.warning(f"Screenshot timeout: {uri}")
                return None

            if proc.returncode != 0:
                logger.warning(f"Screenshot failed ({proc.returncode}): {uri}")
                if stderr:
                    logger.debug(f"stderr: {stderr.decode()}")
                return None

            # Read screenshot data
            path = Path(tmp_path)

            def _check_file():
                if not path.exists():
                    return None
                size = path.stat().st_size
                if size == 0:
                    return None
                return path.read_bytes()

            data = await asyncio.to_thread(_check_file)
            if data is None:
                logger.warning(f"Screenshot empty: {uri}")
                return None
            md5_hash = hashlib.md5(data).hexdigest()

            logger.debug(f"Screenshot captured: {uri} -> {md5_hash}")

            return {
                'data': data,
                'md5_hash': md5_hash,
                'mime_type': 'image/png',
                'size': len(data),
                'source_uri': uri,
            }

        except Exception as e:
            logger.warning(f"Screenshot error for {uri}: {e}")
            return None

        finally:
            # Cleanup temp file (sync unlink is fine in finally - small operation)
            def _cleanup():
                try:
                    Path(tmp_path).unlink(missing_ok=True)
                except Exception:
                    pass

            try:
                await asyncio.to_thread(_cleanup)
            except Exception:
                pass

    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