mcp-forge/src/pod_executor/security/resource_limits.py

"""
Resource limit parser and validator.

Parses and validates resource limit strings (memory, CPU, storage).
All values must be positive and within reasonable limits.
"""

import re
from typing import Dict, Any


def parse_memory_string(memory: str) -> int:
    """
    Parse memory string to bytes.
    
    Supports: k, m, g suffixes (case-insensitive)
    Examples: "512m" → 536870912, "2g" → 2147483648
    
    Args:
        memory: Memory string with suffix (e.g., "512m", "2g", "1024k")
        
    Returns:
        Memory in bytes
        
    Raises:
        ValueError: If format is invalid or value is <= 0
    """
    memory = memory.strip()
    
    # Pattern: optional sign, number (int or float), suffix (k/m/g)
    pattern = r'^(-?\d+(?:\.\d+)?)\s*([kmgKMG])$'
    match = re.match(pattern, memory)
    
    if not match:
        raise ValueError(
            f"Invalid memory format: '{memory}'. "
            f"Expected format: <number><k|m|g> (e.g., '512m', '2g')"
        )
    
    value_str, suffix = match.groups()
    value = float(value_str)
    
    if value <= 0:
        raise ValueError(
            f"Memory value must be positive, got: {value}"
        )
    
    # Convert to bytes
    suffix_lower = suffix.lower()
    multipliers = {
        'k': 1024,
        'm': 1024 * 1024,
        'g': 1024 * 1024 * 1024,
    }
    
    bytes_value = int(value * multipliers[suffix_lower])
    
    return bytes_value


def parse_cpu_quota(cpu_quota: int) -> int:
    """
    Validate CPU quota value.
    
    CPU quota is in microseconds per 100ms period.
    100000 = 100% of one CPU core
    
    Args:
        cpu_quota: CPU quota in microseconds (e.g., 50000 for 50% of one core)
        
    Returns:
        Validated CPU quota value
        
    Raises:
        ValueError: If quota <= 0 or > 1000000 (10 cores max)
    """
    if cpu_quota <= 0:
        raise ValueError(
            f"CPU quota must be positive, got: {cpu_quota}"
        )
    
    # Maximum of 10 cores (1000000 microseconds)
    if cpu_quota > 1000000:
        raise ValueError(
            f"CPU quota exceeds maximum of 1000000 (10 cores), got: {cpu_quota}"
        )
    
    return cpu_quota


def parse_storage_string(storage: str) -> int:
    """
    Parse storage string to bytes (same as memory).
    
    Args:
        storage: Storage string with suffix (e.g., "1g", "512m")
        
    Returns:
        Storage in bytes
        
    Raises:
        ValueError: If format is invalid or value is <= 0
    """
    return parse_memory_string(storage)


class ResourceLimits:
    """
    Resource limits with validation.
    
    Encapsulates memory, storage, CPU, and timeout limits with validation.
    Provides conversion to Podman container parameters.
    """
    
    def __init__(
        self,
        memory: str,
        storage: str,
        cpu_quota: int,
        timeout: int = 300
    ):
        """
        Initialize resource limits with validation.
        
        Args:
            memory: Memory limit string (e.g., "512m", "2g")
            storage: Storage limit string (e.g., "1g", "10g")
            cpu_quota: CPU quota in microseconds per 100ms period
            timeout: Execution timeout in seconds (default: 300)
            
        Raises:
            ValueError: If any limit is invalid
        """
        self.memory_bytes = parse_memory_string(memory)
        self.storage_bytes = parse_storage_string(storage)
        self.cpu_quota = parse_cpu_quota(cpu_quota)
        self.timeout = timeout
        
    def to_podman_params(self) -> Dict[str, Any]:
        """
        Convert to Podman container create parameters.
        
        Returns:
            Dictionary of parameters suitable for Podman container creation
        """
        return {
            "mem_limit": str(self.memory_bytes),  # Podman expects string
            "cpu_quota": self.cpu_quota
            # Note: storage_bytes tracked internally but not passed to Podman (not supported)
        }