"""Screenshot capture utilities for Windows."""
import io
import base64
from datetime import datetime
from pathlib import Path
from typing import Optional, Tuple
import mss
import mss.tools
from PIL import Image
class ScreenshotManager:
"""Manages screenshot capture and storage."""
def __init__(self, temp_dir: Optional[Path] = None):
"""Initialize screenshot manager.
Args:
temp_dir: Directory to store temporary screenshots
"""
self.temp_dir = temp_dir or Path.home() / ".zoho-timeline-mcp" / "screenshots"
self.temp_dir.mkdir(parents=True, exist_ok=True)
self.screenshots = []
def capture_screenshot(self, monitor_num: int = 1) -> Tuple[str, Path]:
"""Capture a screenshot of specified monitor.
Args:
monitor_num: Monitor number (1 = primary, 2+ = additional monitors)
Returns:
Tuple of (screenshot_id, file_path)
"""
with mss.mss() as sct:
# Get the monitor (0 = all monitors, 1+ = specific monitor)
monitor = sct.monitors[monitor_num]
# Capture the screenshot
sct_img = sct.grab(monitor)
# Generate filename with timestamp
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
filename = f"screenshot_{timestamp}.png"
filepath = self.temp_dir / filename
# Save screenshot
mss.tools.to_png(sct_img.rgb, sct_img.size, output=str(filepath))
# Store reference
screenshot_id = f"ss_{len(self.screenshots) + 1}"
self.screenshots.append({
"id": screenshot_id,
"path": filepath,
"timestamp": timestamp,
"monitor": monitor_num
})
return screenshot_id, filepath
def capture_active_window(self) -> Tuple[str, Path]:
"""Capture screenshot of the active window (Windows only).
Returns:
Tuple of (screenshot_id, file_path)
"""
try:
import win32gui
import win32ui
import win32con
from ctypes import windll
# Get active window
hwnd = win32gui.GetForegroundWindow()
# Get window dimensions
left, top, right, bottom = win32gui.GetWindowRect(hwnd)
width = right - left
height = bottom - top
# Get device contexts
hwndDC = win32gui.GetWindowDC(hwnd)
mfcDC = win32ui.CreateDCFromHandle(hwndDC)
saveDC = mfcDC.CreateCompatibleDC()
# Create bitmap
saveBitMap = win32ui.CreateBitmap()
saveBitMap.CreateCompatibleBitmap(mfcDC, width, height)
saveDC.SelectObject(saveBitMap)
# Copy window content
result = windll.user32.PrintWindow(hwnd, saveDC.GetSafeHdc(), 3)
# Convert to PIL Image
bmpinfo = saveBitMap.GetInfo()
bmpstr = saveBitMap.GetBitmapBits(True)
img = Image.frombuffer(
'RGB',
(bmpinfo['bmWidth'], bmpinfo['bmHeight']),
bmpstr, 'raw', 'BGRX', 0, 1
)
# Clean up
win32gui.DeleteObject(saveBitMap.GetHandle())
saveDC.DeleteDC()
mfcDC.DeleteDC()
win32gui.ReleaseDC(hwnd, hwndDC)
# Save screenshot
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
filename = f"screenshot_{timestamp}.png"
filepath = self.temp_dir / filename
img.save(filepath)
# Store reference
screenshot_id = f"ss_{len(self.screenshots) + 1}"
self.screenshots.append({
"id": screenshot_id,
"path": filepath,
"timestamp": timestamp,
"window": "active"
})
return screenshot_id, filepath
except ImportError:
# Fallback to full screen if pywin32 not available
print("pywin32 not available, falling back to full screen capture")
return self.capture_screenshot()
def get_screenshot_base64(self, screenshot_id: str) -> Optional[str]:
"""Get screenshot as base64 encoded string.
Args:
screenshot_id: ID of screenshot to retrieve
Returns:
Base64 encoded PNG image or None if not found
"""
for ss in self.screenshots:
if ss["id"] == screenshot_id:
with open(ss["path"], "rb") as f:
return base64.b64encode(f.read()).decode("utf-8")
return None
def get_all_screenshots(self):
"""Get list of all captured screenshots."""
return self.screenshots.copy()
def clear_screenshots(self):
"""Clear all captured screenshots and delete files."""
for ss in self.screenshots:
try:
ss["path"].unlink(missing_ok=True)
except Exception as e:
print(f"Error deleting {ss['path']}: {e}")
self.screenshots.clear()
def get_screenshot_count(self) -> int:
"""Get number of captured screenshots."""
return len(self.screenshots)
