utils.py

import os
output_dir = "/mnt/agents/output"

# 2. utils.py - Same as before, works for both
utils_content = r'''import time
import hashlib
import re
import subprocess
import threading
from functools import wraps
from typing import Any, Callable, Optional

# Rate limit storage: {ip: [(timestamp, count)]}
_rate_limit_store = {}
_rate_limit_lock = threading.Lock()

def rate_limit(key: str, max_requests: int = 60, window: int = 60) -> bool:
    """
    Check if the key (IP or user ID) has exceeded rate limits.
    Returns True if allowed, False if rate limited.
    """
    now = time.time()
    with _rate_limit_lock:
        if key not in _rate_limit_store:
            _rate_limit_store[key] = []
        # Clean old entries outside the window
        _rate_limit_store[key] = [t for t in _rate_limit_store[key] if now - t < window]
        if len(_rate_limit_store[key]) >= max_requests:
            return False
        _rate_limit_store[key].append(now)
        return True

def safe_exec(
    code: str,
    timeout: int = 30,
    allowed_imports: Optional[list] = None,
    blocked_builtins: Optional[list] = None
) -> dict:
    """
    Safely execute Python code in a restricted environment.
    Returns dict with stdout, stderr, result, and success flag.
    """
    import sys
    import io
    import signal
    
    allowed_imports = allowed_imports or [
        'math', 'random', 'datetime', 'json', 're', 'string', 'collections',
        'itertools', 'functools', 'statistics', 'hashlib', 'base64', 'typing',
        'fractions', 'decimal', 'inspect', 'textwrap', 'uuid', 'time', 'calendar',
        'bisect', 'heapq', 'copy', 'pprint', 'numbers', 'enum', 'dataclasses',
        'pathlib', 'urllib.parse', 'html'
    ]
    blocked_builtins = blocked_builtins or ['open', 'exec', 'eval', 'compile', '__import__']
    
    # Create restricted globals
    restricted_globals = {
        '__builtins__': {k: v for k, v in __builtins__.items() if k not in blocked_builtins}
    }
    
    # Add allowed imports
    for mod_name in allowed_imports:
        try:
            restricted_globals[mod_name] = __import__(mod_name)
        except ImportError:
            pass
    
    old_stdout = sys.stdout
    old_stderr = sys.stderr
    stdout_buffer = io.StringIO()
    stderr_buffer = io.StringIO()
    
    result = {"stdout": "", "stderr": "", "result": None, "success": False}
    
    def timeout_handler(signum, frame):
        raise TimeoutError("Code execution timed out")
    
    try:
        sys.stdout = stdout_buffer
        sys.stderr = stderr_buffer
        
        # Set alarm for timeout (Unix only; Windows uses different approach)
        if hasattr(signal, 'SIGALRM'):
            signal.signal(signal.SIGALRM, timeout_handler)
            signal.alarm(timeout)
        
        exec(code, restricted_globals)
        result["success"] = True
        
        if hasattr(signal, 'SIGALRM'):
            signal.alarm(0)
            
    except TimeoutError as e:
        result["stderr"] = f"Timeout Error: {str(e)}"
    except Exception as e:
        result["stderr"] = f"{type(e).__name__}: {str(e)}"
    finally:
        sys.stdout = old_stdout
        sys.stderr = old_stderr
        if hasattr(signal, 'SIGALRM'):
            signal.alarm(0)
    
    result["stdout"] = stdout_buffer.getvalue()
    result["stderr"] += stderr_buffer.getvalue()
    
    # Try to capture the last expression result
    try:
        # Re-run just the last expression to get its value
        lines = [l for l in code.split('\n') if l.strip() and not l.strip().startswith('#')]
        if lines:
            last_line = lines[-1].strip()
            if not last_line.startswith(('import ', 'from ', 'def ', 'class ', 'if ', 'for ', 'while ', 'try:', 'with ', 'print(')):
                local_vars = {}
                exec(f"_last_result = {last_line}", restricted_globals, local_vars)
                result["result"] = local_vars.get("_last_result")
    except:
        pass
    
    return result

def safe_command_exec(command: str, timeout: int = 30) -> dict:
    """
    Safely execute a shell command with timeout and output capture.
    """
    result = {"stdout": "", "stderr": "", "returncode": -1, "success": False}
    try:
        proc = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=timeout
        )
        result["stdout"] = proc.stdout
        result["stderr"] = proc.stderr
        result["returncode"] = proc.returncode
        result["success"] = proc.returncode == 0
    except subprocess.TimeoutExpired:
        result["stderr"] = f"Command timed out after {timeout} seconds"
    except Exception as e:
        result["stderr"] = f"{type(e).__name__}: {str(e)}"
    return result

def generate_password(length: int = 16, include_special: bool = True) -> str:
    """Generate a secure random password."""
    import random
    import string
    chars = string.ascii_letters + string.digits
    if include_special:
        chars += string.punctuation
    return ''.join(random.choice(chars) for _ in range(length))

def hash_string(data: str, algorithm: str = "sha256") -> str:
    """Hash a string using various algorithms."""
    if algorithm.lower() == "md5":
        return hashlib.md5(data.encode()).hexdigest()
    elif algorithm.lower() == "sha1":
        return hashlib.sha1(data.encode()).hexdigest()
    elif algorithm.lower() == "sha256":
        return hashlib.sha256(data.encode()).hexdigest()
    elif algorithm.lower() == "sha512":
        return hashlib.sha512(data.encode()).hexdigest()
    elif algorithm.lower() == "blake2b":
        return hashlib.blake2b(data.encode()).hexdigest()
    else:
        return hashlib.sha256(data.encode()).hexdigest()

def retry_with_backoff(
    func: Callable,
    max_retries: int = 3,
    initial_delay: float = 1.0,
    backoff_factor: float = 2.0,
    exceptions: tuple = (Exception,)
) -> Any:
    """Retry a function with exponential backoff."""
    delay = initial_delay
    for attempt in range(max_retries):
        try:
            return func()
        except exceptions as e:
            if attempt == max_retries - 1:
                raise
            time.sleep(delay)
            delay *= backoff_factor
    return None
'''

with open(f"{output_dir}/utils.py", "w") as f:
    f.write(utils_content)

print("utils.py written")