""" EU-Utility - Screenshot Service Core Module (Security Hardened) Fast, reliable screen capture functionality with path validation. """ import io import os import time import platform import threading from collections import deque from datetime import datetime from pathlib import Path from typing import Dict, List, Optional, Tuple, Any, Union from PIL import Image from core.security_utils import PathValidator, InputValidator, SecurityError class ScreenshotService: """ Core screenshot service with cross-platform support (Security Hardened). Features: - Singleton pattern for single instance across app - Fast screen capture using PIL.ImageGrab (Windows) or pyautogui (cross-platform) - Configurable auto-save with timestamps - Screenshot history (last 20 in memory) - PNG by default, JPG quality settings - Thread-safe operations - Path traversal protection """ _instance = None _lock = threading.Lock() def __new__(cls): if cls._instance is None: with cls._lock: if cls._instance is None: cls._instance = super().__new__(cls) cls._instance._initialized = False return cls._instance def __init__(self): if self._initialized: return self._initialized = True self._lock = threading.Lock() # Configuration self._auto_save: bool = True self._format: str = "PNG" self._quality: int = 95 # For JPEG self._history_size: int = 20 # Screenshot history (thread-safe deque) self._history: deque = deque(maxlen=self._history_size) self._last_screenshot: Optional[Image.Image] = None # Platform detection - MUST be before _get_default_save_path() self._platform = platform.system().lower() self._use_pil = self._platform == "windows" # Set save path AFTER platform detection self._save_path: Path = self._get_default_save_path() # Lazy init for capture backends self._pil_available: Optional[bool] = None self._pyautogui_available: Optional[bool] = None # Resolve base path for validation self._base_save_path = self._save_path.resolve() # Ensure save directory exists self._ensure_save_directory() print(f"[Screenshot] Service initialized (auto_save={self._auto_save}, format={self._format})") def _get_default_save_path(self) -> Path: """Get default save path for screenshots.""" # Use Documents/Entropia Universe/Screenshots/ as default if self._platform == "windows": documents = Path.home() / "Documents" else: documents = Path.home() / "Documents" return documents / "Entropia Universe" / "Screenshots" def _ensure_save_directory(self) -> None: """Ensure the save directory exists.""" try: self._save_path.mkdir(parents=True, exist_ok=True) except Exception as e: print(f"[Screenshot] Warning: Could not create save directory: {e}") def _check_pil_grab(self) -> bool: """Check if PIL.ImageGrab is available.""" if self._pil_available is not None: return self._pil_available try: from PIL import ImageGrab self._pil_available = True return True except ImportError: self._pil_available = False return False def _check_pyautogui(self) -> bool: """Check if pyautogui is available.""" if self._pyautogui_available is not None: return self._pyautogui_available try: import pyautogui self._pyautogui_available = True return True except ImportError: self._pyautogui_available = False return False def capture(self, full_screen: bool = True) -> Image.Image: """ Capture screenshot. Args: full_screen: If True, capture entire screen. If False, use default region. Returns: PIL Image object Raises: RuntimeError: If no capture backend is available """ with self._lock: screenshot = self._do_capture(full_screen=full_screen) # Store in history self._last_screenshot = screenshot.copy() self._history.append({ 'image': screenshot.copy(), 'timestamp': datetime.now(), 'region': None if full_screen else 'custom' }) # Auto-save if enabled if self._auto_save: self._auto_save_screenshot(screenshot) return screenshot def capture_region(self, x: int, y: int, width: int, height: int) -> Image.Image: """ Capture specific screen region. Args: x: Left coordinate y: Top coordinate width: Region width height: Region height Returns: PIL Image object Raises: SecurityError: If region parameters are invalid """ # Validate region parameters from core.security_utils import InputValidator InputValidator.validate_region_coordinates(x, y, width, height) with self._lock: screenshot = self._do_capture(region=(x, y, x + width, y + height)) # Store in history self._last_screenshot = screenshot.copy() self._history.append({ 'image': screenshot.copy(), 'timestamp': datetime.now(), 'region': (x, y, width, height) }) # Auto-save if enabled if self._auto_save: self._auto_save_screenshot(screenshot) return screenshot def capture_window(self, window_handle: int) -> Optional[Image.Image]: """ Capture specific window by handle (Windows only). Args: window_handle: Window handle (HWND on Windows) Returns: PIL Image object or None if capture failed """ if self._platform != "windows": print("[Screenshot] capture_window is Windows-only") return None # Validate window handle if not isinstance(window_handle, int) or window_handle <= 0: print("[Screenshot] Invalid window handle") return None try: import win32gui import win32ui import win32con from ctypes import windll # Get window dimensions left, top, right, bottom = win32gui.GetWindowRect(window_handle) width = right - left height = bottom - top # Sanity check dimensions if width <= 0 or height <= 0 or width > 7680 or height > 4320: print("[Screenshot] Invalid window dimensions") return None # Create device context hwndDC = win32gui.GetWindowDC(window_handle) mfcDC = win32ui.CreateDCFromHandle(hwndDC) saveDC = mfcDC.CreateCompatibleDC() # Create bitmap saveBitMap = win32ui.CreateBitmap() saveBitMap.CreateCompatibleBitmap(mfcDC, width, height) saveDC.SelectObject(saveBitMap) # Copy screen into bitmap result = windll.user32.PrintWindow(window_handle, saveDC.GetSafeHdc(), 3) # Convert to PIL Image bmpinfo = saveBitMap.GetInfo() bmpstr = saveBitMap.GetBitmapBits(True) screenshot = Image.frombuffer( 'RGB', (bmpinfo['bmWidth'], bmpinfo['bmHeight']), bmpstr, 'raw', 'BGRX', 0, 1 ) # Cleanup win32gui.DeleteObject(saveBitMap.GetHandle()) saveDC.DeleteDC() mfcDC.DeleteDC() win32gui.ReleaseDC(window_handle, hwndDC) if result != 1: return None with self._lock: self._last_screenshot = screenshot.copy() self._history.append({ 'image': screenshot.copy(), 'timestamp': datetime.now(), 'region': 'window', 'window_handle': window_handle }) if self._auto_save: self._auto_save_screenshot(screenshot) return screenshot except Exception as e: print(f"[Screenshot] Window capture failed: {e}") return None def _do_capture(self, full_screen: bool = True, region: Optional[Tuple[int, int, int, int]] = None) -> Image.Image: """Internal capture method.""" # Try PIL.ImageGrab first (Windows, faster) if self._use_pil and self._check_pil_grab(): from PIL import ImageGrab if region: return ImageGrab.grab(bbox=region) else: return ImageGrab.grab() # Fall back to pyautogui (cross-platform) if self._check_pyautogui(): import pyautogui if region: x1, y1, x2, y2 = region return pyautogui.screenshot(region=(x1, y1, x2 - x1, y2 - y1)) else: return pyautogui.screenshot() raise RuntimeError( "No screenshot backend available. " "Install pillow (Windows) or pyautogui (cross-platform)." ) def _auto_save_screenshot(self, image: Image.Image) -> Optional[Path]: """Automatically save screenshot with timestamp.""" try: timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S_%f")[:-3] filename = f"screenshot_{timestamp}.{self._format.lower()}" return self.save_screenshot(image, filename) except Exception as e: print(f"[Screenshot] Auto-save failed: {e}") return None def save_screenshot(self, image: Image.Image, filename: Optional[str] = None) -> Path: """ Save screenshot to file with path validation. Args: image: PIL Image to save filename: Optional filename (auto-generated if None) Returns: Path to saved file Raises: SecurityError: If filename is invalid """ if filename is None: timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S_%f")[:-3] filename = f"screenshot_{timestamp}.{self._format.lower()}" # Sanitize filename safe_filename = PathValidator.sanitize_filename(filename, '_') # Ensure correct extension if not safe_filename.lower().endswith(('.png', '.jpg', '.jpeg')): safe_filename += f".{self._format.lower()}" filepath = self._save_path / safe_filename # Security check: ensure resolved path is within save_path try: resolved_path = filepath.resolve() if not str(resolved_path).startswith(str(self._base_save_path)): raise SecurityError("Path traversal detected in filename") except (OSError, ValueError) as e: print(f"[Screenshot] Security error: {e}") # Fallback to safe default safe_filename = f"screenshot_{int(time.time())}.{self._format.lower()}" filepath = self._save_path / safe_filename # Save with appropriate settings if safe_filename.lower().endswith('.jpg') or safe_filename.lower().endswith('.jpeg'): image = image.convert('RGB') # JPEG doesn't support alpha image.save(filepath, 'JPEG', quality=self._quality, optimize=True) else: image.save(filepath, 'PNG', optimize=True) return filepath def get_last_screenshot(self) -> Optional[Image.Image]: """ Get the most recent screenshot. Returns: PIL Image or None if no screenshots taken yet """ with self._lock: return self._last_screenshot.copy() if self._last_screenshot else None def get_history(self, limit: Optional[int] = None) -> List[Dict[str, Any]]: """ Get screenshot history. Args: limit: Maximum number of entries (default: all) Returns: List of dicts with 'timestamp', 'region', 'image' keys """ with self._lock: history = list(self._history) if limit: history = history[-limit:] return [ { 'timestamp': entry['timestamp'], 'region': entry['region'], 'image': entry['image'].copy() } for entry in history ] def clear_history(self) -> None: """Clear screenshot history from memory.""" with self._lock: self._history.clear() self._last_screenshot = None # ========== Configuration ========== @property def auto_save(self) -> bool: """Get auto-save setting.""" return self._auto_save @auto_save.setter def auto_save(self, value: bool) -> None: """Set auto-save setting.""" self._auto_save = bool(value) @property def save_path(self) -> Path: """Get current save path.""" return self._save_path @save_path.setter def save_path(self, path: Union[str, Path]) -> None: """Set save path.""" self._save_path = Path(path) self._base_save_path = self._save_path.resolve() self._ensure_save_directory() @property def format(self) -> str: """Get image format (PNG or JPEG).""" return self._format @format.setter def format(self, fmt: str) -> None: """Set image format.""" fmt = fmt.upper() if fmt in ('PNG', 'JPG', 'JPEG'): self._format = 'PNG' if fmt == 'PNG' else 'JPEG' else: raise ValueError(f"Unsupported format: {fmt}") @property def quality(self) -> int: """Get JPEG quality (1-100).""" return self._quality @quality.setter def quality(self, value: int) -> None: """Set JPEG quality (1-100).""" self._quality = max(1, min(100, int(value))) def configure(self, auto_save: Optional[bool] = None, save_path: Optional[Union[str, Path]] = None, format: Optional[str] = None, quality: Optional[int] = None) -> Dict[str, Any]: """ Configure screenshot service settings. Args: auto_save: Enable/disable auto-save save_path: Directory to save screenshots format: Image format (PNG or JPEG) quality: JPEG quality (1-100) Returns: Current configuration as dict """ if auto_save is not None: self.auto_save = auto_save if save_path is not None: self.save_path = save_path if format is not None: self.format = format if quality is not None: self.quality = quality return self.get_config() def get_config(self) -> Dict[str, Any]: """Get current configuration.""" return { 'auto_save': self._auto_save, 'save_path': str(self._save_path), 'format': self._format, 'quality': self._quality, 'history_size': self._history_size, 'platform': self._platform, 'backend': 'PIL' if self._use_pil else 'pyautogui' } # ========== Utility Methods ========== def image_to_bytes(self, image: Image.Image, format: Optional[str] = None) -> bytes: """ Convert PIL Image to bytes. Args: image: PIL Image format: Output format (default: current format setting) Returns: Image as bytes """ fmt = (format or self._format).upper() buffer = io.BytesIO() if fmt == 'JPEG': image = image.convert('RGB') image.save(buffer, 'JPEG', quality=self._quality) else: image.save(buffer, 'PNG') return buffer.getvalue() def get_available_backends(self) -> List[str]: """Get list of available capture backends.""" backends = [] if self._check_pil_grab(): backends.append('PIL.ImageGrab') if self._check_pyautogui(): backends.append('pyautogui') return backends def is_available(self) -> bool: """Check if screenshot service is available (has working backend).""" return self._check_pil_grab() or self._check_pyautogui() # Singleton instance _screenshot_service = None def get_screenshot_service() -> ScreenshotService: """Get the global ScreenshotService instance.""" global _screenshot_service if _screenshot_service is None: _screenshot_service = ScreenshotService() return _screenshot_service # Convenience functions for quick screenshots def quick_capture() -> Image.Image: """Quick full-screen capture.""" return get_screenshot_service().capture(full_screen=True) def quick_capture_region(x: int, y: int, width: int, height: int) -> Image.Image: """Quick region capture.""" return get_screenshot_service().capture_region(x, y, width, height) def quick_save(filename: Optional[str] = None) -> Path: """Quick capture and save.""" service = get_screenshot_service() image = service.capture() return service.save_screenshot(image, filename)