diff --git a/config/runtime_policy.py b/config/runtime_policy.py index c5a3980..6ea5d72 100644 --- a/config/runtime_policy.py +++ b/config/runtime_policy.py @@ -72,6 +72,7 @@ def apply_runtime_environment(policy: RuntimePolicy) -> None: 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]: @@ -103,7 +104,7 @@ def run_global_preflight(config: Any, policy: RuntimePolicy) -> list[PreflightIs return issues - # TODO: (network_mode) !!!!Re-enable strict offline endpoint checks after local + # 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))) 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/policy.py b/tools/security_audit/policy.py index ebedd64..b3eeacc 100644 --- a/tools/security_audit/policy.py +++ b/tools/security_audit/policy.py @@ -1,10 +1,13 @@ 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 @@ -61,35 +64,376 @@ "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") -def validate_selected_checkers( - *, - checker_configs: list[CheckerConfig], + +@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] = [] - for checker_config in checker_configs: - if not checker_config.enabled: + 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 - resource_issues = validate_checker_resource_tier_availability( - checker_config=checker_config, - runtime_policy=runtime_policy, + 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, ) - issues.extend(resource_issues) - if not checker_config.enabled: + 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 - if any(issue.level == "error" for issue in resource_issues): + + 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 - issues.extend(validate_checker_network_availability( - checker_config=checker_config, - runtime_policy=runtime_policy, - context_config=context_config, - )) + + # 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, @@ -104,15 +448,7 @@ def validate_checker_network_availability( name = checker_config.name if name.endswith("LLMJudge") and not _has_local_llm_config(context_config): - return [PreflightIssue( - level="error", - code="offline_checker_missing_llm", - checker_name=name, - message=( - "LLM judge checkers require a configured local or intranet LLM " - "endpoint when deployment.network_mode=offline." - ), - )] + return [_offline_missing_llm_issue(name)] if name in _LOCAL_MODEL_PATH_CHECKERS: model_path = _resolve_model_path(checker_config) @@ -127,6 +463,16 @@ def validate_checker_network_availability( "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", @@ -134,6 +480,26 @@ def validate_checker_network_availability( 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 {} @@ -155,72 +521,12 @@ def validate_checker_network_availability( 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] - # TODO: (network_mode) Extend offline checks for Detoxify local cache, - # Presidio/spaCy model availability, local_files_only propagation, and - # checker implementations that may still trigger implicit downloads. return [] - -def validate_checker_resource_tier_availability( - *, - checker_config: CheckerConfig, - runtime_policy: RuntimePolicy, -) -> list[PreflightIssue]: - - name = checker_config.name - required_tier = _CHECKER_MIN_RESOURCE_TIERS.get(name) - if required_tier is None: - return [] - - 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 [] - - source = getattr(checker_config, "selection_source", "config") - if source != "explicit": - checker_config.enabled = False - return [PreflightIssue( - level="warning", - code="checker_filtered_by_resource_tier", - checker_name=name, - message=( - f"Checker `{name}` requires deployment.resource_tier >= {required_tier!r}, " - f"but current resource_tier is {current_tier!r}; " - f"it was disabled from the {source} checker selection." - ), - )] - - return [PreflightIssue( - level="error", - code="checker_resource_tier_too_low", - checker_name=name, - message=( - f"Checker `{name}` requires deployment.resource_tier >= {required_tier!r}, " - f"but current resource_tier is {current_tier!r}." - ), - )] - - - -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 _has_local_llm_config(context_config: dict[str, Any]) -> bool: tool_llm = _get_section(context_config, "tool_llm") agent = _get_section(context_config, "agent") @@ -239,6 +545,83 @@ 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 diff --git a/tools/security_audit/tool.py b/tools/security_audit/tool.py index 0e8e4ea..964ccff 100644 --- a/tools/security_audit/tool.py +++ b/tools/security_audit/tool.py @@ -11,12 +11,19 @@ import os from typing import Any -from config import build_runtime_policy, handle_preflight_issues +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 resolve_default_checkers_for_resource_tier, validate_selected_checkers +from .policy import ( + ResolvedCheckerPlan, + build_checker_capability_set, + build_checker_request, + resolve_checker_plan, + validate_checker_request, + validate_resolved_checker_plan, +) _DEFAULT_RISK_WEIGHTS = { @@ -44,53 +51,20 @@ def _get_tool_defaults(context_config: dict) -> dict: return security_defaults if isinstance(security_defaults, dict) else {} -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 _resolve_checker_configs(kwargs: dict, tool_defaults: dict, resource_tier: str) -> list[CheckerConfig]: - """Resolve CheckerConfig list. - - Priority: - 1. Explicit ``checker_names`` from tool call. These force-enable the - named checkers while inheriting same-name params from default.yaml. - 2. Enabled ``checkers`` list in tool_defaults (from default.yaml). - 3. Resource-tier default checkers. - """ - default_configs = _load_default_checker_configs(tool_defaults) - defaults_by_name = {config.name: config for config in default_configs} - - names = kwargs.get("checker_names") - if names: - configs = [] - for name in names: - default_config = defaults_by_name.get(name) - params = dict(default_config.params) if default_config else {} - configs.append(CheckerConfig( - name=name, - enabled=True, - params=params, - selection_source="explicit", - )) - return configs - - if default_configs: - return default_configs - - return [ - CheckerConfig(name=name, selection_source="auto") - for name in resolve_default_checkers_for_resource_tier(resource_tier) - ] +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) + + +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: @@ -160,46 +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) runtime_policy = build_runtime_policy(context.config) - checker_configs = _resolve_checker_configs( - kwargs, - tool_defaults, - runtime_policy.resource_tier, + 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}") + + # 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_selected_checkers( - checker_configs=checker_configs, + validate_checker_request( + request=request, runtime_policy=runtime_policy, context_config=context.config, ), strict=runtime_policy.strict_preflight, logger=context.logger, ) - checker_names = [c.name for c in checker_configs if c.enabled] - if kwargs.get("checker_names"): + # 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}") - elif 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 " - f"{runtime_policy.resource_tier} resource-tier defaults: {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}" @@ -214,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) @@ -251,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, },