janito 2.22.0__py3-none-any.whl → 2.24.0__py3-none-any.whl

janito/tools/loop_protection_decorator.py

@@ -0,0 +1,110 @@
+import functools
+import time
+import threading
+from typing import Callable, Any
+from janito.tools.loop_protection import LoopProtection
+from janito.tools.tool_use_tracker import normalize_path
+
+
+# Global tracking for decorator-based loop protection
+_decorator_call_tracker = {}
+_decorator_call_tracker_lock = threading.Lock()
+
+
+def protect_against_loops(
+    max_calls: int = 5,
+    time_window: float = 10.0
+):
+    """
+    Decorator that adds loop protection to tool run methods.
+    
+    This decorator monitors tool executions and prevents excessive calls within
+    a configurable time window. It helps prevent infinite loops or excessive
+    resource consumption when tools are called repeatedly.
+    
+    When the configured limits are exceeded, the decorator raises a RuntimeError
+    with a descriptive message. This exception will propagate up the call stack
+    unless caught by a try/except block in the calling code.
+    
+    The decorator works by:
+    1. Tracking the number of calls to the decorated function
+    2. Checking if the calls exceed the configured limits
+    3. Raising a RuntimeError if a potential loop is detected
+    4. Allowing the method to proceed normally if the operation is safe
+    
+    Args:
+        max_calls (int): Maximum number of calls allowed within the time window.
+                        Defaults to 5 calls.
+        time_window (float): Time window in seconds for detecting excessive calls.
+                            Defaults to 10.0 seconds.
+                              
+    Example:
+        >>> @protect_against_loops(max_calls=3, time_window=5.0)
+        >>> def run(self, path: str) -> str:
+        >>>     # Implementation here
+        >>>     pass
+          
+        >>> @protect_against_loops(max_calls=10, time_window=30.0)
+        >>> def run(self, file_paths: list) -> str:
+        >>>     # Implementation here
+        >>>     pass
+        
+    Note:
+        When loop protection is triggered, a RuntimeError will be raised with a
+        descriptive message. This exception will propagate up the call stack
+        unless caught by a try/except block in the calling code.
+    """
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            # Get the tool instance (self)
+            if not args:
+                # This shouldn't happen in normal usage as methods need self
+                return func(*args, **kwargs)
+            
+            # Use the function name as the operation name
+            op_name = func.__name__
+            
+            # Check call limits
+            current_time = time.time()
+            
+            with _decorator_call_tracker_lock:
+                # Clean up old entries outside the time window
+                if op_name in _decorator_call_tracker:
+                    _decorator_call_tracker[op_name] = [
+                        timestamp for timestamp in _decorator_call_tracker[op_name]
+                        if current_time - timestamp <= time_window
+                    ]
+                
+                # Check if we're exceeding the limit
+                if op_name in _decorator_call_tracker:
+                    if len(_decorator_call_tracker[op_name]) >= max_calls:
+                        # Check if all recent calls are within the time window
+                        if all(current_time - timestamp <= time_window 
+                              for timestamp in _decorator_call_tracker[op_name]):
+                            # Define the error reporting function
+                            def _report_error_and_raise(args, operation_type):
+                                # Get the tool instance to access report_error method if available
+                                tool_instance = args[0] if args else None
+                                error_msg = f"Loop protection: Too many {operation_type} operations in a short time period ({max_calls} calls in {time_window}s)"
+                                
+                                # Try to report the error through the tool's reporting mechanism
+                                if hasattr(tool_instance, 'report_error'):
+                                    try:
+                                        tool_instance.report_error(error_msg)
+                                    except Exception:
+                                        pass  # If reporting fails, we still raise the error
+                                
+                                raise RuntimeError(error_msg)
+                            
+                            _report_error_and_raise(args, op_name)
+                
+                # Record this call
+                if op_name not in _decorator_call_tracker:
+                    _decorator_call_tracker[op_name] = []
+                _decorator_call_tracker[op_name].append(current_time)
+            
+            # Proceed with the original function
+            return func(*args, **kwargs)
+        return wrapper
+    return decorator

