tweek 0.3.1__py3-none-any.whl → 0.4.0__py3-none-any.whl

tweek/security/local_reviewer.py

@@ -41,6 +41,16 @@ class LocalModelReviewProvider(ReviewProvider):
         self._model_name = model_name
         self._escalation_provider = escalation_provider
 
+    # Tools where the local prompt-injection classifier is effective.
+    # The DeBERTa model was trained on natural-language prompt injection,
+    # NOT on shell command evaluation.  For Bash/Edit/Write the model
+    # produces severe false positives (e.g. classifying "./run.sh 2>&1"
+    # as injection at 100% confidence).  Those tools should be handled by
+    # pattern matching + cloud LLM escalation instead.
+    _CONTENT_TOOLS: frozenset = frozenset({
+        "Read", "WebFetch", "Grep", "WebSearch",
+    })
+
     def call(self, system_prompt: str, user_prompt: str, max_tokens: int = 256) -> str:
         """Run local inference and return JSON result.
 
@@ -48,8 +58,11 @@ class LocalModelReviewProvider(ReviewProvider):
         runs local inference, and returns a JSON string in the same format
         that LLMReviewer._parse_response() expects.
 
-        If the local model is uncertain and an escalation provider is
-        available, the request is forwarded to the cloud LLM.
+        The local model is only used for content-screening tools (Read,
+        WebFetch, Grep, WebSearch) where the input is natural-language text
+        that the classifier was trained on.  For command-execution tools
+        (Bash, Edit, Write, etc.) the request is forwarded to the
+        escalation provider or returned as low-confidence safe.
 
         Args:
             system_prompt: System-level instructions (used for escalation only).
@@ -61,6 +74,23 @@ class LocalModelReviewProvider(ReviewProvider):
         """
         from tweek.security.local_model import get_local_model
 
+        # Detect the tool from the analysis prompt (e.g. "Tool: Bash")
+        tool_name = self._extract_tool(user_prompt)
+
+        # The DeBERTa prompt-injection model only works on natural-language
+        # content.  For shell commands and code, defer to cloud LLM or
+        # pattern matching.
+        if tool_name and tool_name not in self._CONTENT_TOOLS:
+            if self._escalation_provider:
+                return self._escalation_provider.call(
+                    system_prompt, user_prompt, max_tokens
+                )
+            return json.dumps({
+                "risk_level": "safe",
+                "reason": f"Local model not applicable for {tool_name} commands",
+                "confidence": 0.0,
+            })
+
         # Extract command from untrusted_command tags
         command = self._extract_command(user_prompt)
         if not command:
@@ -124,6 +154,18 @@ class LocalModelReviewProvider(ReviewProvider):
     def model_name(self) -> str:
         return self._model_name
 
