agent_os_kernel 3.1.0__py3-none-any.whl

atr/tools/safe/http_client.py

@@ -0,0 +1,316 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Safe HTTP Client Tool.
+
+Provides HTTP request capabilities with security controls:
+- URL whitelisting
+- Rate limiting
+- Request timeout enforcement
+- Response size limits
+- No redirect following to external domains
+"""
+
+import asyncio
+import time
+from collections import deque
+from typing import Any, Dict, List, Optional, Set
+from urllib.parse import urlparse
+
+from atr.decorator import tool
+
+
+class RateLimiter:
+    """Simple rate limiter using sliding window."""
+    
+    def __init__(self, max_requests: int, window_seconds: float = 60.0):
+        self.max_requests = max_requests
+        self.window_seconds = window_seconds
+        self._requests: deque = deque()
+    
+    def check(self) -> bool:
+        """Check if request is allowed."""
+        now = time.monotonic()
+        
+        # Remove old requests
+        while self._requests and self._requests[0] < now - self.window_seconds:
+            self._requests.popleft()
+        
+        return len(self._requests) < self.max_requests
+    
+    def record(self):
+        """Record a request."""
+        self._requests.append(time.monotonic())
+
+
+class HttpClientTool:
+    """
+    Safe HTTP client with security controls.
+    
+    Features:
+    - Domain whitelisting (only specified domains allowed)
+    - Rate limiting (prevent abuse)
+    - Timeout enforcement (prevent hanging)
+    - Response size limits (prevent memory exhaustion)
+    - Safe redirect handling (no external redirects)
+    
+    Example:
+        ```python
+        http = HttpClientTool(
+            allowed_domains=["api.github.com", "httpbin.org"],
+            rate_limit=30,  # 30 requests per minute
+            timeout=10.0,
+            max_response_size=1_000_000  # 1MB
+        )
+        
+        # Register with ATR
+        registry.register(http.get)
+        registry.register(http.post)
+        
+        # Use from agent
+        response = await http.get("https://api.github.com/users/octocat")
+        ```
+    """
+    
+    def __init__(
+        self,
+        allowed_domains: Optional[List[str]] = None,
+        blocked_domains: Optional[List[str]] = None,
+        rate_limit: int = 60,
+        timeout: float = 30.0,
+        max_response_size: int = 10_000_000,  # 10MB
+        allow_redirects: bool = True,
+        max_redirects: int = 5
+    ):
+        """
+        Initialize HTTP client tool.
+        
+        Args:
+            allowed_domains: Whitelist of allowed domains (if set, only these allowed)
+            blocked_domains: Blacklist of blocked domains
+            rate_limit: Maximum requests per minute
+            timeout: Request timeout in seconds
+            max_response_size: Maximum response size in bytes
+            allow_redirects: Whether to follow redirects
+            max_redirects: Maximum number of redirects to follow
+        """
+        self.allowed_domains: Set[str] = set(allowed_domains or [])
+        self.blocked_domains: Set[str] = set(blocked_domains or [
+            "localhost",
+            "127.0.0.1",
+            "0.0.0.0",
+            "169.254.169.254",  # AWS metadata
+            "metadata.google.internal",  # GCP metadata
+        ])
+        self.timeout = timeout
+        self.max_response_size = max_response_size
+        self.allow_redirects = allow_redirects
+        self.max_redirects = max_redirects
+        self._rate_limiter = RateLimiter(rate_limit)
+    
+    def _validate_url(self, url: str) -> str:
+        """Validate and normalize URL."""
+        parsed = urlparse(url)
+        
+        # Must have scheme
+        if not parsed.scheme:
+            raise ValueError("URL must include scheme (http:// or https://)")
+        
+        # Only allow http/https
+        if parsed.scheme not in ("http", "https"):
+            raise ValueError(f"Scheme '{parsed.scheme}' not allowed. Use http or https.")
+        
+        # Extract domain
+        domain = parsed.netloc.lower()
+        if ":" in domain:
+            domain = domain.split(":")[0]  # Remove port
+        
+        # Check blocked domains
+        if domain in self.blocked_domains:
+            raise ValueError(f"Domain '{domain}' is blocked")
+        
+        # Check for private IP ranges (basic check)
+        if self._is_private_domain(domain):
+            raise ValueError(f"Private/internal domains not allowed: {domain}")
+        
+        # Check allowed domains (if whitelist set)
+        if self.allowed_domains:
+            if not any(domain.endswith(allowed) for allowed in self.allowed_domains):
+                raise ValueError(
+                    f"Domain '{domain}' not in allowed list. "
+                    f"Allowed: {', '.join(self.allowed_domains)}"
+                )
+        
+        return url
+    
+    def _is_private_domain(self, domain: str) -> bool:
+        """Check if domain resolves to private IP (basic check)."""
+        # Check common private patterns
+        private_patterns = [
+            "10.",
+            "192.168.",
+            "172.16.", "172.17.", "172.18.", "172.19.",
+            "172.20.", "172.21.", "172.22.", "172.23.",
+            "172.24.", "172.25.", "172.26.", "172.27.",
+            "172.28.", "172.29.", "172.30.", "172.31.",
+            "internal",
+            ".local",
+        ]
+        return any(domain.startswith(p) or domain.endswith(p) for p in private_patterns)
+    
+    def _check_rate_limit(self):
+        """Check and enforce rate limit."""
+        if not self._rate_limiter.check():
+            raise RuntimeError("Rate limit exceeded. Please wait before making more requests.")
+        self._rate_limiter.record()
+    
+    @tool(
+        name="http_get",
+        description="Make a safe HTTP GET request to retrieve data from a URL",
+        tags=["http", "network", "safe"]
+    )
+    async def get(
+        self,
+        url: str,
+        headers: Optional[Dict[str, str]] = None
+    ) -> Dict[str, Any]:
+        """
+        Make a GET request.
+        
+        Args:
+            url: URL to request (must be in allowed domains)
+            headers: Optional request headers
+        
+        Returns:
+            Dict with status_code, headers, and body
+        """
+        try:
+            import aiohttp
+        except ImportError:
+            raise ImportError("aiohttp required. Install with: pip install aiohttp")
+        
+        # Validate
+        url = self._validate_url(url)
+        self._check_rate_limit()
+        
+        # Make request
+        timeout = aiohttp.ClientTimeout(total=self.timeout)
+        
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            async with session.get(
+                url,
+                headers=headers,
+                allow_redirects=self.allow_redirects,
+                max_redirects=self.max_redirects
+            ) as response:
+                # Check response size
+                content_length = response.headers.get("Content-Length")
+                if content_length and int(content_length) > self.max_response_size:
+                    raise ValueError(f"Response too large: {content_length} bytes")
+                
+                # Read response (with size limit)
+                body = await response.text()
+                if len(body) > self.max_response_size:
+                    body = body[:self.max_response_size] + "... [truncated]"
+                
+                return {
+                    "status_code": response.status,
+                    "headers": dict(response.headers),
+                    "body": body,
+                    "url": str(response.url)
+                }
+    
+    @tool(
+        name="http_post",
+        description="Make a safe HTTP POST request to send data to a URL",
+        tags=["http", "network", "safe"]
+    )
+    async def post(
+        self,
+        url: str,
+        data: Optional[Dict[str, Any]] = None,
+        json_body: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, str]] = None
+    ) -> Dict[str, Any]:
+        """
+        Make a POST request.
+        
+        Args:
+            url: URL to request (must be in allowed domains)
+            data: Form data to send
+            json_body: JSON body to send
+            headers: Optional request headers
+        
+        Returns:
+            Dict with status_code, headers, and body
+        """
+        try:
+            import aiohttp
+        except ImportError:
+            raise ImportError("aiohttp required. Install with: pip install aiohttp")
+        
+        # Validate
+        url = self._validate_url(url)
+        self._check_rate_limit()
+        
+        # Make request
+        timeout = aiohttp.ClientTimeout(total=self.timeout)
+        
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            async with session.post(
+                url,
+                data=data,
+                json=json_body,
+                headers=headers,
+                allow_redirects=self.allow_redirects,
+                max_redirects=self.max_redirects
+            ) as response:
+                # Check response size
+                content_length = response.headers.get("Content-Length")
+                if content_length and int(content_length) > self.max_response_size:
+                    raise ValueError(f"Response too large: {content_length} bytes")
+                
+                # Read response
+                body = await response.text()
+                if len(body) > self.max_response_size:
+                    body = body[:self.max_response_size] + "... [truncated]"
+                
+                return {
+                    "status_code": response.status,
+                    "headers": dict(response.headers),
+                    "body": body,
+                    "url": str(response.url)
+                }
+    
+    @tool(
+        name="http_head",
+        description="Make a safe HTTP HEAD request to check if a URL exists",
+        tags=["http", "network", "safe"]
+    )
+    async def head(self, url: str) -> Dict[str, Any]:
+        """
+        Make a HEAD request (useful for checking if resource exists).
+        
+        Args:
+            url: URL to check
+        
+        Returns:
+            Dict with status_code and headers
+        """
+        try:
+            import aiohttp
+        except ImportError:
+            raise ImportError("aiohttp required. Install with: pip install aiohttp")
+        
+        url = self._validate_url(url)
+        self._check_rate_limit()
+        
+        timeout = aiohttp.ClientTimeout(total=self.timeout)
+        
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            async with session.head(url, allow_redirects=self.allow_redirects) as response:
+                return {
+                    "status_code": response.status,
+                    "headers": dict(response.headers),
+                    "url": str(response.url)
+                }

atr/tools/safe/json_parser.py

@@ -0,0 +1,374 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Safe JSON/YAML Parser Tool.
+
+Provides safe parsing of JSON and YAML with:
+- Size limits to prevent memory exhaustion
+- Depth limits to prevent stack overflow
+- Safe YAML loading (no arbitrary code execution)
+- Schema validation support
+"""
+
+import json
+from typing import Any, Dict, List, Optional, Union
+
+from atr.decorator import tool
+
+
+class JsonParserTool:
+    """
+    Safe JSON and YAML parser.
+    
+    Features:
+    - Input size limits
+    - Nesting depth limits
+    - Safe YAML loading (yaml.safe_load)
+    - Schema validation (optional)
+    - Pretty printing
+    
+    Example:
+        ```python
+        parser = JsonParserTool(
+            max_size=1_000_000,  # 1MB
+            max_depth=50
+        )
+        
+        # Parse JSON
+        data = parser.parse_json('{"key": "value"}')
+        
+        # Parse YAML  
+        data = parser.parse_yaml("key: value")
+        
+        # Validate schema
+        parser.validate(data, schema={"type": "object"})
+        ```
+    """
+    
+    def __init__(
+        self,
+        max_size: int = 10_000_000,  # 10MB
+        max_depth: int = 100,
+        max_keys: int = 10000
+    ):
+        """
+        Initialize parser tool.
+        
+        Args:
+            max_size: Maximum input size in characters
+            max_depth: Maximum nesting depth
+            max_keys: Maximum number of keys in objects
+        """
+        self.max_size = max_size
+        self.max_depth = max_depth
+        self.max_keys = max_keys
+    
+    def _check_size(self, data: str, name: str = "input"):
+        """Check input size."""
+        if len(data) > self.max_size:
+            raise ValueError(
+                f"{name} too large: {len(data)} chars. Maximum: {self.max_size}"
+            )
+    
+    def _check_depth(self, obj: Any, current_depth: int = 0):
+        """Check nesting depth recursively."""
+        if current_depth > self.max_depth:
+            raise ValueError(f"Nesting too deep. Maximum depth: {self.max_depth}")
+        
+        if isinstance(obj, dict):
+            if len(obj) > self.max_keys:
+                raise ValueError(f"Too many keys: {len(obj)}. Maximum: {self.max_keys}")
+            for value in obj.values():
+                self._check_depth(value, current_depth + 1)
+        elif isinstance(obj, list):
+            for item in obj:
+                self._check_depth(item, current_depth + 1)
+    
+    @tool(
+        name="parse_json",
+        description="Safely parse a JSON string into a Python object",
+        tags=["json", "parse", "safe"]
+    )
+    def parse_json(self, data: str) -> Dict[str, Any]:
+        """
+        Parse JSON string.
+        
+        Args:
+            data: JSON string to parse
+        
+        Returns:
+            Dict with parsed data and metadata
+        """
+        self._check_size(data, "JSON")
+        
+        try:
+            parsed = json.loads(data)
+        except json.JSONDecodeError as e:
+            return {
+                "success": False,
+                "error": f"Invalid JSON: {e.msg} at line {e.lineno}, column {e.colno}",
+                "data": None
+            }
+        
+        # Check depth
+        try:
+            self._check_depth(parsed)
+        except ValueError as e:
+            return {
+                "success": False,
+                "error": str(e),
+                "data": None
+            }
+        
+        return {
+            "success": True,
+            "data": parsed,
+            "type": type(parsed).__name__
+        }
+    
+    @tool(
+        name="parse_yaml",
+        description="Safely parse a YAML string into a Python object",
+        tags=["yaml", "parse", "safe"]
+    )
+    def parse_yaml(self, data: str) -> Dict[str, Any]:
+        """
+        Parse YAML string safely.
+        
+        Uses yaml.safe_load to prevent arbitrary code execution.
+        
+        Args:
+            data: YAML string to parse
+        
+        Returns:
+            Dict with parsed data and metadata
+        """
+        try:
+            import yaml
+        except ImportError:
+            return {
+                "success": False,
+                "error": "PyYAML not installed. Install with: pip install pyyaml",
+                "data": None
+            }
+        
+        self._check_size(data, "YAML")
+        
+        try:
+            # Always use safe_load to prevent code execution
+            parsed = yaml.safe_load(data)
+        except yaml.YAMLError as e:
+            return {
+                "success": False,
+                "error": f"Invalid YAML: {e}",
+                "data": None
+            }
+        
+        # Check depth
+        try:
+            self._check_depth(parsed)
+        except ValueError as e:
+            return {
+                "success": False,
+                "error": str(e),
+                "data": None
+            }
+        
+        return {
+            "success": True,
+            "data": parsed,
+            "type": type(parsed).__name__ if parsed is not None else "null"
+        }
+    
+    @tool(
+        name="to_json",
+        description="Convert a Python object to a JSON string",
+        tags=["json", "serialize", "safe"]
+    )
+    def to_json(
+        self,
+        data: Any,
+        indent: Optional[int] = 2,
+        sort_keys: bool = False
+    ) -> Dict[str, Any]:
+        """
+        Convert object to JSON string.
+        
+        Args:
+            data: Object to serialize
+            indent: Indentation level (None for compact)
+            sort_keys: Whether to sort dictionary keys
+        
+        Returns:
+            Dict with JSON string
+        """
+        try:
+            result = json.dumps(data, indent=indent, sort_keys=sort_keys, ensure_ascii=False)
+            
+            if len(result) > self.max_size:
+                return {
+                    "success": False,
+                    "error": f"Output too large: {len(result)} chars",
+                    "json": None
+                }
+            
+            return {
+                "success": True,
+                "json": result,
+                "size": len(result)
+            }
+        except (TypeError, ValueError) as e:
+            return {
+                "success": False,
+                "error": f"Cannot serialize: {e}",
+                "json": None
+            }
+    
+    @tool(
+        name="to_yaml",
+        description="Convert a Python object to a YAML string",
+        tags=["yaml", "serialize", "safe"]
+    )
+    def to_yaml(
+        self,
+        data: Any,
+        default_flow_style: bool = False
+    ) -> Dict[str, Any]:
+        """
+        Convert object to YAML string.
+        
+        Args:
+            data: Object to serialize
+            default_flow_style: Use flow style (compact) formatting
+        
+        Returns:
+            Dict with YAML string
+        """
+        try:
+            import yaml
+        except ImportError:
+            return {
+                "success": False,
+                "error": "PyYAML not installed. Install with: pip install pyyaml",
+                "yaml": None
+            }
+        
+        try:
+            result = yaml.safe_dump(
+                data,
+                default_flow_style=default_flow_style,
+                allow_unicode=True
+            )
+            
+            if len(result) > self.max_size:
+                return {
+                    "success": False,
+                    "error": f"Output too large: {len(result)} chars",
+                    "yaml": None
+                }
+            
+            return {
+                "success": True,
+                "yaml": result,
+                "size": len(result)
+            }
+        except yaml.YAMLError as e:
+            return {
+                "success": False,
+                "error": f"Cannot serialize: {e}",
+                "yaml": None
+            }
+    
+    @tool(
+        name="validate_json_schema",
+        description="Validate data against a JSON schema",
+        tags=["json", "validate", "schema", "safe"]
+    )
+    def validate_schema(
+        self,
+        data: Any,
+        schema: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """
+        Validate data against JSON schema.
+        
+        Args:
+            data: Data to validate
+            schema: JSON schema
+        
+        Returns:
+            Dict with validation result
+        """
+        try:
+            import jsonschema
+        except ImportError:
+            return {
+                "valid": False,
+                "error": "jsonschema not installed. Install with: pip install jsonschema"
+            }
+        
+        try:
+            jsonschema.validate(instance=data, schema=schema)
+            return {
+                "valid": True,
+                "error": None
+            }
+        except jsonschema.ValidationError as e:
+            return {
+                "valid": False,
+                "error": e.message,
+                "path": list(e.absolute_path)
+            }
+        except jsonschema.SchemaError as e:
+            return {
+                "valid": False,
+                "error": f"Invalid schema: {e.message}"
+            }
+    
+    @tool(
+        name="json_query",
+        description="Query JSON data using JSONPath or simple dot notation",
+        tags=["json", "query", "safe"]
+    )
+    def query(
+        self,
+        data: Any,
+        path: str
+    ) -> Dict[str, Any]:
+        """
+        Query JSON data using dot notation.
+        
+        Args:
+            data: JSON data to query
+            path: Dot-notation path (e.g., "users.0.name")
+        
+        Returns:
+            Dict with query result
+        """
+        parts = path.split(".")
+        current = data
+        
+        try:
+            for part in parts:
+                if not part:
+                    continue
+                
+                if isinstance(current, dict):
+                    current = current[part]
+                elif isinstance(current, list):
+                    idx = int(part)
+                    current = current[idx]
+                else:
+                    raise KeyError(f"Cannot access '{part}' on {type(current).__name__}")
+            
+            return {
+                "success": True,
+                "value": current,
+                "path": path
+            }
+        except (KeyError, IndexError, ValueError) as e:
+            return {
+                "success": False,
+                "error": f"Path not found: {e}",
+                "value": None
+            }