decodingtrust-agent-sdk 0.1.0__py3-none-any.whl

utils/task_helpers.py

@@ -0,0 +1,270 @@
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional, Union
+
+import yaml
+
+from .config import PROJECT_ROOT
+
+
+def extract_dataset_path(task_path: Union[str, Path]) -> str:
+    """
+    Extract the relative dataset path from a task directory or config file path.
+
+    Expected input formats:
+        dataset/<domain>/benign/<task_id>/config.yaml
+        dataset/<domain>/malicious/<threat_model>/<category>/<task_id>/
+        /abs/path/to/dataset/<domain>/malicious/<threat_model>/<category>/<task_id>
+
+    Returns path like:
+        crm/benign/1
+        crm/malicious/indirect/data-exfiltration/1
+
+    Args:
+        task_path: Path to task directory or config file
+
+    Returns:
+        Relative path after 'dataset/' directory
+    """
+    # Normalize path
+    if isinstance(task_path, str):
+        task_path = Path(task_path)
+    task_path = task_path.resolve()
+
+    # If it's a file, get the parent directory
+    if task_path.is_file():
+        task_path = task_path.parent
+
+    parts = task_path.parts
+    try:
+        dataset_idx = parts.index('dataset')
+        # Return the path after 'dataset/'
+        return str(Path(*parts[dataset_idx + 1:]))
+    except ValueError:
+        # 'dataset' not found in path, use the last components as fallback
+        return str(Path(*parts[-3:])) if len(parts) >= 3 else task_path.name
+
+
+@dataclass
+class TaskSpec:
+    task_dir: Path
+    domain: Optional[str] = None
+    task_type: Optional[str] = None  # "benign" or "malicious"
+    threat_model: Optional[str] = None  # "direct" or "indirect" (for malicious)
+    risk_category: Optional[str] = None  # e.g., "social-scoring", "data-exfiltration"
+    task_id: Optional[str] = None
+
+
+def build_task_dir(obj: dict, root: Optional[Path] = None) -> Optional[Path]:
+    """
+    Build task directory path from task specification.
+
+    Format: {"domain": "workflow", "type": "benign", "task_id": 100}
+    For malicious: {"domain": "crm", "type": "malicious", "threat_model": "indirect",
+                   "risk_category": "social-scoring", "task_id": 1}
+
+    Directory structure:
+    - Benign: dataset/<domain>/benign/<task_id>/
+    - Malicious: dataset/<domain>/malicious/<threat_model>/<risk_category>/<task_id>/
+
+    Args:
+        obj: Dictionary with task specification
+        root: Root directory for resolving paths (defaults to PROJECT_ROOT)
+
+    Returns:
+        Resolved Path to task directory, or None if unable to determine
+    """
+    if root is None:
+        root = PROJECT_ROOT
+
+    domain = obj.get("domain")
+    task_type = obj.get("type")
+    task_id = obj.get("task_id")
+
+    if not all([domain, task_type, task_id]):
+        return None
+
+    if task_type == "benign":
+        # dataset/<domain>/benign/<task_id>/
+        task_path = root / "dataset" / domain / "benign" / str(task_id)
+    elif task_type == "malicious":
+        threat_model = obj.get("threat_model")
+        risk_category = obj.get("risk_category")
+        if not all([threat_model, risk_category]):
+            return None
+        # dataset/<domain>/malicious/<threat_model>/<risk_category>/<task_id>/
+        task_path = root / "dataset" / domain / "malicious" / threat_model / risk_category / str(task_id)
+    else:
+        return None
+
+    return task_path.resolve()
+
+
+def parse_task_list(
+    path: Path,
+    domain: Optional[str] = None,
+    task_type: Optional[str] = None,
+    threat_model: Optional[str] = None,
+    risk_category: Optional[str] = None,
+) -> List[TaskSpec]:
+    """
+    Parse a JSONL task list file with optional filters.
+
+    Args:
+        path: Path to the JSONL task file
+        domain: Optional domain to filter tasks by (e.g., "workflow", "crm")
+        task_type: Optional type filter ("benign" or "malicious")
+        threat_model: Optional threat model filter ("direct" or "indirect")
+        risk_category: Optional risk category filter (e.g., "social-scoring", "data-exfiltration")
+
+    Returns:
+        List of TaskSpec objects matching all specified filters
+    """
+    specs: List[TaskSpec] = []
+    with path.open("r", encoding="utf-8") as f:
+        for line_num, raw_line in enumerate(f, 1):
+            line = raw_line.strip()
+            if not line:
+                continue
+            try:
+                obj = json.loads(line)
+            except json.JSONDecodeError:
+                print(f"[WARN] Skipping invalid JSON at line {line_num}: {line!r}")
+                continue
+            if not isinstance(obj, dict):
+                continue
+
+            # Apply filters
+            domain_val = obj.get("domain")
+            type_val = obj.get("type")
+            threat_model_val = obj.get("threat_model")
+            risk_category_val = obj.get("risk_category")
+
+            if domain and domain_val != domain:
+                continue
+            if task_type and type_val != task_type:
+                continue
+            if threat_model and threat_model_val != threat_model:
+                continue
+            if risk_category and risk_category_val != risk_category:
+                continue
+
+            # Build task directory path
+            task_path = build_task_dir(obj)
+            if not task_path:
+                print(f"[WARN] Skipping line {line_num}: unable to determine task directory")
+                continue
+
+            specs.append(TaskSpec(
+                task_dir=task_path,
+                domain=domain_val,
+                task_type=type_val,
+                threat_model=threat_model_val,
+                risk_category=risk_category_val,
+                task_id=str(obj.get("task_id")) if obj.get("task_id") else None,
+            ))
+    return specs
+
+
+def parse_red_teaming_task_file(
+    path: Path,
+    domain: Optional[str] = None,
+    threat_model: Optional[str] = None,
+    risk_category: Optional[str] = None,
+) -> List[TaskSpec]:
+    """
+    Parse a JSONL task list file for red-teaming tasks.
+
+    This parses malicious tasks only. Environment injection configuration
+    is now loaded from each task's config.yaml instead of the task file.
+
+    Args:
+        path: Path to the JSONL task file
+        domain: Optional domain to filter tasks by (e.g., "travel", "crm")
+        threat_model: Optional threat model filter ("direct" or "indirect")
+        risk_category: Optional risk category filter
+
+    Returns:
+        List of TaskSpec objects matching all specified filters
+    """
+    specs: List[TaskSpec] = []
+    with path.open("r", encoding="utf-8") as f:
+        for line_num, raw_line in enumerate(f, 1):
+            line = raw_line.strip()
+            if not line:
+                continue
+            try:
+                obj = json.loads(line)
+            except json.JSONDecodeError:
+                print(f"[WARN] Skipping invalid JSON at line {line_num}: {line!r}")
+                continue
+            if not isinstance(obj, dict):
+                continue
+
+            # Red-teaming only targets malicious tasks
+            type_val = obj.get("type")
+            if type_val != "malicious":
+                continue
+
+            # Apply filters
+            domain_val = obj.get("domain")
+            threat_model_val = obj.get("threat_model")
+            risk_category_val = obj.get("risk_category")
+
+            if domain and domain_val != domain:
+                continue
+            if threat_model and threat_model_val != threat_model:
+                continue
+            if risk_category and risk_category_val != risk_category:
+                continue
+
+            # Build task directory path
+            task_path = build_task_dir(obj)
+            if not task_path:
+                print(f"[WARN] Skipping line {line_num}: unable to determine task directory")
+                continue
+
+            specs.append(TaskSpec(
+                task_dir=task_path,
+                domain=domain_val,
+                task_type=type_val,
+                threat_model=threat_model_val,
+                risk_category=risk_category_val,
+                task_id=str(obj.get("task_id")) if obj.get("task_id") else None,
+            ))
+    return specs
+
+
+def has_existing_attack(task_dir: Path) -> bool:
+    """
+    Check if a task's config.yaml contains a non-empty attack_turns field.
+
+    This is used to skip tasks where an attack has already been found/defined.
+
+    Args:
+        task_dir: Path to the task directory containing config.yaml
+
+    Returns:
+        True if the config.yaml exists and has non-empty attack_turns, False otherwise
+    """
+    config_path = task_dir / "config.yaml"
+    if not config_path.exists():
+        return False
+
+    try:
+        with open(config_path, "r") as f:
+            data = yaml.safe_load(f)
+
+        if not data:
+            return False
+
+        attack_section = data.get("Attack", {})
+        attack_turns = attack_section.get("attack_turns", [])
+
+        # Consider it as having an existing attack if attack_turns is non-empty
+        return bool(attack_turns)
+
+    except Exception as e:
+        print(f"[WARN] Failed to read config file {config_path}: {e}")
+        return False