+    @staticmethod
+    def _extract_tool(user_prompt: str) -> Optional[str]:
+        """Extract the tool name from the analysis prompt.
+
+        The LLMReviewer ANALYSIS_PROMPT includes a ``Tool: <name>`` line.
+
+        Returns:
+            Tool name (e.g. "Bash", "Read"), or None if not found.
+        """
+        match = re.search(r"^Tool:\s*(\S+)", user_prompt, re.MULTILINE)
+        return match.group(1) if match else None
+
     @staticmethod
     def _extract_command(user_prompt: str) -> str:
         """Extract the command from <untrusted_command> tags.

tweek/security/model_registry.py

@@ -40,7 +40,9 @@ class ModelDefinition:
     license: str = "unknown"
     size_mb: float = 0.0  # approximate download size
     files: List[str] = field(default_factory=list)
+    file_hashes: Dict[str, str] = field(default_factory=dict)  # filename -> sha256
     hf_subfolder: str = ""  # subfolder in the HF repo (e.g., "onnx")
+    hf_revision: str = "main"  # git revision (commit SHA for pinned downloads)
     requires_auth: bool = False
     default: bool = False
 
@@ -73,7 +75,12 @@ MODEL_CATALOG: Dict[str, ModelDefinition] = {
         license="Apache-2.0",
         size_mb=750.0,
         files=["model.onnx", "tokenizer.json"],
+        file_hashes={
+            "model.onnx": "f0ea7f239f765aedbde7c9e163a7cb38a79c5b8853d3f76db5152172047b228c",
+            "tokenizer.json": "752fe5f0d5678ad563e1bd2ecc1ddf7a3ba7e2024d0ac1dba1a72975e26dff2f",
+        },
         hf_subfolder="onnx",
+        hf_revision="e6535ca4ce3ba852083e75ec585d7c8aeb4be4c5",
         requires_auth=False,
         default=True,
         escalate_min_confidence=0.1,
@@ -167,11 +174,15 @@ class ModelDownloadError(Exception):
     pass
 
 
-def _build_hf_url(repo: str, filename: str, subfolder: str = "") -> str:
-    """Build a HuggingFace CDN download URL."""
+def _build_hf_url(repo: str, filename: str, subfolder: str = "", revision: str = "main") -> str:
+    """Build a HuggingFace CDN download URL.
+
+    When *revision* is a commit SHA, the URL points to an immutable
+    snapshot — the same bytes every time, safe to verify with SHA-256.
+    """
     if subfolder:
-        return f"https://huggingface.co/{repo}/resolve/main/{subfolder}/{filename}"
-    return f"https://huggingface.co/{repo}/resolve/main/{filename}"
+        return f"https://huggingface.co/{repo}/resolve/{revision}/{subfolder}/{filename}"
+    return f"https://huggingface.co/{repo}/resolve/{revision}/{filename}"
 
 
 def _get_hf_headers() -> Dict[str, str]:
@@ -234,9 +245,12 @@ def download_model(
     # Create SSL context
     ssl_context = ssl.create_default_context()
 
-    # Download each file
+    # Download each file, pinned to a specific revision for reproducibility
     for filename in definition.files:
-        url = _build_hf_url(definition.hf_repo, filename, definition.hf_subfolder)
+        url = _build_hf_url(
+            definition.hf_repo, filename,
+            definition.hf_subfolder, definition.hf_revision,
+        )
         dest = model_dir / filename
         tmp_dest = model_dir / f".{filename}.tmp"
 
@@ -258,6 +272,20 @@ def download_model(
                     if progress_callback:
                         progress_callback(filename, downloaded, total)
 
+            # Verify SHA-256 if the catalog provides an expected hash
+            expected_hash = definition.file_hashes.get(filename)
+            if expected_hash:
+                actual_hash = hashlib.sha256(tmp_dest.read_bytes()).hexdigest()
+                if actual_hash != expected_hash:
+                    tmp_dest.unlink(missing_ok=True)
+                    raise ModelDownloadError(
+                        f"SHA-256 mismatch for {filename}: "
+                        f"expected {expected_hash[:16]}..., "
+                        f"got {actual_hash[:16]}... "
+                        f"The file may be corrupted or tampered with. "
+                        f"Try again with --force, or report this issue."
+                    )
+
             # Atomic rename
             tmp_dest.rename(dest)
 
@@ -284,6 +312,8 @@ def download_model(
             raise ModelDownloadError(
                 f"Network error downloading {filename}: {e.reason}"
             ) from e
+        except ModelDownloadError:
+            raise  # Re-raise SHA mismatch without wrapping
         except Exception as e:
             tmp_dest.unlink(missing_ok=True)
             raise ModelDownloadError(
@@ -327,7 +357,7 @@ def remove_model(name: str) -> bool:
 
 
 def verify_model(name: str) -> Dict[str, bool]:
-    """Verify a model installation.
+    """Verify a model installation (file existence only).
 
     Args:
         name: Model name.
@@ -347,6 +377,42 @@ def verify_model(name: str) -> Dict[str, bool]:
 
     status["model_meta.yaml"] = (model_dir / "model_meta.yaml").exists()
 
+
+def verify_model_hashes(name: str) -> Dict[str, Optional[str]]:
+    """Verify SHA-256 integrity of an installed model's files.
+
+    Args:
+        name: Model name from the catalog.
+
+    Returns:
+        Dict mapping filename to verification status:
+        - ``"ok"`` — hash matches catalog
+        - ``"mismatch"`` — hash does not match (corrupted or tampered)
+        - ``"missing"`` — file not found on disk
+        - ``"no_hash"`` — catalog has no expected hash for this file
+        Returns empty dict if model is not in the catalog.
+    """
+    definition = MODEL_CATALOG.get(name)
+    if definition is None:
+        return {}
+
+    model_dir = get_model_dir(name)
+    results: Dict[str, Optional[str]] = {}
+
+    for filename in definition.files:
+        expected = definition.file_hashes.get(filename)
+        path = model_dir / filename
+
+        if not path.exists():
+            results[filename] = "missing"
+        elif not expected:
+            results[filename] = "no_hash"
+        else:
+            actual = hashlib.sha256(path.read_bytes()).hexdigest()
+            results[filename] = "ok" if actual == expected else "mismatch"
+
+    return results
+
     return status
 
 

tweek/skill_template/overrides-reference.md

@@ -66,7 +66,7 @@ whitelist:
 
 ## Pattern Toggles
 
-Control which of the 259 detection patterns are active.
+Control which of the 262 detection patterns are active.
 
 ### Globally Disable a Pattern
 

tweek/skills/context.py

@@ -0,0 +1,221 @@
+"""
+Tweek Skill Context Tracking
+
+Detects active skill context from Claude Code's Skill tool invocations.
+When PreToolUse sees tool_name=="Skill", the skill name is extracted from
+tool_input and written to a breadcrumb file. Subsequent tool calls within
+the same session read the breadcrumb to get skill context for tier lookup.
+
+This bridges the gap where Claude Code's hook protocol doesn't include
+skill_name: the Skill tool IS a regular tool, so PreToolUse sees it.
+
+Security properties:
+  - Session-isolated: per-session breadcrumb files (no cross-session leakage)
+  - Auto-expiring: 60-second staleness timeout
+  - Atomic writes: write-to-temp + os.rename (POSIX atomic)
+  - Restricted permissions: 0o600 on breadcrumb files
+  - Fail-safe: any error falls to no-context = default tier
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+import time
+from pathlib import Path
+from typing import Optional
+
+# Breadcrumb location — per-session files for isolation
+TWEEK_STATE_DIR = Path.home() / ".tweek" / "state"
+
+# Breadcrumb expires after 60 seconds of inactivity.
+# Skills typically issue tool calls in rapid succession; 60s is generous
+# while limiting the window for staleness-based attacks.
+STALENESS_TIMEOUT_SECONDS = 60
+
+# Maximum age before a per-session breadcrumb file is considered orphaned
+# and eligible for cleanup (1 hour).
+ORPHAN_CLEANUP_SECONDS = 3600
+
+
+def _breadcrumb_path_for_session(session_id: str, state_dir: Optional[Path] = None) -> Path:
+    """Get the breadcrumb file path for a specific session.
+
+    Uses first 12 chars of session_id to avoid excessively long filenames
+    while maintaining sufficient uniqueness.
+    """
+    prefix = session_id[:12] if session_id else "unknown"
+    base = state_dir or TWEEK_STATE_DIR
+    return base / f"active_skill_{prefix}.json"
+
+
+def write_skill_breadcrumb(
+    skill_name: str,
+    session_id: str,
+    *,
+    breadcrumb_path: Optional[Path] = None,
+) -> None:
+    """Record the active skill for this session.
+
+    Called when PreToolUse detects a Skill tool invocation.
+    Uses atomic write (temp file + rename) and restricts permissions to 0o600.
+
+    Args:
+        skill_name: The skill being invoked (from tool_input["skill"])
+        session_id: Current hook session ID for isolation
+        breadcrumb_path: Override for testing (bypasses per-session naming)
+    """
+    path = breadcrumb_path or _breadcrumb_path_for_session(session_id)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    data = {
+        "skill": skill_name,
+        "session_id": session_id,
+        "timestamp": time.time(),
+    }
+
+    # Atomic write: write to temp file, then rename (POSIX atomic)
+    fd = None
+    tmp_path = None
+    try:
+        fd, tmp_path = tempfile.mkstemp(
+            dir=str(path.parent),
+            prefix=".skill_",
+            suffix=".tmp",
+        )
+        os.write(fd, json.dumps(data).encode("utf-8"))
+        os.close(fd)
+        fd = None  # Mark as closed
+
+        # Restrict permissions before rename (owner read/write only)
+        os.chmod(tmp_path, 0o600)
+
+        # Atomic rename
+        os.rename(tmp_path, str(path))
+        tmp_path = None  # Mark as renamed (don't clean up)
+    finally:
+        # Clean up on failure
+        if fd is not None:
+            try:
+                os.close(fd)
+            except OSError:
+                pass
+        if tmp_path is not None:
+            try:
+                os.unlink(tmp_path)
+            except OSError:
+                pass
+
+
+def read_skill_context(
+    session_id: str,
+    *,
+    breadcrumb_path: Optional[Path] = None,
+    staleness_seconds: float = STALENESS_TIMEOUT_SECONDS,
+) -> Optional[str]:
+    """Read the active skill for the current session, if any.
+
+    Returns the skill name if a fresh, session-matching breadcrumb exists.
+    Returns None if no breadcrumb, wrong session, or stale.
+
+    Args:
+        session_id: Current hook session ID — must match breadcrumb
+        breadcrumb_path: Override for testing (bypasses per-session naming)
+        staleness_seconds: Max age in seconds before breadcrumb is stale
+    """
+    path = breadcrumb_path or _breadcrumb_path_for_session(session_id)
+
+    try:
+        if not path.exists():
+            return None
+
+        data = json.loads(path.read_text(encoding="utf-8"))
+
+        # Session isolation: only match same session
+        if data.get("session_id") != session_id:
+            return None
+
+        # Staleness check
+        ts = data.get("timestamp", 0)
+        if (time.time() - ts) > staleness_seconds:
+            # Expired — clean up
+            _clear_breadcrumb(path)
+            return None
+
+        return data.get("skill")
+
+    except (json.JSONDecodeError, OSError, KeyError):
+        return None
+
+
+def clear_skill_breadcrumb(
+    session_id: Optional[str] = None,
+    *,
+    breadcrumb_path: Optional[Path] = None,
+) -> None:
+    """Clear the active skill breadcrumb.
+
+    Called on session end or UserPromptSubmit if needed.
+
+    Args:
+        session_id: If provided, clears the per-session breadcrumb.
+        breadcrumb_path: Override for testing.
+    """
+    if breadcrumb_path:
+        _clear_breadcrumb(breadcrumb_path)
+    elif session_id:
+        _clear_breadcrumb(_breadcrumb_path_for_session(session_id))
+
+
+def cleanup_orphaned_breadcrumbs(
+    *,
+    state_dir: Optional[Path] = None,
+    max_age_seconds: float = ORPHAN_CLEANUP_SECONDS,
+) -> int:
+    """Remove breadcrumb files older than max_age_seconds.
+
+    Called periodically to prevent accumulation of stale session files.
+    Returns the number of files cleaned up.
+    """
+    base = state_dir or TWEEK_STATE_DIR
+    cleaned = 0
+
+    try:
+        if not base.exists():
+            return 0
+
+        now = time.time()
+        for f in base.glob("active_skill_*.json"):
+            try:
+                age = now - f.stat().st_mtime
+                if age > max_age_seconds:
+                    f.unlink(missing_ok=True)
+                    cleaned += 1
+            except OSError:
+                continue
+    except OSError:
+        pass
+
+    return cleaned
+
+
+def _clear_breadcrumb(path: Path) -> None:
+    """Remove the breadcrumb file silently."""
+    try:
+        path.unlink(missing_ok=True)
+    except OSError:
+        pass
+
+
+def extract_skill_from_tool_input(tool_input: dict) -> Optional[str]:
+    """Extract the skill name from a Skill tool's tool_input.
+
+    The Skill tool sends: {"skill": "commit", "args": "..."}
+
+    Returns the skill name or None if not a valid Skill invocation.
+    """
+    skill = tool_input.get("skill")
+    if isinstance(skill, str) and skill.strip():
+        return skill.strip()
+    return None

