""" Data extraction tools: text, elements, links """ import json import re from typing import Dict, Any, List from urllib.parse import urljoin, urlparse from bs4 import BeautifulSoup from mcp.types import TextContent from ..config import session, authenticated_session async def extract_text_tool(args: Dict[str, Any]) -> List[TextContent]: """Extract clean text content from HTML.""" html = args["html"] preserve_formatting = args.get("preserve_formatting", False) try: soup = BeautifulSoup(html, 'lxml') # Remove script and style elements for script in soup(["script", "style"]): script.decompose() if preserve_formatting: text = soup.get_text(separator='\n', strip=True) else: text = soup.get_text(strip=True) # Clean up extra whitespace text = ' '.join(text.split()) result = { "text": text, "length": len(text), "preserve_formatting": preserve_formatting } 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 extract text: {str(e)}")] async def find_elements_tool(args: Dict[str, Any]) -> List[TextContent]: """Find HTML elements using CSS selectors.""" html = args["html"] selector = args["selector"] attribute = args.get("attribute") limit = args.get("limit", 10) try: soup = BeautifulSoup(html, 'lxml') elements = soup.select(selector)[:limit] results = [] for element in elements: if attribute: value = element.get(attribute) if value: results.append(value) else: results.append({ "tag": element.name, "text": element.get_text(strip=True), "attributes": dict(element.attrs) }) result = { "selector": selector, "found_count": len(elements), "returned_count": len(results), "results": results } 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 find elements: {str(e)}")] async def extract_links_tool(args: Dict[str, Any]) -> List[TextContent]: """Extract and filter links from HTML or URL with advanced options.""" url = args.get("url") html = args.get("html") base_url = args.get("base_url") internal_only = args.get("internal_only", False) external_only = args.get("external_only", False) url_pattern = args.get("url_pattern") text_pattern = args.get("text_pattern") exclude_fragments = args.get("exclude_fragments", True) unique_only = args.get("unique_only", True) limit = args.get("limit", 100) try: # Session info will be set later if authenticated # Get HTML content using efficient API with browser session cookies if url: 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)) base_url = base_url or url # Add session info if authenticated if authenticated_session.get("active"): cookies_info = api_client.get_cookies_info() if cookies_info["cookies_count"] > 0: session_info = { "source": "api_with_browser_session", "login_url": authenticated_session.get("login_url"), "cookies_count": cookies_info["cookies_count"], "session_age_hours": ( (time.time() - authenticated_session.get("login_timestamp", time.time())) / 3600 if authenticated_session.get("login_timestamp") else 0 ), "api_response_time": result_data.get("response_time_seconds", 0) } elif not html: return [TextContent(type="text", text="Either 'url' or 'html' parameter is required")] soup = BeautifulSoup(html, 'lxml') # Get base domain for internal/external filtering base_domain = None if base_url and (internal_only or external_only): base_domain = urlparse(base_url).netloc.lower() links = [] seen_urls = set() for link in soup.find_all('a', href=True): href = link['href'].strip() text = link.get_text(strip=True) # Skip empty links if not href: continue # Skip fragments if requested if exclude_fragments and href.startswith('#'): continue # Resolve relative URLs if base_url: absolute_url = urljoin(base_url, href) else: absolute_url = href # Filter by URL pattern if url_pattern and not re.search(url_pattern, absolute_url): continue # Filter by text pattern if text_pattern and not re.search(text_pattern, text): continue # Filter internal/external links if base_domain: link_domain = urlparse(absolute_url).netloc.lower() if internal_only and link_domain != base_domain: continue if external_only and link_domain == base_domain: continue # Check for uniqueness if unique_only: if absolute_url in seen_urls: continue seen_urls.add(absolute_url) link_data = { "url": absolute_url, "text": text, "original_href": href } # Add domain info for easier filtering parsed_url = urlparse(absolute_url) link_data["domain"] = parsed_url.netloc link_data["scheme"] = parsed_url.scheme link_data["path"] = parsed_url.path links.append(link_data) # Respect limit if len(links) >= limit: break # Compile results result = { "total_found": len(links), "base_url": base_url, "filters_applied": { "internal_only": internal_only, "external_only": external_only, "url_pattern": url_pattern, "text_pattern": text_pattern, "exclude_fragments": exclude_fragments, "unique_only": unique_only }, "links": links } # Add session information if used if session_info: result["session_used"] = session_info 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 extract links: {str(e)}")]