EU-Utility/core/screenshot_vulnerable.py

510 lines
16 KiB
Python

"""
EU-Utility - Screenshot Service Core Module
Fast, reliable screen capture functionality for all plugins.
Part of core - not a plugin. Plugins access via PluginAPI.
"""
import io
import os
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
try:
from PIL import Image
PIL_AVAILABLE = True
except ImportError:
PIL_AVAILABLE = False
Image = None
class ScreenshotService:
"""
Core screenshot service with cross-platform support.
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
"""
_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._save_path: Path = self._get_default_save_path()
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 = None
# Platform detection
self._platform = platform.system().lower()
self._use_pil = self._platform == "windows"
# Lazy init for capture backends
self._pil_available: Optional[bool] = None
self._pyautogui_available: Optional[bool] = None
# 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):
"""
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):
"""
Capture specific screen region.
Args:
x: Left coordinate
y: Top coordinate
width: Region width
height: Region height
Returns:
PIL Image object
"""
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[Any]:
"""
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
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
# 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.
Args:
image: PIL Image to save
filename: Optional filename (auto-generated if None)
Returns:
Path to saved file
"""
if filename is None:
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S_%f")[:-3]
filename = f"screenshot_{timestamp}.{self._format.lower()}"
# Ensure correct extension
if not filename.lower().endswith(('.png', '.jpg', '.jpeg')):
filename += f".{self._format.lower()}"
filepath = self._save_path / filename
# Save with appropriate settings
if filename.lower().endswith('.jpg') or 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._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)