utils/template_helpers.py

@@ -0,0 +1,179 @@
+import os
+import re
+import socket
+import time
+from typing import Any, Dict, List, Set
+from urllib.parse import urlparse
+
+
+def find_template_vars(template: str) -> Set[str]:
+    """Extract ${VAR_NAME} placeholders from a template string.
+
+    Args:
+        template: A string potentially containing ${VAR_NAME} placeholders
+
+    Returns:
+        Set of variable names found in the template
+
+    Example:
+        >>> find_template_vars("http://127.0.0.1:${GMAIL_AUTH_PORT}/api")
+        {'GMAIL_AUTH_PORT'}
+    """
+    return set(re.findall(r'\$\{(\w+)\}', str(template)))
+
+
+def render_template(template: str, values: Dict[str, int]) -> str:
+    """Replace ${VAR_NAME} placeholders with actual values.
+
+    Args:
+        template: A string containing ${VAR_NAME} placeholders
+        values: Dictionary mapping variable names to their values
+
+    Returns:
+        The template string with all placeholders replaced
+
+    Example:
+        >>> render_template("http://127.0.0.1:${PORT}/api", {"PORT": 8080})
+        'http://127.0.0.1:8080/api'
+    """
+    result = str(template)
+    for var_name, value in values.items():
+        result = result.replace(f"${{{var_name}}}", str(value))
+    return result
+
+
+def resolve_port_from_env(var_name: str) -> int:
+    """Resolve a port value from environment variable.
+
+    Args:
+        var_name: Name of the environment variable to read
+
+    Returns:
+        The port number as an integer
+
+    Raises:
+        ValueError: If the environment variable is not set
+
+    Example:
+        >>> os.environ["MY_PORT"] = "8080"
+        >>> resolve_port_from_env("MY_PORT")
+        8080
+    """
+    existing = os.getenv(var_name)
+    if existing:
+        return int(existing)
+    raise ValueError(f"Environment variable {var_name} not set.")
+
+
+def allocate_server_port(
+    server_name: str,
+    resource_manager: Any,
+    task_id: str,
+    prefix: str = "mcp"
+) -> int:
+    """Allocate a port for an MCP server.
+
+    Args:
+        server_name: Name of the server
+        resource_manager: ResourceManager instance for port allocation
+        task_id: Unique task identifier
+        prefix: Prefix for the port resource identifier (default: "mcp")
+
+    Returns:
+        Allocated port number
+    """
+    port = resource_manager.allocate_port(task_id, f"{prefix}.{server_name.lower()}")
+    print(f"[MCP] Allocated port {port} for server '{server_name}'")
+    return port
+
+
+def resolve_server_env_vars(
+    env_config: Dict[str, Any],
+    resolved_ports: Dict[str, int],
+    server_name: str
+) -> Dict[str, str]:
+    """Resolve template variables in server environment configuration.
+
+    Args:
+        env_config: Raw env dict from YAML (may contain ${VAR_NAME} templates)
+        resolved_ports: Cache of already-resolved port variables
+        server_name: Server name for logging
+
+    Returns:
+        Dict with all templates resolved to actual values
+    """
+    result = {}
+
+    for env_key, env_val in env_config.items():
+        template_vars = find_template_vars(str(env_val))
+        if template_vars:
+            # Resolve any new variables we haven't seen yet
+            for var_name in template_vars:
+                if var_name not in resolved_ports:
+                    try:
+                        resolved_ports[var_name] = resolve_port_from_env(var_name)
+                        print(f"[ENV] {var_name}={resolved_ports[var_name]}")
+                    except ValueError as e:
+                        print(f"[WARN] {e}")
+
+            # Render the template with resolved values
+            rendered = render_template(str(env_val), resolved_ports)
+            result[env_key] = rendered
+            print(f"[MCP] {server_name}: {env_key}={rendered}")
+        else:
+            # No template variables, use as-is
+            result[env_key] = str(env_val)
+
+    return result
+
+
+def build_server_name_map(servers: List[Dict[str, Any]]) -> Dict[str, str]:
+    """Build case-insensitive name mapping from server configs.
+
+    Args:
+        servers: List of server config dicts with "name" keys
+
+    Returns:
+        Dict mapping lowercase names to actual names
+    """
+    return {srv["name"].lower(): srv["name"] for srv in servers}
+
+
+def wait_for_servers_ready(
+    server_urls: Dict[str, str],
+    timeout: float = 600.0,
+    interval: float = 0.5
+) -> None:
+    """Wait for MCP servers to accept TCP connections.
+
+    Args:
+        server_urls: Dict mapping server names to URLs
+        timeout: Maximum wait time in seconds
+        interval: Polling interval in seconds
+
+    Raises:
+        RuntimeError: If servers don't become ready within timeout
+    """
+    if not server_urls:
+        return
+
+    pending = server_urls.copy()
+    deadline = time.time() + timeout
+
+    while pending and time.time() < deadline:
+        for name, url in list(pending.items()):
+            try:
+                parsed = urlparse(url)
+                host = parsed.hostname or "127.0.0.1"
+                port = parsed.port or (443 if parsed.scheme == "https" else 80)
+                with socket.create_connection((host, port), timeout=2):
+                    pending.pop(name, None)
+                    print(f"  [OK] {name} ready at {url}")
+            except OSError:
+                continue
+        if pending:
+            time.sleep(interval)
+
+    if pending:
+        names = ", ".join(sorted(pending.keys()))
+        raise RuntimeError(f"MCP servers failed to become ready: {names}")