""" Session management tools """ import json import time from typing import Dict, Any, List from mcp.types import TextContent from ..config import ( session, authenticated_session, memory_cache, cache_timestamps, BROWSER_COOKIE3_AVAILABLE ) from ..session.cookies import import_cookies_from_browser from ..session.manager import clear_authenticated_session async def get_session_info_tool(args: Dict[str, Any]) -> List[TextContent]: """Get current session information including cookies and auth status.""" try: cookies_info = {} for cookie in session.cookies: cookies_info[cookie.name] = { "value": cookie.value, "domain": cookie.domain, "path": cookie.path, "secure": cookie.secure, "expires": cookie.expires } auth_info = None if session.auth: auth_info = { "type": "Basic Authentication", "username": session.auth[0] } # Add authenticated session info auth_session_info = { "active": authenticated_session.get("active", False), "login_url": authenticated_session.get("login_url"), "login_timestamp": authenticated_session.get("login_timestamp"), "cookies_count": authenticated_session.get("cookies_count", 0), "last_activity": authenticated_session.get("last_activity"), "session_age_hours": ( (time.time() - authenticated_session.get("login_timestamp", time.time())) / 3600 if authenticated_session.get("login_timestamp") else 0 ) } result = { "session_cookies_count": len(cookies_info), "session_cookies": cookies_info, "basic_authentication": auth_info, "session_headers": dict(session.headers), "has_auth": bool(session.auth or cookies_info), "authenticated_session": auth_session_info, "total_session_state": { "requests_session_cookies": len(session.cookies), "authenticated_session_cookies": authenticated_session.get("cookies_count", 0), "cache_entries": len(memory_cache) } } return [TextContent(type="text", text=json.dumps(result, indent=2, ensure_ascii=False))] except Exception as e: return [TextContent(type="text", text=f"Failed to get session info: {str(e)}")] async def clear_session_tool(args: Dict[str, Any]) -> List[TextContent]: """Clear authentication session and cookies from all tools.""" confirmation = args.get("confirmation", True) clear_cache = args.get("clear_cache", False) verbose = args.get("verbose", True) if not confirmation: return [TextContent(type="text", text=json.dumps({ "action": "clear_session", "status": "cancelled", "message": "Session clearing cancelled - confirmation=false" }, indent=2, ensure_ascii=False))] try: # Track what's being cleared cleared_items = [] # Store current session info for reporting before_state = { "authenticated_session_active": authenticated_session.get("active", False), "login_url": authenticated_session.get("login_url"), "cookies_count": authenticated_session.get("cookies_count", 0), "session_cookies_count": len(session.cookies), "cache_items": len(memory_cache) if clear_cache else "not_cleared" } # Clear authenticated session state if authenticated_session.get("active"): cleared_items.append(f"Authenticated session for {authenticated_session.get('login_url', 'unknown URL')}") clear_authenticated_session() cleared_items.append("Global authenticated session state") # Clear requests session cookies if session.cookies: cookies_count = len(session.cookies) session.cookies.clear() cleared_items.append(f"Session cookies ({cookies_count} items)") # Clear session headers (except default User-Agent) current_headers = dict(session.headers) session.headers.clear() session.headers.update({ "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36" }) if len(current_headers) > 1: # More than just User-Agent cleared_items.append(f"Session headers ({len(current_headers)} items)") # Clear basic auth if session.auth: session.auth = None cleared_items.append("Basic authentication") # Clear memory cache if requested if clear_cache: cache_count = len(memory_cache) memory_cache.clear() cache_timestamps.clear() if cache_count > 0: cleared_items.append(f"Memory cache ({cache_count} items)") # Build result result = { "action": "clear_session", "status": "completed", "cleared_items": cleared_items, "before_state": before_state if verbose else {}, "after_state": { "authenticated_session_active": False, "session_cookies_count": 0, "cache_items": len(memory_cache) }, "timestamp": time.time(), "message": f"Successfully cleared {len(cleared_items)} session components" } if verbose: result["details"] = { "authenticated_session_reset": True, "requests_session_cookies_cleared": True, "session_headers_reset": True, "basic_auth_cleared": True, "memory_cache_cleared": clear_cache, "tools_affected": [ "fetch_html", "parse_page", "find_elements", "extract_text", "extract_links", "login_form", "browser_login", "auto_paginate", "parallel_fetch", "infinite_scroll", "extract_structured_data", "download_media" ] } return [TextContent(type="text", text=json.dumps(result, indent=2, ensure_ascii=False))] except Exception as e: error_result = { "action": "clear_session", "status": "error", "error": str(e), "error_type": type(e).__name__, "message": "Failed to clear session", "recommendations": [ "Check if session was already cleared", "Restart the MCP server if session state is corrupted", "Verify no other tools are currently using the session" ] } return [TextContent(type="text", text=json.dumps(error_result, indent=2, ensure_ascii=False))] async def import_browser_cookies_tool(args: Dict[str, Any]) -> List[TextContent]: """Import cookies from browser with detailed diagnostics and access analysis.""" browser = args["browser"] domain = args.get("domain") verbose = args.get("verbose", True) if not BROWSER_COOKIE3_AVAILABLE: return [TextContent(type="text", text=json.dumps({ "error": "browser-cookie3 not available", "browser": browser, "domain": domain, "cookies_imported": 0, "diagnostics": { "library_available": False, "error_details": "browser-cookie3 library not installed" }, "recommendations": [ "Install browser-cookie3: pip install browser-cookie3", "Ensure browser is properly installed and has been used", "Check if browser data directory is accessible" ] }, indent=2, ensure_ascii=False))] try: result = await import_cookies_from_browser(browser, domain, verbose) return [TextContent(type="text", text=json.dumps(result, indent=2, ensure_ascii=False))] except Exception as e: error_result = { "error": str(e), "error_type": type(e).__name__, "browser": browser, "domain": domain, "cookies_imported": 0, "diagnostics": { "library_available": BROWSER_COOKIE3_AVAILABLE, "error_details": f"Critical error during cookie import: {str(e)}" }, "recommendations": [ f"Check if {browser} browser is installed and has been used", "Ensure browser is closed before importing cookies", "Verify file permissions for browser data directory", "Try different browser or check browser version compatibility" ] } return [TextContent(type="text", text=json.dumps(error_result, indent=2, ensure_ascii=False))] async def set_basic_auth_tool(args: Dict[str, Any]) -> List[TextContent]: """Set HTTP Basic Authentication for all subsequent requests.""" username = args["username"] password = args["password"] try: # Set basic auth for the session session.auth = (username, password) result = { "username": username, "auth_type": "Basic Authentication", "status": "Authentication set for all subsequent requests" } return [TextContent(type="text", text=json.dumps(result, indent=2, ensure_ascii=False))] except Exception as e: return [TextContent(type="text", text=f"Failed to set basic auth: {str(e)}")] async def set_oauth_token_tool(args: Dict[str, Any]) -> List[TextContent]: """Set OAuth/JWT token for API authentication.""" token = args["token"] token_type = args.get("token_type", "Bearer") header_name = args.get("header_name", "Authorization") provider = args.get("provider", "custom") try: # Format token based on type if token_type.lower() == "bearer": auth_header = f"Bearer {token}" elif token_type.lower() == "jwt": auth_header = f"JWT {token}" else: auth_header = f"{token_type} {token}" # Set token in session headers session.headers[header_name] = auth_header result = { "token_type": token_type, "header_name": header_name, "provider": provider, "token_length": len(token), "status": f"OAuth token set successfully", "headers_updated": dict(session.headers) } return [TextContent(type="text", text=json.dumps(result, indent=2, ensure_ascii=False))] except Exception as e: return [TextContent(type="text", text=f"Failed to set OAuth token: {str(e)}")]