""" Media download and screenshot tools """ import json import os import asyncio from typing import Dict, Any, List from urllib.parse import urljoin, urlparse from bs4 import BeautifulSoup from mcp.types import TextContent from ..config import ( PLAYWRIGHT_AVAILABLE, authenticated_session, logger ) async def _setup_context_with_session(context): """Setup Playwright context with cookies from authenticated session.""" if authenticated_session.get("active"): browser_cookies = authenticated_session.get('browser_cookies', []) if browser_cookies: # Convert cookies to Playwright format and add to context for cookie_data in browser_cookies: try: cookie = { "name": cookie_data.get('name', ''), "value": cookie_data.get('value', ''), "domain": cookie_data.get('domain', ''), "path": cookie_data.get('path', '/'), "httpOnly": cookie_data.get('httpOnly', False), "secure": cookie_data.get('secure', False) } # Add expires if present if cookie_data.get('expires'): cookie["expires"] = cookie_data['expires'] await context.add_cookies([cookie]) except Exception as cookie_error: continue async def download_media_tool(args: Dict[str, Any]) -> List[TextContent]: """Automatically download media files from a page.""" url = args["url"] types = args.get("types", ["images"]) download_path = args.get("download_path", "./downloads/") max_files = args.get("max_files", 50) min_size = args.get("min_size", 1) * 1024 # Convert KB to bytes max_size = args.get("max_size", 100) * 1024 * 1024 # Convert MB to bytes try: # Create download directory os.makedirs(download_path, exist_ok=True) # Get page HTML using efficient API with browser session cookies from ..utils.api_client import get_api_client api_client = get_api_client() # Update API client with latest browser session cookies api_client.ensure_session_updated() result_data = api_client.get(url, timeout=10, headers=headers, debug=debug) if not result_data.get("success", False): error_msg = result_data.get("error", f"HTTP {result_data.get('status_code', 0)}: Failed to fetch {url}") return [TextContent(type="text", text=error_msg)] # Extract HTML from API response html = result_data.get("content", "") if result_data.get("json_data"): # If response is JSON, try to extract HTML from it json_data = result_data["json_data"] if isinstance(json_data, dict): html = json_data.get("html", json_data.get("content", html)) soup = BeautifulSoup(html, 'lxml') downloaded_files = [] download_stats = { "images": 0, "videos": 0, "pdfs": 0, "documents": 0 } # Define selectors and extensions for different media types media_config = { "images": { "selectors": ["img[src]", "img[data-src]", "[style*='background-image']"], "extensions": [".jpg", ".jpeg", ".png", ".gif", ".webp", ".svg", ".bmp"], "attributes": ["src", "data-src", "style"] }, "videos": { "selectors": ["video[src]", "source[src]", "a[href*='.mp4']", "a[href*='.webm']"], "extensions": [".mp4", ".webm", ".avi", ".mov", ".wmv", ".flv"], "attributes": ["src", "href"] }, "pdfs": { "selectors": ["a[href*='.pdf']", "embed[src*='.pdf']", "object[data*='.pdf']"], "extensions": [".pdf"], "attributes": ["href", "src", "data"] }, "documents": { "selectors": ["a[href*='.doc']", "a[href*='.docx']", "a[href*='.xls']", "a[href*='.ppt']"], "extensions": [".doc", ".docx", ".xlsx", ".xls", ".ppt", ".pptx", ".txt"], "attributes": ["href", "src"] } } for media_type in types: if media_type not in media_config: continue config = media_config[media_type] for selector in config["selectors"]: elements = soup.select(selector) for element in elements: if len(downloaded_files) >= max_files: break media_url = None # Extract URL from different attributes for attr in config["attributes"]: if attr == "style": # Extract from background-image style style = element.get("style", "") if "background-image" in style: import re match = re.search(r'url\(["\']?([^"\']+)["\']?\)', style) if match: media_url = match.group(1) break else: media_url = element.get(attr) if media_url: break if not media_url: continue # Resolve relative URLs media_url = urljoin(url, media_url) # Check if URL has valid extension parsed_url = urlparse(media_url) path_lower = parsed_url.path.lower() if not any(path_lower.endswith(ext) for ext in config["extensions"]): continue try: # Download file using aiohttp for efficiency import aiohttp import asyncio # Get cookies from browser session for the request cookies_dict = {} if authenticated_session.get("active"): browser_cookies = authenticated_session.get('browser_cookies', []) for cookie_data in browser_cookies: cookies_dict[cookie_data.get('name', '')] = cookie_data.get('value', '') # Check file size with HEAD request async with aiohttp.ClientSession(cookies=aiohttp.CookieJar()) as session: # Add cookies to session for name, value in cookies_dict.items(): session.cookie_jar.update_cookies({name: value}, urlparse(media_url).netloc) async with session.head(media_url, timeout=aiohttp.ClientTimeout(total=5)) as response: content_length = response.headers.get('content-length') if content_length: file_size = int(content_length) if file_size < min_size or file_size > max_size: continue # Actually download the file async with session.get(media_url, timeout=aiohttp.ClientTimeout(total=30)) as response: response.raise_for_status() content = await response.read() # Generate filename filename = os.path.basename(parsed_url.path) or f"media_{len(downloaded_files)}" if not any(filename.lower().endswith(ext) for ext in config["extensions"]): # Add extension based on content-type content_type = response.headers.get('content-type', '') if media_type == "images" and "image" in content_type: filename += ".jpg" elif media_type == "videos" and "video" in content_type: filename += ".mp4" elif media_type == "pdfs": filename += ".pdf" # Save file file_path = os.path.join(download_path, filename) with open(file_path, 'wb') as f: f.write(content) download_stats[media_type] += 1 downloaded_files.append({ "filename": filename, "url": media_url, "type": media_type, "size": len(media_response.content), "path": file_path }) except Exception as e: continue result = { "source_url": url, "download_path": download_path, "media_types": types, "total_downloaded": len(downloaded_files), "download_stats": download_stats, "max_files": max_files, "size_limits": { "min_size_kb": min_size / 1024, "max_size_mb": max_size / (1024 * 1024) }, "downloaded_files": downloaded_files, "status": f"Downloaded {len(downloaded_files)} files" } return [TextContent(type="text", text=json.dumps(result, indent=2, ensure_ascii=False))] except Exception as e: return [TextContent(type="text", text=f"Media download failed: {str(e)}")] async def debug_screenshot_tool(args: Dict[str, Any]) -> List[TextContent]: """Take screenshot for debugging purposes.""" url = args["url"] path = args.get("path") full_page = args.get("full_page", True) element_selector = args.get("element_selector") wait_for = args.get("wait_for") if not PLAYWRIGHT_AVAILABLE: return [TextContent(type="text", text="Playwright not available. Install with: pip install playwright")] try: # Get unified browser session from ..utils.http_client import get_http_client http_client = await get_http_client() # Navigate to page and wait for content await http_client.ensure_session() # Get page from http_client page = http_client.page # Navigate using browser response = await page.goto(url, wait_until='domcontentloaded', timeout=30000) # Wait for element if specified if wait_for: await page.wait_for_selector(wait_for, timeout=10000) # Additional wait for dynamic content await page.wait_for_load_state('networkidle', timeout=5000) # Generate path if not provided if not path: import datetime timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") path = f"debug_screenshot_{timestamp}.png" # Take screenshot if element_selector: element = await page.locator(element_selector).first await element.screenshot(path=path) screenshot_type = "element" else: await page.screenshot(path=path, full_page=full_page) screenshot_type = "full_page" if full_page else "viewport" # Get file size file_size = os.path.getsize(path) if os.path.exists(path) else 0 result = { "url": url, "screenshot_path": path, "screenshot_type": screenshot_type, "full_page": full_page, "element_selector": element_selector, "file_size": file_size, "wait_for": wait_for, "session_active": authenticated_session.get("active", False), "browser_session_used": True, "status": "Screenshot saved successfully" } return [TextContent(type="text", text=json.dumps(result, indent=2, ensure_ascii=False))] except Exception as e: return [TextContent(type="text", text=f"Screenshot failed: {str(e)}")]