202 lines
6.7 KiB
Python
202 lines
6.7 KiB
Python
"""
|
|
Screenshot capture module for SneakyScanner.
|
|
|
|
Uses Playwright with Chromium to capture screenshots of discovered web services.
|
|
"""
|
|
|
|
import os
|
|
import logging
|
|
from pathlib import Path
|
|
from playwright.sync_api import sync_playwright, TimeoutError as PlaywrightTimeout
|
|
|
|
|
|
class ScreenshotCapture:
|
|
"""
|
|
Handles webpage screenshot capture for web services discovered during scanning.
|
|
|
|
Uses Playwright with Chromium in headless mode to capture viewport screenshots
|
|
of HTTP and HTTPS services. Handles SSL certificate errors gracefully.
|
|
"""
|
|
|
|
def __init__(self, output_dir, scan_timestamp, timeout=15, viewport=None):
|
|
"""
|
|
Initialize the screenshot capture handler.
|
|
|
|
Args:
|
|
output_dir (str): Base output directory for scan reports
|
|
scan_timestamp (str): Timestamp string for this scan (format: YYYYMMDD_HHMMSS)
|
|
timeout (int): Timeout in seconds for page load and screenshot (default: 15)
|
|
viewport (dict): Viewport size dict with 'width' and 'height' keys
|
|
(default: {'width': 1280, 'height': 720})
|
|
"""
|
|
self.output_dir = output_dir
|
|
self.scan_timestamp = scan_timestamp
|
|
self.timeout = timeout * 1000 # Convert to milliseconds for Playwright
|
|
self.viewport = viewport or {'width': 1280, 'height': 720}
|
|
|
|
self.playwright = None
|
|
self.browser = None
|
|
self.screenshot_dir = None
|
|
|
|
# Set up logging
|
|
self.logger = logging.getLogger('SneakyScanner.Screenshot')
|
|
|
|
def _get_screenshot_dir(self):
|
|
"""
|
|
Create and return the screenshots subdirectory for this scan.
|
|
|
|
Returns:
|
|
Path: Path object for the screenshots directory
|
|
"""
|
|
if self.screenshot_dir is None:
|
|
dir_name = f"scan_report_{self.scan_timestamp}_screenshots"
|
|
self.screenshot_dir = Path(self.output_dir) / dir_name
|
|
self.screenshot_dir.mkdir(parents=True, exist_ok=True)
|
|
self.logger.info(f"Created screenshot directory: {self.screenshot_dir}")
|
|
|
|
return self.screenshot_dir
|
|
|
|
def _generate_filename(self, ip, port):
|
|
"""
|
|
Generate a filename for the screenshot.
|
|
|
|
Args:
|
|
ip (str): IP address of the service
|
|
port (int): Port number of the service
|
|
|
|
Returns:
|
|
str: Filename in format: {ip}_{port}.png
|
|
"""
|
|
# Replace dots in IP with underscores for filesystem compatibility
|
|
safe_ip = ip.replace('.', '_')
|
|
return f"{safe_ip}_{port}.png"
|
|
|
|
def _launch_browser(self):
|
|
"""
|
|
Launch Playwright and Chromium browser in headless mode.
|
|
|
|
Returns:
|
|
bool: True if browser launched successfully, False otherwise
|
|
"""
|
|
if self.browser is not None:
|
|
return True # Already launched
|
|
|
|
try:
|
|
self.logger.info("Launching Chromium browser...")
|
|
self.playwright = sync_playwright().start()
|
|
self.browser = self.playwright.chromium.launch(
|
|
headless=True,
|
|
args=[
|
|
'--no-sandbox',
|
|
'--disable-setuid-sandbox',
|
|
'--disable-dev-shm-usage',
|
|
'--disable-gpu',
|
|
]
|
|
)
|
|
self.logger.info("Chromium browser launched successfully")
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to launch browser: {e}")
|
|
return False
|
|
|
|
def _close_browser(self):
|
|
"""
|
|
Close the browser and cleanup Playwright resources.
|
|
"""
|
|
if self.browser:
|
|
try:
|
|
self.browser.close()
|
|
self.logger.info("Browser closed")
|
|
except Exception as e:
|
|
self.logger.warning(f"Error closing browser: {e}")
|
|
finally:
|
|
self.browser = None
|
|
|
|
if self.playwright:
|
|
try:
|
|
self.playwright.stop()
|
|
except Exception as e:
|
|
self.logger.warning(f"Error stopping playwright: {e}")
|
|
finally:
|
|
self.playwright = None
|
|
|
|
def capture(self, ip, port, protocol):
|
|
"""
|
|
Capture a screenshot of a web service.
|
|
|
|
Args:
|
|
ip (str): IP address of the service
|
|
port (int): Port number of the service
|
|
protocol (str): Protocol to use ('http' or 'https')
|
|
|
|
Returns:
|
|
str: Relative path to the screenshot file, or None if capture failed
|
|
"""
|
|
# Validate protocol
|
|
if protocol not in ['http', 'https']:
|
|
self.logger.warning(f"Invalid protocol '{protocol}' for {ip}:{port}")
|
|
return None
|
|
|
|
# Launch browser if not already running
|
|
if not self._launch_browser():
|
|
return None
|
|
|
|
# Build URL
|
|
url = f"{protocol}://{ip}:{port}"
|
|
|
|
# Generate screenshot filename
|
|
filename = self._generate_filename(ip, port)
|
|
screenshot_dir = self._get_screenshot_dir()
|
|
screenshot_path = screenshot_dir / filename
|
|
|
|
try:
|
|
self.logger.info(f"Capturing screenshot: {url}")
|
|
|
|
# Create new browser context with viewport and SSL settings
|
|
context = self.browser.new_context(
|
|
viewport=self.viewport,
|
|
ignore_https_errors=True, # Handle self-signed certs
|
|
user_agent='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
|
|
)
|
|
|
|
# Create new page
|
|
page = context.new_page()
|
|
|
|
# Set default timeout
|
|
page.set_default_timeout(self.timeout)
|
|
|
|
# Navigate to URL
|
|
page.goto(url, wait_until='networkidle', timeout=self.timeout)
|
|
|
|
# Take screenshot (viewport only)
|
|
page.screenshot(path=str(screenshot_path), type='png')
|
|
|
|
# Close page and context
|
|
page.close()
|
|
context.close()
|
|
|
|
self.logger.info(f"Screenshot saved: {screenshot_path}")
|
|
|
|
# Return relative path (relative to output directory)
|
|
relative_path = f"{screenshot_dir.name}/{filename}"
|
|
return relative_path
|
|
|
|
except PlaywrightTimeout:
|
|
self.logger.warning(f"Timeout capturing screenshot for {url}")
|
|
return None
|
|
|
|
except Exception as e:
|
|
self.logger.warning(f"Failed to capture screenshot for {url}: {e}")
|
|
return None
|
|
|
|
def __enter__(self):
|
|
"""Context manager entry."""
|
|
self._launch_browser()
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Context manager exit - cleanup browser resources."""
|
|
self._close_browser()
|
|
return False # Don't suppress exceptions
|