tweek 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

tweek/security/model_registry.py

@@ -0,0 +1,371 @@
+#!/usr/bin/env python3
+"""
+Tweek Local Model Registry
+
+Manages the catalog of local security models, downloads from HuggingFace,
+and handles model directory lifecycle.
+
+Models are stored in ~/.tweek/models/<model-name>/ with:
+- model.onnx       — ONNX model file
+- tokenizer.json   — Tokenizer configuration
+- model_meta.yaml  — Metadata (catalog info + download timestamps)
+"""
+
+import hashlib
+import shutil
+import urllib.request
+import urllib.error
+import os
+import ssl
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable, Dict, List, Optional
+
+import yaml
+
+
+@dataclass
+class ModelDefinition:
+    """Definition of a model in the catalog."""
+
+    name: str
+    display_name: str
+    hf_repo: str
+    description: str
+    num_labels: int
+    label_map: Dict[int, str]
+    risk_map: Dict[str, str]  # label -> risk level (safe/suspicious/dangerous)
+    max_length: int = 512
+    license: str = "unknown"
+    size_mb: float = 0.0  # approximate download size
+    files: List[str] = field(default_factory=list)
+    hf_subfolder: str = ""  # subfolder in the HF repo (e.g., "onnx")
+    requires_auth: bool = False
+    default: bool = False
+
+    # Confidence thresholds for escalation
+    escalate_min_confidence: float = 0.1
+    escalate_max_confidence: float = 0.9
+
+
+# ============================================================================
+# MODEL CATALOG
+# ============================================================================
+
+MODEL_CATALOG: Dict[str, ModelDefinition] = {
+    "deberta-v3-injection": ModelDefinition(
+        name="deberta-v3-injection",
+        display_name="ProtectAI DeBERTa v3 Prompt Injection v2",
+        hf_repo="protectai/deberta-v3-base-prompt-injection-v2",
+        description=(
+            "Binary prompt injection classifier based on DeBERTa-v3-base. "
+            "Detects prompt injection attacks in English text. "
+            "Apache 2.0 license, no authentication required."
+        ),
+        num_labels=2,
+        label_map={0: "safe", 1: "injection"},
+        risk_map={
+            "safe": "safe",
+            "injection": "dangerous",
+        },
+        max_length=512,
+        license="Apache-2.0",
+        size_mb=750.0,
+        files=["model.onnx", "tokenizer.json"],
+        hf_subfolder="onnx",
+        requires_auth=False,
+        default=True,
+        escalate_min_confidence=0.1,
+        escalate_max_confidence=0.9,
+    ),
+}
+
+DEFAULT_MODEL = "deberta-v3-injection"
+
+
+# ============================================================================
+# DIRECTORY MANAGEMENT
+# ============================================================================
+
+
+def get_models_dir() -> Path:
+    """Get the models directory (~/.tweek/models/)."""
+    models_dir = Path.home() / ".tweek" / "models"
+    return models_dir
+
+
+def get_model_dir(name: str) -> Path:
+    """Get the directory for a specific model."""
+    return get_models_dir() / name
+
+
+def get_default_model_name() -> str:
+    """Get the configured default model name.
+
+    Checks user config first, falls back to catalog default.
+    """
+    config_path = Path.home() / ".tweek" / "config.yaml"
+    if config_path.exists():
+        try:
+            with open(config_path) as f:
+                config = yaml.safe_load(f) or {}
+            local_model_cfg = config.get("local_model", {})
+            model = local_model_cfg.get("model", "auto")
+            if model != "auto" and model in MODEL_CATALOG:
+                return model
+        except Exception:
+            pass
+
+    return DEFAULT_MODEL
+
+
+def is_model_installed(name: str) -> bool:
+    """Check if a model is installed with all required files."""
+    if name not in MODEL_CATALOG:
+        return False
+
+    model_dir = get_model_dir(name)
+    if not model_dir.exists():
+        return False
+
+    definition = MODEL_CATALOG[name]
+    for filename in definition.files:
+        if not (model_dir / filename).exists():
+            return False
+
+    return True
+
+
+def list_installed_models() -> List[str]:
+    """List all installed model names."""
+    models_dir = get_models_dir()
+    if not models_dir.exists():
+        return []
+
+    installed = []
+    for name in MODEL_CATALOG:
+        if is_model_installed(name):
+            installed.append(name)
+
+    return installed
+
+
+def get_model_definition(name: str) -> Optional[ModelDefinition]:
+    """Get the catalog definition for a model."""
+    return MODEL_CATALOG.get(name)
+
+
+# ============================================================================
+# MODEL DOWNLOAD
+# ============================================================================
+
+
+class ModelDownloadError(Exception):
+    """Error during model download."""
+
+    pass
+
+
+def _build_hf_url(repo: str, filename: str, subfolder: str = "") -> str:
+    """Build a HuggingFace CDN download URL."""
+    if subfolder:
+        return f"https://huggingface.co/{repo}/resolve/main/{subfolder}/{filename}"
+    return f"https://huggingface.co/{repo}/resolve/main/{filename}"
+
+
+def _get_hf_headers() -> Dict[str, str]:
+    """Get HTTP headers for HuggingFace requests."""
+    headers = {
+        "User-Agent": "tweek/0.1.0",
+    }
+
+    # Support HF_TOKEN for gated models (like Prompt Guard)
+    hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+    if hf_token:
+        headers["Authorization"] = f"Bearer {hf_token}"
+
+    return headers
+
+
+def download_model(
+    name: str,
+    progress_callback: Optional[Callable[[str, int, int], None]] = None,
+    force: bool = False,
+) -> Path:
+    """Download a model from HuggingFace.
+
+    Args:
+        name: Model name from the catalog.
+        progress_callback: Optional callback(filename, bytes_downloaded, total_bytes).
+        force: If True, re-download even if already installed.
+
+    Returns:
+        Path to the model directory.
+
+    Raises:
+        ModelDownloadError: If the model is not in the catalog or download fails.
+    """
+    definition = MODEL_CATALOG.get(name)
+    if definition is None:
+        available = ", ".join(MODEL_CATALOG.keys())
+        raise ModelDownloadError(
+            f"Unknown model '{name}'. Available models: {available}"
+        )
+
+    model_dir = get_model_dir(name)
+
+    if is_model_installed(name) and not force:
+        return model_dir
+
+    # Create directory
+    model_dir.mkdir(parents=True, exist_ok=True)
+
+    headers = _get_hf_headers()
+
+    if definition.requires_auth and "Authorization" not in headers:
+        raise ModelDownloadError(
+            f"Model '{name}' requires HuggingFace authentication. "
+            f"Set HF_TOKEN environment variable with a token that has "
+            f"access to {definition.hf_repo}. "
+            f"Get a token at https://huggingface.co/settings/tokens"
+        )
+
+    # Create SSL context
+    ssl_context = ssl.create_default_context()
+
+    # Download each file
+    for filename in definition.files:
+        url = _build_hf_url(definition.hf_repo, filename, definition.hf_subfolder)
+        dest = model_dir / filename
+        tmp_dest = model_dir / f".{filename}.tmp"
+
+        try:
+            request = urllib.request.Request(url, headers=headers)
+            response = urllib.request.urlopen(request, context=ssl_context)
+
+            total = int(response.headers.get("Content-Length", 0))
+            downloaded = 0
+            chunk_size = 1024 * 1024  # 1MB chunks
+
+            with open(tmp_dest, "wb") as f:
+                while True:
+                    chunk = response.read(chunk_size)
+                    if not chunk:
+                        break
+                    f.write(chunk)
+                    downloaded += len(chunk)
+                    if progress_callback:
+                        progress_callback(filename, downloaded, total)
+
+            # Atomic rename
+            tmp_dest.rename(dest)
+
+        except urllib.error.HTTPError as e:
+            tmp_dest.unlink(missing_ok=True)
+            if e.code == 401:
+                raise ModelDownloadError(
+                    f"Authentication failed for '{name}'. "
+                    f"Check your HF_TOKEN has access to {definition.hf_repo}. "
+                    f"You may need to accept the license at "
+                    f"https://huggingface.co/{definition.hf_repo}"
+                ) from e
+            elif e.code == 404:
+                raise ModelDownloadError(
+                    f"File '{filename}' not found in {definition.hf_repo}. "
+                    f"The model may have been moved or renamed."
+                ) from e
+            else:
+                raise ModelDownloadError(
+                    f"HTTP {e.code} downloading {filename}: {e.reason}"
+                ) from e
+        except urllib.error.URLError as e:
+            tmp_dest.unlink(missing_ok=True)
+            raise ModelDownloadError(
+                f"Network error downloading {filename}: {e.reason}"
+            ) from e
+        except Exception as e:
+            tmp_dest.unlink(missing_ok=True)
+            raise ModelDownloadError(
+                f"Failed to download {filename}: {e}"
+            ) from e
+
+    # Write metadata
+    meta = {
+        "name": definition.name,
+        "display_name": definition.display_name,
+        "hf_repo": definition.hf_repo,
+        "num_labels": definition.num_labels,
+        "label_map": definition.label_map,
+        "risk_map": definition.risk_map,
+        "max_length": definition.max_length,
+        "license": definition.license,
+        "downloaded_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+        "files": definition.files,
+    }
+
+    with open(model_dir / "model_meta.yaml", "w") as f:
+        yaml.dump(meta, f, default_flow_style=False, sort_keys=False)
+
+    return model_dir
+
+
+def remove_model(name: str) -> bool:
+    """Remove a downloaded model.
+
+    Args:
+        name: Model name.
+
+    Returns:
+        True if the model was removed, False if not found.
+    """
+    model_dir = get_model_dir(name)
+    if model_dir.exists():
+        shutil.rmtree(model_dir)
+        return True
+    return False
+
+
+def verify_model(name: str) -> Dict[str, bool]:
+    """Verify a model installation.
+
+    Args:
+        name: Model name.
+
+    Returns:
+        Dict mapping filename to exists status.
+    """
+    definition = MODEL_CATALOG.get(name)
+    if definition is None:
+        return {}
+
+    model_dir = get_model_dir(name)
+    status = {}
+
+    for filename in definition.files:
+        status[filename] = (model_dir / filename).exists()
+
+    status["model_meta.yaml"] = (model_dir / "model_meta.yaml").exists()
+
+    return status
+
+
+def get_model_size(name: str) -> Optional[int]:
+    """Get the total size of an installed model in bytes.
+
+    Args:
+        name: Model name.
+
+    Returns:
+        Total size in bytes, or None if not installed.
+    """
+    model_dir = get_model_dir(name)
+    if not model_dir.exists():
+        return None
+
+    total = 0
+    for path in model_dir.iterdir():
+        if path.is_file():
+            total += path.stat().st_size
+
+    return total