tweek/skills/scanner.py

@@ -6,7 +6,7 @@ installation. Reuses existing Tweek infrastructure where possible.
 
 Layers:
 1. Structure Validation    — file types, size, depth, symlinks
-2. Pattern Matching        — 259 regex patterns (reuses audit.py)
+2. Pattern Matching        — 262 regex patterns (reuses audit.py)
 3. Secret Scanning         — credential detection (reuses secret_scanner.py)
 4. AST Analysis            — forbidden imports/calls (reuses git_security.py)
 5. Prompt Injection Scan   — skill-specific instruction injection patterns
@@ -349,7 +349,7 @@ class SkillScanner:
     def _scan_patterns(
         self, skill_dir: Path, text_files: List[Path]
     ) -> ScanLayerResult:
-        """Run 259 regex patterns against all text files."""
+        """Run 262 regex patterns against all text files."""
         result = ScanLayerResult(layer_name="patterns", passed=True)
 
         try:

{tweek-0.3.1.dist-info → tweek-0.4.0.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: tweek
-Version: 0.3.1
+Version: 0.4.0
 Summary: Defense-in-depth security for AI coding assistants - protect credentials, code, and system from prompt injection attacks
 Author: Tommy Mancino
 License-Expression: Apache-2.0
 Project-URL: Homepage, https://gettweek.com
 Project-URL: Repository, https://github.com/gettweek/tweek
 Project-URL: Issues, https://github.com/gettweek/tweek/issues
-Keywords: claude,security,sandbox,ai,llm,tweek,claude-code,prompt-injection,mcp,credential-theft
+Keywords: claude,security,dry-run,ai,llm,tweek,claude-code,prompt-injection,mcp,credential-theft
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
 Classifier: Operating System :: MacOS :: MacOS X
@@ -45,6 +45,7 @@ Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pytest-cov>=4.0; extra == "dev"
 Requires-Dist: pytest-xdist>=3.5.0; extra == "dev"
+Requires-Dist: pytest-timeout>=2.2.0; extra == "dev"
 Requires-Dist: hypothesis>=6.98.0; extra == "dev"
 Requires-Dist: black>=23.0; extra == "dev"
 Requires-Dist: ruff>=0.1.0; extra == "dev"
@@ -124,7 +125,7 @@ tweek proxy setup                       # Cursor, Windsurf, Continue.dev (HTTP p
 tweek doctor
 ```
 
-That's it. Tweek auto-detects your tools, applies all 259 attack patterns across 6 defense layers, and runs 100% locally. Your code never leaves your machine.
+That's it. Tweek auto-detects your tools, applies all 262 attack patterns across 6 defense layers, and runs 100% locally. Your code never leaves your machine.
 
 ---
 
@@ -166,7 +167,7 @@ Turn 3: cat ~/.ssh/id_rsa → BLOCKED: path_escalation anomaly
 
 **Response injection** — Malicious instructions hidden in tool responses are caught at ingestion.
 
-See the full [Attack Patterns Reference](docs/ATTACK_PATTERNS.md) for all 259 patterns across 11 categories.
+See the full [Attack Patterns Reference](docs/ATTACK_PATTERNS.md) for all 262 patterns across 11 categories.
 
 ---
 
@@ -217,7 +218,7 @@ Every tool call passes through six independent screening layers. An attacker wou
 
 | Layer | What It Does |
 |-------|-------------|
-| **1. Pattern Matching** | 259 regex signatures catch known credential theft, exfiltration, and injection attacks instantly |
+| **1. Pattern Matching** | 262 regex signatures catch known credential theft, exfiltration, and injection attacks instantly |
 | **2. Rate Limiting** | Detects burst attacks, automated probing, and resource theft sequences |
 | **3. Local Prompt Injection AI** | Custom-trained AI models built specifically to classify and detect prompt injection. Run 100% on your machine — no API calls, no cloud, no latency. Small enough to be fast, accurate enough to catch what regex can't. |
 | **4. Session Tracking** | Behavioral analysis across turns detects multi-step attacks that look innocent individually |
@@ -235,7 +236,7 @@ See [Defense Layers](docs/DEFENSE_LAYERS.md) for the deep dive and [Architecture
 | [Full Feature List](docs/FEATURES.md) | Complete feature inventory |
 | [Architecture](docs/ARCHITECTURE.md) | System design and interception layers |
 | [Defense Layers](docs/DEFENSE_LAYERS.md) | Screening pipeline deep dive |
-| [Attack Patterns](docs/ATTACK_PATTERNS.md) | Full 259-pattern library reference |
+| [Attack Patterns](docs/ATTACK_PATTERNS.md) | Full 262-pattern library reference |
 | [Configuration](docs/CONFIGURATION.md) | Config files, tiers, and presets |
 | [CLI Reference](docs/CLI_REFERENCE.md) | All commands, flags, and examples |
 | [MCP Integration](docs/MCP_INTEGRATION.md) | MCP proxy and gateway setup |
@@ -244,7 +245,7 @@ See [Defense Layers](docs/DEFENSE_LAYERS.md) for the deep dive and [Architecture
 | [Credential Vault](docs/VAULT.md) | Vault setup and migration |
 | [Plugins](docs/PLUGINS.md) | Plugin development and registry |
 | [Logging](docs/LOGGING.md) | Event logging and audit trail |
-| [Sandbox](docs/SANDBOX.md) | Sandbox preview configuration |
+| [Dry-Run](docs/DRY_RUN.md) | Dry-run preview configuration |
 | [Tweek vs. Claude Code](docs/COMPARISON.md) | Feature comparison with native security |
 | [Troubleshooting](docs/TROUBLESHOOTING.md) | Common issues and fixes |