diff --git a/cli/common.py b/cli/common.py index f5e2de0..7a3d31f 100644 --- a/cli/common.py +++ b/cli/common.py @@ -8,7 +8,13 @@ from tools.tool_registry import ToolRegistry from agentic import AssetManager -from config import load_config +from config import ( + apply_runtime_environment, + build_runtime_policy, + handle_preflight_issues, + load_config, + run_global_preflight, +) from database import create_database_strategy from llm import LLMTraceRecorder, OpenAIProvider, TracingLLMProvider from runtime import JobManager, RuntimeExecutor @@ -24,6 +30,13 @@ def bootstrap_environment( include_candidate_tools: bool = False, ) -> dict[str, Any]: cfg = load_config(config_path=config_path, prefix=prefix) + runtime_policy = build_runtime_policy(cfg) + apply_runtime_environment(runtime_policy) + handle_preflight_issues( + run_global_preflight(cfg, runtime_policy), + strict=runtime_policy.strict_preflight, + logger=logger, + ) trace_recorder = LLMTraceRecorder( env_id=_resolve_env_id(config_path, prefix), enabled=cfg.llm_tracing.enabled, @@ -61,6 +74,7 @@ def bootstrap_environment( return { "config": cfg, + "runtime_policy": runtime_policy, "database": db, "registry": registry, "job_manager": job_manager, diff --git a/config.yaml b/config.yaml index 9217691..1730236 100644 --- a/config.yaml +++ b/config.yaml @@ -1,3 +1,8 @@ +deployment: + network_mode: offline # Options: online | offline + resource_tier: full # Options: light | standard | full + + database: type: local_file # Options: local_file, mock path: ./test_data diff --git a/config/__init__.py b/config/__init__.py index 606a01a..8a9d0db 100644 --- a/config/__init__.py +++ b/config/__init__.py @@ -1,3 +1,24 @@ from .config_loader import Config, load_config +from .runtime_policy import ( + PreflightIssue, + RuntimePolicy, + VALID_NETWORK_MODES, + VALID_RESOURCE_TIERS, + apply_runtime_environment, + build_runtime_policy, + handle_preflight_issues, + run_global_preflight, +) -__all__ = ["Config", "load_config"] +__all__ = [ + "Config", + "PreflightIssue", + "RuntimePolicy", + "VALID_NETWORK_MODES", + "VALID_RESOURCE_TIERS", + "apply_runtime_environment", + "build_runtime_policy", + "handle_preflight_issues", + "load_config", + "run_global_preflight", +] diff --git a/config/config_loader.py b/config/config_loader.py index c554c08..56fdcad 100644 --- a/config/config_loader.py +++ b/config/config_loader.py @@ -59,6 +59,12 @@ class RuntimeConfig: timeout_seconds: int = 300 +@dataclass +class DeploymentConfig: + network_mode: str = "online" + resource_tier: str = "light" + + @dataclass class AgentConfig: type: str = "opencode" @@ -95,6 +101,7 @@ class Config: database: DatabaseConfig = field(default_factory=DatabaseConfig) execution: ExecutionConfig = field(default_factory=ExecutionConfig) runtime: RuntimeConfig = field(default_factory=RuntimeConfig) + deployment: DeploymentConfig = field(default_factory=DeploymentConfig) agent: AgentConfig = field(default_factory=AgentConfig) tool_llm: ToolLLMConfig = field(default_factory=ToolLLMConfig) llm_tracing: LLMTracingConfig = field(default_factory=LLMTracingConfig) @@ -109,6 +116,7 @@ def from_dict(cls, data: dict[str, Any]) -> "Config": db_data = data.get("database", {}) exec_data = data.get("execution", {}) runtime_data = data.get("runtime", {}) + deployment_data = data.get("deployment", {}) agent_data = _expand_env_placeholders(data.get("agent", {})) tool_llm_data = _expand_env_placeholders(data.get("tool_llm", {})) llm_tracing_data = data.get("llm_tracing", {}) @@ -118,6 +126,7 @@ def from_dict(cls, data: dict[str, Any]) -> "Config": database=DatabaseConfig(**db_data), execution=ExecutionConfig(**exec_data), runtime=RuntimeConfig(**runtime_data), + deployment=DeploymentConfig(**deployment_data), agent=AgentConfig(**agent_data), tool_llm=ToolLLMConfig(**tool_llm_data), llm_tracing=LLMTracingConfig(**llm_tracing_data), diff --git a/config/runtime_policy.py b/config/runtime_policy.py new file mode 100644 index 0000000..6ea5d72 --- /dev/null +++ b/config/runtime_policy.py @@ -0,0 +1,204 @@ +from __future__ import annotations + +import ipaddress +import os +from dataclasses import dataclass +from typing import Any +from urllib.parse import urlparse + + +VALID_NETWORK_MODES = {"online", "offline"} +VALID_RESOURCE_TIERS = {"light", "standard", "full"} + + +@dataclass(frozen=True) +class RuntimePolicy: + network_mode: str = "online" + resource_tier: str = "light" + + @property + def online(self) -> bool: + return self.network_mode == "online" + + @property + def offline(self) -> bool: + return self.network_mode == "offline" + + @property + def light_resource_tier(self) -> bool: + return self.resource_tier == "light" + + @property + def standard_resource_tier(self) -> bool: + return self.resource_tier == "standard" + + @property + def full_resource_tier(self) -> bool: + return self.resource_tier == "full" + + @property + def model_policy(self) -> str: + return "local_only" if self.offline else "allow_download" + + @property + def strict_preflight(self) -> bool: + return self.offline + + +@dataclass(frozen=True) +class PreflightIssue: + level: str + code: str + message: str + checker_name: str | None = None + + def format(self) -> str: + prefix = f"[{self.code}]" + if self.checker_name: + prefix += f" {self.checker_name}:" + return f"{prefix} {self.message}" + + +def build_runtime_policy(config: Any) -> RuntimePolicy: + deployment = _get_section(config, "deployment") + network_mode = str(_get_value(deployment, "network_mode", "online") or "online").strip().lower() + resource_tier = str(_get_value(deployment, "resource_tier", "light") or "light").strip().lower() + return RuntimePolicy(network_mode=network_mode, resource_tier=resource_tier) + + +def apply_runtime_environment(policy: RuntimePolicy) -> None: + if not policy.offline: + return + os.environ.setdefault("HF_HUB_OFFLINE", "1") + os.environ.setdefault("TRANSFORMERS_OFFLINE", "1") + os.environ.setdefault("HF_DATASETS_OFFLINE", "1") + os.environ.setdefault("DATAELF_OFFLINE_MODE", "1") + + +def run_global_preflight(config: Any, policy: RuntimePolicy) -> list[PreflightIssue]: + issues: list[PreflightIssue] = [] + + if policy.network_mode not in VALID_NETWORK_MODES: + issues.append(PreflightIssue( + level="error", + code="invalid_network_mode", + message=( + f"deployment.network_mode must be one of {sorted(VALID_NETWORK_MODES)}, " + f"got {policy.network_mode!r}." + ), + )) + + if policy.resource_tier not in VALID_RESOURCE_TIERS: + issues.append(PreflightIssue( + level="error", + code="invalid_resource_tier", + message=( + f"deployment.resource_tier must be one of {sorted(VALID_RESOURCE_TIERS)}, " + f"got {policy.resource_tier!r}." + ), + )) + + if not policy.offline: + # TODO: (network_mode) Add best-effort online dependency checks here as tools + # expose cheap, non-network preflight hooks. + return issues + + + # TODO: (network_mode) Re-enable strict offline endpoint checks after local + # LLM deployment is ready. Development still uses an external relay LLM. + # issues.extend(_validate_offline_llm_endpoint(config, "agent", required=_agent_requires_llm(config))) + # issues.extend(_validate_offline_llm_endpoint(config, "tool_llm", required=_tool_llm_configured(config))) + return issues + + +def handle_preflight_issues( + issues: list[PreflightIssue], + *, + strict: bool, + logger: Any | None = None, +) -> None: + if not issues: + return + + errors = [issue for issue in issues if issue.level == "error"] + warnings = [issue for issue in issues if issue.level != "error"] + + for issue in warnings: + if logger is not None and hasattr(logger, "warning"): + logger.warning(f"Preflight warning: {issue.format()}") + + blocking = errors + (warnings if strict else []) + if blocking: + message = "Preflight failed:\n" + "\n".join(f" - {issue.format()}" for issue in blocking) + raise RuntimeError(message) + + +def _validate_offline_llm_endpoint( + config: Any, + section_name: str, + *, + required: bool, +) -> list[PreflightIssue]: + section = _get_section(config, section_name) + base_url = _get_value(section, "base_url", None) + if not base_url: + if not required: + return [] + return [PreflightIssue( + level="error", + code="offline_missing_llm_endpoint", + message=( + f"deployment.network_mode=offline requires {section_name}.base_url " + "to point to a local or intranet OpenAI-compatible service." + ), + )] + + if _is_public_ip_endpoint(str(base_url)): + return [PreflightIssue( + level="error", + code="offline_public_llm_endpoint", + message=( + f"deployment.network_mode=offline cannot use public endpoint " + f"{section_name}.base_url={base_url!r}." + ), + )] + # TODO: (network_mode) Add optional endpoint allowlisting once the first + # deployment config needs approved intranet hostnames beyond private IPs. + return [] + + +def _agent_requires_llm(config: Any) -> bool: + agent = _get_section(config, "agent") + return _get_value(agent, "type", "opencode") == "opencode" + + +def _tool_llm_configured(config: Any) -> bool: + tool_llm = _get_section(config, "tool_llm") + return bool(_get_value(tool_llm, "model", None) or _get_value(tool_llm, "base_url", None)) + + +def _is_public_ip_endpoint(url: str) -> bool: + parsed = urlparse(url) + host = parsed.hostname + if not host: + return False + if host in {"localhost", "127.0.0.1", "::1"}: + return False + try: + ip = ipaddress.ip_address(host) + except ValueError: + # Hostnames may be intranet DNS names; keep the first version flexible. + return False + return not (ip.is_private or ip.is_loopback or ip.is_link_local) + + +def _get_section(config: Any, name: str) -> Any: + if isinstance(config, dict): + return config.get(name, {}) + return getattr(config, name, {}) + + +def _get_value(section: Any, name: str, default: Any = None) -> Any: + if isinstance(section, dict): + return section.get(name, default) + return getattr(section, name, default) diff --git a/runtime/executor.py b/runtime/executor.py index 7478dcf..079a9c1 100644 --- a/runtime/executor.py +++ b/runtime/executor.py @@ -382,7 +382,7 @@ def execute(self, job_id: str, pipeline: str) -> dict[str, Any]: # Save pipeline to file for reference pipeline_dir = Path("pipelines") pipeline_dir.mkdir(exist_ok=True) - pipeline_file = pipeline_dir / f"{job_id}.py" + pipeline_file = pipeline_dir / f"{job_id}.dsl" with open(pipeline_file, "w") as f: f.write(pipeline) diff --git a/tools/security_audit/checker/heuristic/graceful_backdoor_detection.py b/tools/security_audit/checker/heuristic/graceful_backdoor_detection.py index 0fe0c6a..92fd72d 100644 --- a/tools/security_audit/checker/heuristic/graceful_backdoor_detection.py +++ b/tools/security_audit/checker/heuristic/graceful_backdoor_detection.py @@ -34,6 +34,7 @@ def __init__( min_batch_size: int = _MIN_BATCH_SIZE, dct_keep_ratio: float = 0.125, device: str = "auto", + local_files_only: bool = False, **kwargs, ): """ @@ -45,6 +46,7 @@ def __init__( min_batch_size: Minimum samples required for clustering. dct_keep_ratio: Keep top-left ratio for DCT (default 1/8). device: "auto" | "cuda" | "cpu". + local_files_only: only load local model files; used by offline mode. """ super().__init__(**kwargs) self.victim_config = victim_config or {} @@ -54,6 +56,7 @@ def __init__( self.min_batch_size = min_batch_size self.dct_keep_ratio = dct_keep_ratio self.device = device + self.local_files_only = local_files_only self._victim: Optional[CasualLLMVictim] = None def check(self, sample: DataSample) -> CheckResult: @@ -142,6 +145,7 @@ def _load_victim(self) -> CasualLLMVictim: config["device"] = "gpu" elif self.device == "cpu": config["device"] = "cpu" + config["local_files_only"] = self.local_files_only self._victim = load_victim(config) return self._victim diff --git a/tools/security_audit/checker/heuristic/graceful_helpers.py b/tools/security_audit/checker/heuristic/graceful_helpers.py index 7b1737d..427644f 100644 --- a/tools/security_audit/checker/heuristic/graceful_helpers.py +++ b/tools/security_audit/checker/heuristic/graceful_helpers.py @@ -59,6 +59,7 @@ def __init__( model: Optional[str] = "llama", path: Optional[str] = None, max_len: Optional[int] = 4096, + local_files_only: bool = False, **kwargs, ): if not path: @@ -67,16 +68,18 @@ def __init__( "cuda" if device in {"gpu", "cuda"} and torch.cuda.is_available() else "cpu" ) self.model_type = model - self.model_config = AutoConfig.from_pretrained(path) + self.local_files_only = local_files_only + self.model_config = AutoConfig.from_pretrained(path, local_files_only=local_files_only) self.llm = AutoModelForCausalLM.from_pretrained( path, config=self.model_config, trust_remote_code=True, device_map="auto" if self.device.type == "cuda" else None, + local_files_only=local_files_only, ) if self.device.type != "cuda": self.llm = self.llm.to(self.device) - self.tokenizer = AutoTokenizer.from_pretrained(path) + self.tokenizer = AutoTokenizer.from_pretrained(path, local_files_only=local_files_only) if self.model_type == "llama": pad_token = self.tokenizer.unk_token or self.tokenizer.eos_token self.llm.config.pad_token_id = self.tokenizer.convert_tokens_to_ids(pad_token) diff --git a/tools/security_audit/checker/model_based/bias_classifier.py b/tools/security_audit/checker/model_based/bias_classifier.py index af30d98..6d362a5 100644 --- a/tools/security_audit/checker/model_based/bias_classifier.py +++ b/tools/security_audit/checker/model_based/bias_classifier.py @@ -31,6 +31,7 @@ def __init__( threshold: float = 0.5, max_length: int = 512, device: str = "auto", + local_files_only: bool = False, ): """ Args: @@ -38,19 +39,21 @@ def __init__( threshold: per-category score above which text is flagged as biased. max_length: max token length for truncation. device: "auto", "cuda", or "cpu". + local_files_only: only load local model files; used by offline mode. """ super().__init__() self.model_name_or_path = model_name_or_path self.threshold = threshold self.max_length = max_length self._device = device + self.local_files_only = local_files_only self.model = None def load_model(self): if self.model is not None: return try: - from transformers import pipeline + from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline import torch if self._device == "auto": @@ -61,9 +64,18 @@ def load_model(self): device = -1 self._log.info(f"BiasClassifier: loading model '{self.model_name_or_path}' ...") + tokenizer = AutoTokenizer.from_pretrained( + self.model_name_or_path, + local_files_only=self.local_files_only, + ) + model = AutoModelForSequenceClassification.from_pretrained( + self.model_name_or_path, + local_files_only=self.local_files_only, + ) self.model = pipeline( "text-classification", - model=self.model_name_or_path, + model=model, + tokenizer=tokenizer, top_k=None, # return all label scores device=device, truncation=True, diff --git a/tools/security_audit/checker/model_based/harmful_classifier.py b/tools/security_audit/checker/model_based/harmful_classifier.py index 990302e..56682ed 100644 --- a/tools/security_audit/checker/model_based/harmful_classifier.py +++ b/tools/security_audit/checker/model_based/harmful_classifier.py @@ -44,6 +44,7 @@ def __init__( device: str = "auto", dtype: str = "bfloat16", threshold: float = 0.5, + local_files_only: bool = False, ): """ Args: @@ -51,12 +52,14 @@ def __init__( device: "auto", "cuda", "cpu". dtype: torch dtype string, "bfloat16" or "float16". threshold: unsafe probability above which the sample is flagged. + local_files_only: only load local model files; used by offline mode. """ super().__init__() self.model_name_or_path = model_name_or_path self.device = device self.dtype = getattr(torch, dtype, torch.bfloat16) self.threshold = threshold + self.local_files_only = local_files_only self._tokenizer = None self.model = None self._safe_token_id = None @@ -67,19 +70,24 @@ def load_model(self): return try: self._log.info("HarmfulContentClassifier: loading Llama Guard model ...") - self._tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path) + self._tokenizer = AutoTokenizer.from_pretrained( + self.model_name_or_path, + local_files_only=self.local_files_only, + ) if self.device in ("auto", "cuda"): # Use device_map for GPU (handles multi-GPU sharding) self.model = AutoModelForCausalLM.from_pretrained( self.model_name_or_path, torch_dtype=self.dtype, device_map=self.device, + local_files_only=self.local_files_only, ) else: # CPU: load without device_map to avoid meta tensor issues self.model = AutoModelForCausalLM.from_pretrained( self.model_name_or_path, torch_dtype=self.dtype, + local_files_only=self.local_files_only, ).to(self.device) self.model.eval() # Pre-compute token IDs for "safe" and "unsafe" diff --git a/tools/security_audit/checker/model_based/jailbreak_classifier.py b/tools/security_audit/checker/model_based/jailbreak_classifier.py index cf3d13b..dcf129c 100644 --- a/tools/security_audit/checker/model_based/jailbreak_classifier.py +++ b/tools/security_audit/checker/model_based/jailbreak_classifier.py @@ -94,17 +94,20 @@ def __init__( model_name_or_path: str = "allenai/wildguard", max_new_tokens: int = 32, device: str = "auto", + local_files_only: bool = False, ): """ Args: model_name_or_path: HuggingFace model ID or local path. max_new_tokens: token budget for WildGuard's generated verdict. device: "auto", "cuda", or "cpu". + local_files_only: only load local model files; used by offline mode. """ super().__init__() self.model_name_or_path = model_name_or_path self.max_new_tokens = max_new_tokens self._device = device + self.local_files_only = local_files_only self.model = None self._tokenizer = None @@ -123,11 +126,16 @@ def load_model(self): dtype = torch.float16 if device == "cuda" else torch.float32 self._log.info(f"JailbreakClassifier: loading model '{self.model_name_or_path}' on {device} ...") - self._tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=False) + self._tokenizer = AutoTokenizer.from_pretrained( + self.model_name_or_path, + use_fast=False, + local_files_only=self.local_files_only, + ) self.model = AutoModelForCausalLM.from_pretrained( self.model_name_or_path, torch_dtype=dtype, low_cpu_mem_usage=True, + local_files_only=self.local_files_only, ).to(device) self.model.eval() self._device_obj = device diff --git a/tools/security_audit/checker/model_based/pii_ner.py b/tools/security_audit/checker/model_based/pii_ner.py index 37cdf27..7b84b3e 100644 --- a/tools/security_audit/checker/model_based/pii_ner.py +++ b/tools/security_audit/checker/model_based/pii_ner.py @@ -1,7 +1,8 @@ # Model: Microsoft Presidio (regex + spaCy NER) # https://github.com/microsoft/presidio -from typing import Dict, List, Optional, Tuple +import os +from typing import Dict, List, Optional, Tuple from ..base import ModelBasedChecker from ..registry import CheckerRegistry @@ -38,10 +39,15 @@ def __init__( super().__init__() try: - spacy_model = "en_core_web_lg" if language == "en" else f"{language}_core_web_sm" - if not spacy.util.is_package(spacy_model): - self._log.info(f"spaCy model '{spacy_model}' not found, downloading ...") - spacy.cli.download(spacy_model) + spacy_model = "en_core_web_lg" if language == "en" else f"{language}_core_web_sm" + if not spacy.util.is_package(spacy_model): + if os.environ.get("DATAELF_OFFLINE_MODE") == "1": + raise RuntimeError( + f"spaCy model '{spacy_model}' is required in offline mode; " + "install it before running PIINERDetector." + ) + self._log.info(f"spaCy model '{spacy_model}' not found, downloading ...") + spacy.cli.download(spacy_model) self.model = AnalyzerEngine() self.anonymizer = AnonymizerEngine() diff --git a/tools/security_audit/checker/model_based/prompt_injection_classifier.py b/tools/security_audit/checker/model_based/prompt_injection_classifier.py index 2e078c8..718eda2 100644 --- a/tools/security_audit/checker/model_based/prompt_injection_classifier.py +++ b/tools/security_audit/checker/model_based/prompt_injection_classifier.py @@ -23,6 +23,7 @@ def __init__( threshold: float = 0.5, max_length: int = 512, device: str = "auto", + local_files_only: bool = False, ): """ Args: @@ -30,12 +31,14 @@ def __init__( threshold: injection probability score above which input is flagged. max_length: max token length passed to the tokenizer (PIGuard supports up to 2048). device: "auto", "cuda", or "cpu". + local_files_only: only load local model files; used by offline mode. """ super().__init__() self.model_name_or_path = model_name_or_path self.threshold = threshold self.max_length = max_length self._device = device + self.local_files_only = local_files_only self.model = None def load_model(self): @@ -56,10 +59,12 @@ def load_model(self): tokenizer = AutoTokenizer.from_pretrained( self.model_name_or_path, model_max_length=self.max_length, + local_files_only=self.local_files_only, ) model = AutoModelForSequenceClassification.from_pretrained( self.model_name_or_path, trust_remote_code=True, + local_files_only=self.local_files_only, ) self.model = pipeline( "text-classification", diff --git a/tools/security_audit/config.py b/tools/security_audit/config.py index 62474d7..44c6908 100644 --- a/tools/security_audit/config.py +++ b/tools/security_audit/config.py @@ -1,40 +1,41 @@ -from typing import Dict, List, Optional - -from pydantic import BaseModel - - -class LLMConfig(BaseModel): - """LLM service config for LLM-as-Judge checkers.""" - model: str = "" - api_key: str = "" - api_url: str = "" - temperature: float = 0.0 - max_tokens: int = 2048 - - -class CheckerConfig(BaseModel): - """Config for a single checker.""" - name: str - enabled: bool = True - params: Dict = {} - - -class ExecutorConfig(BaseModel): - """Executor engine config.""" - max_workers: int = 4 - batch_size: int = 100 - start_index: int = 0 - end_index: int = -1 # -1 means process all - - -class AuditConfig(BaseModel): - """Top-level audit config.""" - task_name: str = "security_audit" - output_path: str = "outputs/" - log_level: str = "INFO" - - executor: ExecutorConfig = ExecutorConfig() - llm: Optional[LLMConfig] = None # llm model name (llm-based checkers required) - models: Dict[str, str] = {} # model name or path (model-based checkers required) - checkers: List[CheckerConfig] = [] - checker_tags: List[str] = [] +from typing import Dict, List, Optional + +from pydantic import BaseModel + + +class LLMConfig(BaseModel): + """LLM service config for LLM-as-Judge checkers.""" + model: str = "" + api_key: str = "" + api_url: str = "" + temperature: float = 0.0 + max_tokens: int = 2048 + + +class CheckerConfig(BaseModel): + """Config for a single checker.""" + name: str + enabled: bool = True + params: Dict = {} + selection_source: str = "config" # explicit | config | auto + + +class ExecutorConfig(BaseModel): + """Executor engine config.""" + max_workers: int = 4 + batch_size: int = 100 + start_index: int = 0 + end_index: int = -1 # -1 means process all + + +class AuditConfig(BaseModel): + """Top-level audit config.""" + task_name: str = "security_audit" + output_path: str = "outputs/" + log_level: str = "INFO" + + executor: ExecutorConfig = ExecutorConfig() + llm: Optional[LLMConfig] = None # llm model name (llm-based checkers required) + models: Dict[str, str] = {} # model name or path (model-based checkers required) + checkers: List[CheckerConfig] = [] + checker_tags: List[str] = [] diff --git a/tools/security_audit/default.yaml b/tools/security_audit/default.yaml index f853012..1b47c5d 100644 --- a/tools/security_audit/default.yaml +++ b/tools/security_audit/default.yaml @@ -23,6 +23,26 @@ checkers: # - DPOLabelFlipLLMJudge # - name: JailbreakLLMJudge # enabled: false + - name: PIINERDetector + enabled: false + params: + language: en + - name: JailbreakClassifier + enabled: false + params: + model_name_or_path: /mnt/shared-storage-gpfs2/gpfs2-shared-public/huggingface/zskj-hub/models--allenai--wildguard + - name: PromptInjectionClassifier + enabled: false + params: + model_name_or_path: leolee99/PIGuard + - name: HarmfulContentClassifier + enabled: false + params: + model_name_or_path: /mnt/shared-storage-gpfs2/gpfs2-shared-public/huggingface/zskj-hub/models-meta-llama-Llama-Guard-3-8B + - name: BiasClassifier + enabled: false + params: + model_name_or_path: cirimus/modernbert-large-bias-type-classifier - name: GraCeFulBackdoorDefender enabled: false params: @@ -48,9 +68,3 @@ risk_weights: jailbreak: 10 sycophancy: 2 -# model name/path for model-based checkers (HuggingFace model ID or local path) -models: - jailbreak_classifier: /mnt/shared-storage-gpfs2/gpfs2-shared-public/huggingface/zskj-hub/models--allenai--wildguard - prompt_injection_classifier: leolee99/PIGuard - harmful_content_classifier: /mnt/shared-storage-gpfs2/gpfs2-shared-public/huggingface/zskj-hub/models-meta-llama-Llama-Guard-3-8B - bias_classifier: cirimus/modernbert-large-bias-type-classifier diff --git a/tools/security_audit/policy.py b/tools/security_audit/policy.py new file mode 100644 index 0000000..b3eeacc --- /dev/null +++ b/tools/security_audit/policy.py @@ -0,0 +1,633 @@ +from __future__ import annotations + +import importlib.util +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from config import PreflightIssue, RuntimePolicy + +from .checker.registry import CheckerRegistry +from .config import CheckerConfig + + +_RULE_BASED_CHECKERS = [ + "PIIRule", + "SecretRule", + "ToxicityKeywordRule", + "HarmfulKeywordRule", + "BiasKeywordRule", +] + +_LLM_JUDGE_CHECKERS = [ + "HarmfulContentLLMJudge", + "BiasLLMJudge", + "ToxicityLLMJudge", + "PIILLMJudge", + "SycophancyLLMJudge", + "PromptInjectionLLMJudge", + "JailbreakLLMJudge", + "FactualInconsistancyLLMJudge", + "SelfContradictionLLMJudge", + "InstructionMismatchLLMJudge", + "DPOLabelFlipLLMJudge", +] + +_STANDARD_MODEL_CHECKERS = [ + "PIINERDetector", +] + +_HEAVY_MODEL_CHECKERS = [ + "HarmfulContentClassifier", + "ToxicityClassifier", + "BiasClassifier", + "JailbreakClassifier", + "PromptInjectionClassifier", + "GraCeFulBackdoorDefender", +] + +_HEAVY_CHECKERS = set(_HEAVY_MODEL_CHECKERS) + +_RESOURCE_TIER_ORDER = {"light": 0, "standard": 1, "full": 2} + +_CHECKER_MIN_RESOURCE_TIERS = { + **{name: "light" for name in _RULE_BASED_CHECKERS}, + **{name: "standard" for name in _LLM_JUDGE_CHECKERS}, + **{name: "standard" for name in _STANDARD_MODEL_CHECKERS}, + **{name: "full" for name in _HEAVY_MODEL_CHECKERS}, +} + +_LOCAL_MODEL_PATH_CHECKERS = { + "HarmfulContentClassifier", + "JailbreakClassifier", + "PromptInjectionClassifier", + "BiasClassifier", +} + +_SPACY_MODEL_BY_LANGUAGE = { + "en": "en_core_web_lg", +} + +_TOKENIZER_FILES = { + "tokenizer.json", + "tokenizer.model", + "vocab.json", + "vocab.txt", + "merges.txt", + "spiece.model", +} + +_WEIGHT_SUFFIXES = (".safetensors", ".bin") + + +@dataclass(frozen=True) +class CheckerRequest: + raw_checker_names: Any = None + has_checker_names: bool = False + strategy: str = "single_stage" + + @property + def explicit(self) -> bool: + return self.has_checker_names + + @property + def checker_names(self) -> list[str]: + if isinstance(self.raw_checker_names, list): + return [name for name in self.raw_checker_names if isinstance(name, str)] + return [] + + +@dataclass +class CheckerCapability: + name: str + allowed: bool + required_tier: str | None = None + params: dict[str, Any] = field(default_factory=dict) + issues: list[PreflightIssue] = field(default_factory=list) + + @property + def blocked_reason(self) -> str | None: + if self.allowed or not self.issues: + return None + return self.issues[0].code + + +@dataclass +class CheckerCapabilitySet: + capabilities: dict[str, CheckerCapability] + + def get(self, checker_name: str) -> CheckerCapability | None: + return self.capabilities.get(checker_name) + + def is_allowed(self, checker_name: str) -> bool: + capability = self.get(checker_name) + return bool(capability and capability.allowed) + + +@dataclass +class ResolvedCheckerPlan: + strategy: str + checker_configs: list[CheckerConfig] + skipped_checkers: list[dict[str, str]] = field(default_factory=list) + source: str = "auto" + + +def build_checker_request(kwargs: dict, tool_defaults: dict) -> CheckerRequest: + # TODO: (resource_tier) Let security_audit strategies accept explicit + # strategy names once funnel policies are implemented. + strategy = str(kwargs.get("strategy") or tool_defaults.get("strategy") or "single_stage") + return CheckerRequest( + raw_checker_names=kwargs.get("checker_names"), + has_checker_names="checker_names" in kwargs, + strategy=strategy, + ) + + +def validate_checker_request( + request: CheckerRequest, + runtime_policy: RuntimePolicy, + context_config: dict[str, Any], +) -> list[PreflightIssue]: + issues: list[PreflightIssue] = [] + if not isinstance(request.strategy, str) or not request.strategy.strip(): + issues.append(PreflightIssue( + level="error", + code="invalid_security_audit_strategy", + message="security_audit strategy must be a non-empty string.", + )) + + if not request.has_checker_names: + return issues + + if not isinstance(request.raw_checker_names, list): + return issues + [PreflightIssue( + level="error", + code="invalid_checker_names", + message="checker_names must be a list of checker class names.", + )] + + if not request.raw_checker_names: + return issues + [PreflightIssue( + level="error", + code="empty_checker_names", + message="checker_names cannot be empty when explicitly provided.", + )] + + available = _available_checker_names() + for raw_name in request.raw_checker_names: + if not isinstance(raw_name, str) or not raw_name.strip(): + issues.append(PreflightIssue( + level="error", + code="invalid_checker_name", + message=f"checker_names entries must be non-empty strings, got {raw_name!r}.", + )) + continue + name = raw_name.strip() + if name not in available: + issues.append(PreflightIssue( + level="error", + code="unknown_checker", + checker_name=name, + message=f"Checker `{name}` is not registered.", + )) + continue + resource_issue = _validate_checker_resource_tier(name, runtime_policy) + if resource_issue: + issues.append(resource_issue) + if runtime_policy.offline and name.endswith("LLMJudge") and not _has_local_llm_config(context_config): + issues.append(_offline_missing_llm_issue(name)) + return issues + + +def build_checker_capability_set( + *, + tool_defaults: dict, + runtime_policy: RuntimePolicy, + context_config: dict[str, Any], +) -> CheckerCapabilitySet: + default_configs = load_default_checker_configs(tool_defaults) + params_by_name = {config.name: dict(config.params) for config in default_configs} + capabilities: dict[str, CheckerCapability] = {} + + for name in sorted(_available_checker_names()): + issues: list[PreflightIssue] = [] + resource_issue = _validate_checker_resource_tier(name, runtime_policy) + if resource_issue: + issues.append(resource_issue) + else: + # Resource tier is checked first because it is cheap. Network/offline + # checks only run for resource-eligible checkers. + issues.extend(validate_checker_network_availability( + checker_config=CheckerConfig( + name=name, + params=dict(params_by_name.get(name, {})), + selection_source="capability", + ), + runtime_policy=runtime_policy, + context_config=context_config, + )) + + capabilities[name] = CheckerCapability( + name=name, + allowed=not any(issue.level == "error" for issue in issues), + required_tier=_CHECKER_MIN_RESOURCE_TIERS.get(name), + params=dict(params_by_name.get(name, {})), + issues=issues, + ) + return CheckerCapabilitySet(capabilities=capabilities) + + +def resolve_checker_plan( + *, + request: CheckerRequest, + capability_set: CheckerCapabilitySet, + tool_defaults: dict, + resource_tier: str, +) -> ResolvedCheckerPlan: + default_configs = load_default_checker_configs(tool_defaults) + defaults_by_name = {config.name: config for config in default_configs} + + if request.explicit: + checker_configs = [] + for name in request.checker_names: + default_config = defaults_by_name.get(name) + checker_configs.append(CheckerConfig( + name=name, + enabled=True, + params=dict(default_config.params) if default_config else {}, + selection_source="explicit", + )) + return ResolvedCheckerPlan( + strategy=request.strategy, + checker_configs=checker_configs, + source="explicit", + ) + + if default_configs: + checker_configs: list[CheckerConfig] = [] + skipped: list[dict[str, str]] = [] + for config in default_configs: + if not config.enabled: + continue + if capability_set.is_allowed(config.name): + checker_configs.append(config.copy(deep=True)) + continue + capability = capability_set.get(config.name) + skipped.append({ + "name": config.name, + "reason": capability.blocked_reason if capability else "unknown_checker", + }) + return ResolvedCheckerPlan( + strategy=request.strategy, + checker_configs=checker_configs, + skipped_checkers=skipped, + source="config", + ) + + checker_configs = [] + skipped = [] + for name in resolve_default_checkers_for_resource_tier(resource_tier): + if capability_set.is_allowed(name): + checker_configs.append(CheckerConfig(name=name, selection_source="auto")) + else: + capability = capability_set.get(name) + skipped.append({ + "name": name, + "reason": capability.blocked_reason if capability else "unknown_checker", + }) + return ResolvedCheckerPlan( + strategy=request.strategy, + checker_configs=checker_configs, + skipped_checkers=skipped, + source="auto", + ) + + +def validate_resolved_checker_plan( + *, + plan: ResolvedCheckerPlan, + capability_set: CheckerCapabilitySet, + runtime_policy: RuntimePolicy, + context_config: dict[str, Any], +) -> list[PreflightIssue]: + issues: list[PreflightIssue] = [] + enabled_configs = [config for config in plan.checker_configs if config.enabled] + if not enabled_configs: + issues.append(PreflightIssue( + level="error", + code="no_enabled_checkers", + message="security_audit resolved plan has no enabled checkers.", + )) + return issues + + available = _available_checker_names() + for checker_config in enabled_configs: + name = checker_config.name + if name not in available: + issues.append(PreflightIssue( + level="error", + code="unknown_checker", + checker_name=name, + message=f"Checker `{name}` is not registered.", + )) + continue + + capability = capability_set.get(name) + if capability is None or not capability.allowed: + capability_issues = capability.issues if capability else [] + issues.extend(capability_issues or [PreflightIssue( + level="error", + code="checker_not_allowed", + checker_name=name, + message=f"Checker `{name}` is not allowed by the current deployment policy.", + )]) + continue + + # TODO: (network_mode) Extend resolved-plan checks to cover checker + # implementations that can still perform implicit downloads internally. + if runtime_policy.offline and _requires_transformers_local_files_only(name): + if checker_config.params.get("local_files_only") is not True: + issues.append(PreflightIssue( + level="error", + code="offline_checker_missing_local_files_only", + checker_name=name, + message=( + f"Checker `{name}` must receive params.local_files_only=True " + "when deployment.network_mode=offline." + ), + )) + return issues + + +def load_default_checker_configs(tool_defaults: dict) -> list[CheckerConfig]: + raw = tool_defaults.get("checkers") + if not isinstance(raw, list): + return [] + + configs: list[CheckerConfig] = [] + for item in raw: + if isinstance(item, str): + configs.append(CheckerConfig(name=item, selection_source="config")) + elif isinstance(item, dict) and "name" in item: + configs.append(CheckerConfig(**{**item, "selection_source": "config"})) + return configs + + +def _available_checker_names() -> set[str]: + return set(CheckerRegistry.list_all()) + + +def _validate_checker_resource_tier( + checker_name: str, + runtime_policy: RuntimePolicy, +) -> PreflightIssue | None: + required_tier = _CHECKER_MIN_RESOURCE_TIERS.get(checker_name) + if required_tier is None: + return None + + current_tier = runtime_policy.resource_tier + current_rank = _RESOURCE_TIER_ORDER.get(current_tier) + required_rank = _RESOURCE_TIER_ORDER[required_tier] + if current_rank is None or current_rank >= required_rank: + return None + + return PreflightIssue( + level="error", + code="checker_resource_tier_too_low", + checker_name=checker_name, + message=( + f"Checker `{checker_name}` requires deployment.resource_tier >= {required_tier!r}, " + f"but current resource_tier is {current_tier!r}." + ), + ) + + +def _requires_transformers_local_files_only(checker_name: str) -> bool: + return checker_name in _LOCAL_MODEL_PATH_CHECKERS or checker_name == "GraCeFulBackdoorDefender" + + +def _offline_missing_llm_issue(checker_name: str) -> PreflightIssue: + return PreflightIssue( + level="error", + code="offline_checker_missing_llm", + checker_name=checker_name, + message=( + "LLM judge checkers require a configured local or intranet LLM " + "endpoint when deployment.network_mode=offline." + ), + ) + + +def resolve_default_checkers_for_resource_tier(resource_tier: str) -> list[str]: + normalized = (resource_tier or "light").strip().lower() + if normalized == "standard": + return [ + *_STANDARD_MODEL_CHECKERS, + *_LLM_JUDGE_CHECKERS, + ] + if normalized == "full": + return [ + *_STANDARD_MODEL_CHECKERS, + *_LLM_JUDGE_CHECKERS, + *_HEAVY_MODEL_CHECKERS, + ] + return list(_RULE_BASED_CHECKERS) + + +def validate_checker_network_availability( + *, + checker_config: CheckerConfig, + runtime_policy: RuntimePolicy, + context_config: dict[str, Any], +) -> list[PreflightIssue]: + if not runtime_policy.offline: + # TODO: (network_mode) Add best-effort dependency warnings for online + # mode once checker dependency metadata is finalized. + return [] + + name = checker_config.name + + if name.endswith("LLMJudge") and not _has_local_llm_config(context_config): + return [_offline_missing_llm_issue(name)] + + if name in _LOCAL_MODEL_PATH_CHECKERS: + model_path = _resolve_model_path(checker_config) + if not model_path: + return [PreflightIssue( + level="error", + code="offline_checker_missing_model_path", + checker_name=name, + message=( + "Offline model-based checker requires a local model path. " + "Set checker params.model_name_or_path in " + "tools/security_audit/default.yaml." + ), + )] + if _looks_like_hf_model_id(model_path): + return [PreflightIssue( + level="error", + code="offline_checker_uses_hf_model_id", + checker_name=name, + message=( + f"Configured model_name_or_path={model_path!r} looks like a HuggingFace model id. " + "Offline mode requires a local model directory." + ), + )] + if not _path_exists(model_path): + return [PreflightIssue( + level="error", + code="offline_checker_model_path_not_found", + checker_name=name, + message=f"Configured local model path does not exist: {model_path!r}.", + )] + issue = _validate_transformers_model_dir(name, model_path) + if issue: + return [issue] + + if name == "ToxicityClassifier": + return [PreflightIssue( + level="error", + code="offline_checker_not_local_only_ready", + checker_name=name, + message=( + "ToxicityClassifier uses Detoxify, which may fetch weights from package caches " + "or remote sources. It is disabled in offline mode until an explicit local " + "weights/cache contract is implemented." + ), + )] + + if name == "PIINERDetector": + issue = _validate_pii_ner_offline(checker_config) + if issue: + return [issue] + + if name == "GraCeFulBackdoorDefender": + victim_config = checker_config.params.get("victim_config") or {} + victim_path = victim_config.get("path") if isinstance(victim_config, dict) else None + if not victim_path: + return [PreflightIssue( + level="error", + code="offline_checker_missing_victim_model", + checker_name=name, + message=( + "GraCeFulBackdoorDefender requires params.victim_config.path " + "to point to a local victim/surrogate model in offline mode." + ), + )] + if not _path_exists(victim_path): + return [PreflightIssue( + level="error", + code="offline_checker_victim_model_not_found", + checker_name=name, + message=f"Configured victim model path does not exist: {victim_path!r}.", + )] + issue = _validate_transformers_model_dir(name, victim_path) + if issue: + return [issue] + + return [] + +def _has_local_llm_config(context_config: dict[str, Any]) -> bool: + tool_llm = _get_section(context_config, "tool_llm") + agent = _get_section(context_config, "agent") + return bool( + (_get_value(tool_llm, "model") and _get_value(tool_llm, "base_url")) + or (_get_value(agent, "model") and _get_value(agent, "base_url")) + ) + + +def _resolve_model_path(checker_config: CheckerConfig) -> str | None: + explicit = checker_config.params.get("model_name_or_path") + return str(explicit) if explicit else None + + +def _path_exists(value: str) -> bool: + return Path(value).expanduser().exists() + + +def _validate_transformers_model_dir(checker_name: str, model_path: str) -> PreflightIssue | None: + path = Path(model_path).expanduser() + if not path.is_dir(): + return PreflightIssue( + level="error", + code="offline_checker_model_path_not_directory", + checker_name=checker_name, + message=f"Configured model path must be a local directory: {model_path!r}.", + ) + if not (path / "config.json").exists(): + return PreflightIssue( + level="error", + code="offline_checker_model_missing_config", + checker_name=checker_name, + message=f"Local model directory is missing config.json: {model_path!r}.", + ) + if not any((path / filename).exists() for filename in _TOKENIZER_FILES): + return PreflightIssue( + level="error", + code="offline_checker_model_missing_tokenizer", + checker_name=checker_name, + message=f"Local model directory is missing tokenizer files: {model_path!r}.", + ) + if not any(file.is_file() and file.name.endswith(_WEIGHT_SUFFIXES) for file in path.rglob("*")): + return PreflightIssue( + level="error", + code="offline_checker_model_missing_weights", + checker_name=checker_name, + message=f"Local model directory is missing .safetensors or .bin weights: {model_path!r}.", + ) + return None + + +def _validate_pii_ner_offline(checker_config: CheckerConfig) -> PreflightIssue | None: + if importlib.util.find_spec("spacy") is None: + return PreflightIssue( + level="error", + code="offline_checker_missing_dependency", + checker_name=checker_config.name, + message="PIINERDetector requires spaCy to be installed in offline mode.", + ) + if importlib.util.find_spec("presidio_analyzer") is None: + return PreflightIssue( + level="error", + code="offline_checker_missing_dependency", + checker_name=checker_config.name, + message="PIINERDetector requires presidio-analyzer to be installed in offline mode.", + ) + if importlib.util.find_spec("presidio_anonymizer") is None: + return PreflightIssue( + level="error", + code="offline_checker_missing_dependency", + checker_name=checker_config.name, + message="PIINERDetector requires presidio-anonymizer to be installed in offline mode.", + ) + + import spacy + + language = str(checker_config.params.get("language", "en")) + spacy_model = _SPACY_MODEL_BY_LANGUAGE.get(language, f"{language}_core_web_sm") + if not spacy.util.is_package(spacy_model): + return PreflightIssue( + level="error", + code="offline_checker_missing_spacy_model", + checker_name=checker_config.name, + message=( + f"PIINERDetector requires spaCy model {spacy_model!r} to be installed. " + "Offline mode will not run spacy.cli.download()." + ), + ) + return None + + +def _looks_like_hf_model_id(value: str) -> bool: + return "/" in value and not value.startswith(("/", "./", "../", "~")) + + +def _get_section(config: dict[str, Any], name: str) -> Any: + section = config.get(name, {}) if isinstance(config, dict) else {} + return section + + +def _get_value(section: Any, name: str) -> Any: + if isinstance(section, dict): + return section.get(name) + return getattr(section, name, None) diff --git a/tools/security_audit/tool.py b/tools/security_audit/tool.py index 12442b5..964ccff 100644 --- a/tools/security_audit/tool.py +++ b/tools/security_audit/tool.py @@ -11,10 +11,19 @@ import os from typing import Any +from config import apply_runtime_environment, build_runtime_policy, handle_preflight_issues from tools.base_tool import BaseTool, ToolContext from .config import AuditConfig, CheckerConfig, ExecutorConfig, LLMConfig from .executor import Executor from .loader import load_samples +from .policy import ( + ResolvedCheckerPlan, + build_checker_capability_set, + build_checker_request, + resolve_checker_plan, + validate_checker_request, + validate_resolved_checker_plan, +) _DEFAULT_RISK_WEIGHTS = { @@ -42,29 +51,20 @@ def _get_tool_defaults(context_config: dict) -> dict: return security_defaults if isinstance(security_defaults, dict) else {} -def _resolve_checker_configs(kwargs: dict, tool_defaults: dict) -> list[CheckerConfig]: - """Resolve CheckerConfig list. +def _get_config_value(section: Any, name: str, default: Any = None) -> Any: + if isinstance(section, dict): + return section.get(name, default) + return getattr(section, name, default) - Priority: - 1. ``checkers`` list in tool_defaults (from default.yaml). - 2. Built-in fallback: [PIIRule]. - """ - names = kwargs.get("checker_names") - if names: # non-empty list → explicit override - return [CheckerConfig(name=n) for n in names] - raw = tool_defaults.get("checkers") - if raw and isinstance(raw, list): - configs = [] - for item in raw: - if isinstance(item, str): - configs.append(CheckerConfig(name=item)) - elif isinstance(item, dict) and "name" in item: - configs.append(CheckerConfig(**item)) - if configs: - return configs - - return [CheckerConfig(name="PIIRule")] +def _apply_runtime_policy_to_checker_configs( + checker_configs: list[CheckerConfig], + runtime_policy: Any, +) -> None: + if runtime_policy.model_policy != "local_only": + return + for checker_config in checker_configs: + checker_config.params["local_files_only"] = True def _calc_security_score(risk_distribution: dict, risk_weights: dict) -> float: @@ -134,31 +134,129 @@ def run(self, context: ToolContext, **kwargs: Any) -> dict[str, Any]: data: list[dict] = kwargs.get("data", []) max_workers: int = kwargs.get("max_workers", 4) + # Step 1: Load tool defaults and deployment runtime policy. tool_defaults = _get_tool_defaults(context.config) - checker_configs = _resolve_checker_configs(kwargs, tool_defaults) + runtime_policy = build_runtime_policy(context.config) + apply_runtime_environment(runtime_policy) + + # Step 2: Validate the request, build capabilities, and resolve a checker plan. + plan = self._build_execution_plan( + context=context, + kwargs=kwargs, + tool_defaults=tool_defaults, + runtime_policy=runtime_policy, + ) + + # Step 3: Log the resolved single-stage plan. + checker_configs = plan.checker_configs checker_names = [c.name for c in checker_configs if c.enabled] + self._log_plan(context, plan, checker_names, bool(kwargs.get("checker_names"))) + context.log(f"SecurityAuditTool: {len(data)} records, checkers={checker_names}") - if kwargs.get("checker_names"): + # TODO:Step 4: Execute the resolved plan. + return self._execute_single_stage_plan( + context=context, + plan=plan, + data=data, + max_workers=max_workers, + tool_defaults=tool_defaults, + checker_configs=checker_configs, + checker_names=checker_names, + ) + + def _build_execution_plan( + self, + *, + context: ToolContext, + kwargs: dict[str, Any], + tool_defaults: dict, + runtime_policy: Any, + ) -> ResolvedCheckerPlan: + # Step 2.1: Build and validate the raw checker request from tool args. + request = build_checker_request(kwargs, tool_defaults) + handle_preflight_issues( + validate_checker_request( + request=request, + runtime_policy=runtime_policy, + context_config=context.config, + ), + strict=runtime_policy.strict_preflight, + logger=context.logger, + ) + + # Step 2.2: Build the checker capability set under resource/network policy. + capability_set = build_checker_capability_set( + tool_defaults=tool_defaults, + runtime_policy=runtime_policy, + context_config=context.config, + ) + + # TODO:Step 2.3: Resolve the execution plan within the capability set. + plan = resolve_checker_plan( + request=request, + capability_set=capability_set, + tool_defaults=tool_defaults, + resource_tier=runtime_policy.resource_tier, + ) + + # TODO:Step 2.4: Inject runtime params and run final resolved-plan preflight. + _apply_runtime_policy_to_checker_configs(plan.checker_configs, runtime_policy) + handle_preflight_issues( + validate_resolved_checker_plan( + plan=plan, + capability_set=capability_set, + runtime_policy=runtime_policy, + context_config=context.config, + ), + strict=runtime_policy.strict_preflight, + logger=context.logger, + ) + return plan + + def _log_plan( + self, + context: ToolContext, + plan: ResolvedCheckerPlan, + checker_names: list[str], + explicit_checkers: bool, + ) -> None: + if explicit_checkers: context.log(f"SecurityAuditTool: using user-specified checkers: {checker_names}") - if tool_defaults.get("checkers"): + elif plan.source == "config": context.log( f"SecurityAuditTool: using checkers from default.yaml: {checker_names}. " ) else: context.log( - f"SecurityAuditTool: no checkers configured, using built-in default: {checker_names}" + f"SecurityAuditTool: using {plan.strategy} defaults: {checker_names}" ) - context.log(f"SecurityAuditTool: {len(data)} records, checkers={checker_names}") + for skipped in plan.skipped_checkers: + context.log( + f"SecurityAuditTool: skipped checker {skipped['name']} ({skipped['reason']}).", + "warning", + ) - # 1. Convert raw dicts → DataSample + def _execute_single_stage_plan( + self, + *, + context: ToolContext, + plan: ResolvedCheckerPlan, + data: list[dict], + max_workers: int, + tool_defaults: dict, + checker_configs: list[CheckerConfig], + checker_names: list[str], + ) -> dict[str, Any]: + # TODO: (resource_tier) Replace this single-stage executor with a + # strategy-aware multi-stage executor when funnel plans are implemented. samples = load_samples(data) - # 2. Build AuditConfig agent_cfg = context.config.get("agent") tool_llm_cfg = context.config.get("tool_llm", {}) llm_name: str = ( - getattr(tool_llm_cfg, "model", "") or getattr(agent_cfg, "model", "") + _get_config_value(tool_llm_cfg, "model", "") + or _get_config_value(agent_cfg, "model", "") ) task_name = f"security_audit_{context.job_id}" @@ -173,11 +271,9 @@ def run(self, context: ToolContext, **kwargs: Any) -> dict[str, Any]: models=tool_defaults.get("models") or {}, ) - # 3. Log LLM model info if cfg.llm: context.log(f"Tool LLM model: {llm_name}", "info") - # 4. Run the audit engine (writes artifacts to output_path) engine = Executor(cfg, logger=context.logger, llm=context.llm, job_id=context.job_id, mode=context.mode) engine.setup() task_report, _ = engine.run(samples) @@ -210,6 +306,21 @@ def run(self, context: ToolContext, **kwargs: Any) -> dict[str, Any]: "metadata": { "task_name": task_report.task_name, "checker_names": checker_names, + "resolved_plan": { + "schema_version": "security_audit.resolved_plan.v1", + "strategy": plan.strategy, + "source": plan.source, + "skipped_checkers": plan.skipped_checkers, + "stages": [ + { + "id": "stage_1", + "name": "single_stage", + "type": "single_stage", + "checkers": checker_names, + "input_scope": {"type": "all_samples"}, + } + ], + }, "create_time": task_report.create_time, "finish_time": task_report.finish_time, },