tweek/security/rate_limiter.py

@@ -120,7 +120,7 @@ class CircuitBreaker:
     - OPEN: Too many failures, requests blocked, waiting for timeout
     - HALF_OPEN: Testing recovery, limited requests allowed
 
-    Based on moltbot's circuit breaker implementation for resilience.
+    Based on OpenClaw's circuit breaker implementation for resilience.
     """
 
     def __init__(self, config: Optional[CircuitBreakerConfig] = None):
@@ -459,8 +459,12 @@ class RateLimiter:
             RateLimitResult with allowed status and any violations
         """
         if not session_id:
-            # No session tracking - allow but log
-            return RateLimitResult(allowed=True, message="No session ID for rate limiting")
+            # No session ID - generate unique one per process invocation
+            import os as _os
+            import uuid as _uuid
+            session_id = hashlib.sha256(
+                f"tweek-{_os.getpid()}-{_os.getcwd()}-{_uuid.getnode()}".encode()
+            ).hexdigest()[:16]
 
         # Check circuit breaker first
         circuit_key = f"session:{session_id}"
@@ -531,10 +535,11 @@ class RateLimiter:
                         details["velocity_ratio"] = round(current / baseline, 2)
 
         except Exception as e:
-            # Database error - fail open but log
+            # Database error - fail closed for safety
             return RateLimitResult(
-                allowed=True,
-                message=f"Rate limit check failed: {e}",
+                allowed=False,
+                violations=[RateLimitViolation.BURST],
+                message=f"Rate limit check failed (blocking for safety): {e}",
                 details={"error": str(e)}
             )
 