janito/tools/url_whitelist.py

@@ -0,0 +1,121 @@
+"""URL whitelist management for fetch_url tool."""
+
+import json
+from pathlib import Path
+from typing import Set, List, Optional
+from urllib.parse import urlparse
+
+
+class UrlWhitelistManager:
+    """Manages allowed sites for the fetch_url tool."""
+    
+    def __init__(self):
+        self.config_path = Path.home() / ".janito" / "url_whitelist.json"
+        self._allowed_sites = self._load_whitelist()
+        self._unrestricted_mode = False
+    
+    def set_unrestricted_mode(self, enabled: bool = True):
+        """Enable or disable unrestricted mode (bypasses whitelist)."""
+        self._unrestricted_mode = enabled
+    
+    def _load_whitelist(self) -> Set[str]:
+        """Load the whitelist from config file."""
+        if not self.config_path.exists():
+            return set()
+        
+        try:
+            with open(self.config_path, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+                return set(data.get('allowed_sites', []))
+        except (json.JSONDecodeError, IOError):
+            return set()
+    
+    def _save_whitelist(self):
+        """Save the whitelist to config file."""
+        self.config_path.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            with open(self.config_path, 'w', encoding='utf-8') as f:
+                json.dump({'allowed_sites': list(self._allowed_sites)}, f, indent=2)
+        except IOError:
+            pass  # Silently fail if we can't write
+    
+    def is_url_allowed(self, url: str) -> bool:
+        """Check if a URL is allowed based on the whitelist."""
+        if self._unrestricted_mode:
+            return True  # Unrestricted mode bypasses all whitelist checks
+        
+        if not self._allowed_sites:
+            return True  # No whitelist means all sites allowed
+        
+        try:
+            parsed = urlparse(url)
+            domain = parsed.netloc.lower()
+            
+            # Check exact matches and subdomain matches
+            for allowed in self._allowed_sites:
+                allowed = allowed.lower()
+                if domain == allowed or domain.endswith('.' + allowed):
+                    return True
+            
+            return False
+        except Exception:
+            return False  # Invalid URLs are blocked
+    
+    def add_allowed_site(self, site: str) -> bool:
+        """Add a site to the whitelist."""
+        # Clean up the site format
+        site = site.strip().lower()
+        if site.startswith('http://') or site.startswith('https://'):
+            parsed = urlparse(site)
+            site = parsed.netloc
+        
+        if site and site not in self._allowed_sites:
+            self._allowed_sites.add(site)
+            self._save_whitelist()
+            return True
+        return False
+    
+    def remove_allowed_site(self, site: str) -> bool:
+        """Remove a site from the whitelist."""
+        site = site.strip().lower()
+        if site.startswith('http://') or site.startswith('https://'):
+            parsed = urlparse(site)
+            site = parsed.netloc
+        
+        if site in self._allowed_sites:
+            self._allowed_sites.remove(site)
+            self._save_whitelist()
+            return True
+        return False
+    
+    def get_allowed_sites(self) -> List[str]:
+        """Get the list of allowed sites."""
+        return sorted(self._allowed_sites)
+    
+    def set_allowed_sites(self, sites: List[str]):
+        """Set the complete list of allowed sites."""
+        self._allowed_sites = set()
+        for site in sites:
+            site = site.strip().lower()
+            if site.startswith('http://') or site.startswith('https://'):
+                parsed = urlparse(site)
+                site = parsed.netloc
+            if site:
+                self._allowed_sites.add(site)
+        self._save_whitelist()
+    
+    def clear_whitelist(self):
+        """Clear all allowed sites."""
+        self._allowed_sites.clear()
+        self._save_whitelist()
+
+
+# Global singleton
+_url_whitelist_manager = None
+
+def get_url_whitelist_manager() -> UrlWhitelistManager:
+    """Get the global URL whitelist manager instance."""
+    global _url_whitelist_manager
+    if _url_whitelist_manager is None:
+        _url_whitelist_manager = UrlWhitelistManager()
+    return _url_whitelist_manager