tweek/security/secret_scanner.py

@@ -5,7 +5,7 @@ Tweek Secret Scanner
 Scans configuration files for hardcoded secrets and credentials.
 Enforces environment-variable-only secrets policy.
 
-Based on moltbot's secret-guard security hardening initiative.
+Based on OpenClaw's secret-guard security hardening initiative.
 """
 
 import os
@@ -150,6 +150,68 @@ class SecretScanner:
         # Slack tokens
         (r'xox[baprs]-[0-9]{10,13}-[0-9]{10,13}[a-zA-Z0-9-]*', SecretType.TOKEN, "critical"),
 
+        # Stripe keys (live and test)
+        (r'sk_live_[A-Za-z0-9]{24,}', SecretType.API_KEY, "critical"),
+        (r'pk_live_[A-Za-z0-9]{24,}', SecretType.API_KEY, "high"),
+        (r'rk_live_[A-Za-z0-9]{24,}', SecretType.API_KEY, "critical"),
+
+        # SendGrid API key
+        (r'SG\.[A-Za-z0-9_-]{22}\.[A-Za-z0-9_-]{43}', SecretType.API_KEY, "critical"),
+
+        # Anthropic API key
+        (r'sk-ant-[A-Za-z0-9_-]{36,}', SecretType.API_KEY, "critical"),
+
+        # OpenAI API key
+        (r'sk-[A-Za-z0-9]{20}T3BlbkFJ[A-Za-z0-9]{20}', SecretType.API_KEY, "critical"),
+
+        # npm token
+        (r'npm_[A-Za-z0-9]{36,}', SecretType.TOKEN, "critical"),
+
+        # PyPI token
+        (r'pypi-[A-Za-z0-9_-]{50,}', SecretType.TOKEN, "critical"),
+
+        # Azure subscription key / connection string
+        (r'(?i)DefaultEndpointProtocol=https;AccountName=[^;]+;AccountKey=[A-Za-z0-9+/=]{60,}', SecretType.CONNECTION_STRING, "critical"),
+
+        # Google API key
+        (r'AIza[A-Za-z0-9_-]{35}', SecretType.API_KEY, "high"),
+
+        # Google OAuth client secret
+        (r'GOCSPX-[A-Za-z0-9_-]{28}', SecretType.OAUTH_SECRET, "critical"),
+
+        # Twilio API key
+        (r'SK[a-f0-9]{32}', SecretType.API_KEY, "high"),
+
+        # Mailchimp API key
+        (r'[a-f0-9]{32}-us[0-9]{1,2}', SecretType.API_KEY, "high"),
+
+        # Discord bot token
+        (r'[MN][A-Za-z0-9]{23,}\.[A-Za-z0-9_-]{6}\.[A-Za-z0-9_-]{27,}', SecretType.TOKEN, "critical"),
+
+        # Heroku API key
+        (r'[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}', SecretType.API_KEY, "medium"),
+
+        # Vercel token
+        (r'vercel_[A-Za-z0-9_-]{24,}', SecretType.TOKEN, "critical"),
+
+        # Supabase key
+        (r'sbp_[A-Za-z0-9]{40,}', SecretType.API_KEY, "critical"),
+
+        # Databricks token
+        (r'dapi[a-f0-9]{32}', SecretType.TOKEN, "critical"),
+
+        # Hashicorp Vault token
+        (r'hvs\.[A-Za-z0-9_-]{24,}', SecretType.TOKEN, "critical"),
+
+        # GitLab tokens (personal, pipeline, runner)
+        (r'glpat-[A-Za-z0-9_-]{20,}', SecretType.TOKEN, "critical"),
+
+        # Figma token
+        (r'figd_[A-Za-z0-9_-]{40,}', SecretType.TOKEN, "high"),
+
+        # Linear API key
+        (r'lin_api_[A-Za-z0-9]{40,}', SecretType.API_KEY, "high"),
+
         # Generic high-entropy strings that look like secrets
         (r'[\'"][A-Za-z0-9+/]{40,}={0,2}[\'"]', SecretType.API_KEY, "medium"),
     ]
@@ -347,15 +409,19 @@ class SecretScanner:
         return False
 
     def _redact_value(self, value: Any) -> str:
-        """Redact a secret value for safe display."""
+        """Redact a secret value for safe display.
+
+        Only shows first 4 characters (provider prefix) — never reveals suffix
+        to minimize information leakage.
+        """
         if not isinstance(value, str):
             return "<redacted>"
 
         if len(value) <= 8:
             return "*" * len(value)
 
-        # Show first 4 and last 4 chars
-        return f"{value[:4]}{'*' * (len(value) - 8)}{value[-4:]}"
+        # Show first 4 chars only (identifies provider without revealing key material)
+        return f"{value[:4]}{'*' * (len(value) - 4)}"
 
     def _redact_line(self, line: str) -> str:
         """Redact sensitive parts of a line."""

tweek/security/session_analyzer.py

@@ -475,6 +475,28 @@ class SessionAnalyzer:
                     anomalies.append(AnomalyType.CAPABILITY_AGGREGATION)
                 all_details["capability_aggregation"] = aggregation_details
 
+                # Memory: read cross-session workflow baselines
+                try:
+                    from tweek.memory.queries import memory_get_workflow_baseline
+                    from tweek.memory.store import hash_project
+                    # Use session_id as a proxy for project context
+                    baseline = memory_get_workflow_baseline(session_id)
+                    if baseline:
+                        all_details["memory_baseline"] = baseline
+                        # Flag if current denial ratio significantly exceeds baseline
+                        current_denial_ratio = len(
+                            [e for e in events if e.get("decision") in ("block", "ask")]
+                        ) / max(len(events), 1)
+                        baseline_ratio = baseline.get("denial_ratio", 0)
+                        if (
+                            current_denial_ratio > baseline_ratio * 2
+                            and current_denial_ratio > 0.2
+                            and baseline.get("total_invocations", 0) >= 20
+                        ):
+                            all_details["memory_baseline_exceeded"] = True
+                except Exception:
+                    pass  # Memory is best-effort
+
                 # Calculate risk score
                 risk_score = self._calculate_risk_score(anomalies, events, all_details)
 
@@ -523,10 +545,12 @@ class SessionAnalyzer:
                 )
 
         except Exception as e:
+            # Fail closed: analysis failure is itself suspicious
             return SessionAnalysis(
                 session_id=session_id,
-                risk_score=0.0,
-                details={"error": str(e)}
+                risk_score=0.5,
+                anomalies=[AnomalyType.SUSPICIOUS_PATTERN],
+                details={"error": str(e), "fail_closed": True}
             )
 
     def _update_session_profile(