osintengine 1.0.2__py3-none-any.whl

src/watson/infra/__init__.py

@@ -0,0 +1,43 @@
+"""Infrastructure utilities — caching, retry, circuit breakers, rate limiting, resilience."""
+
+from .retry import (
+    retry,
+    CircuitBreaker,
+    CircuitOpenError,
+    get_circuit,
+    _is_retryable,
+)
+from .ratelimit import (
+    TokenBucket,
+    get_bucket,
+    RateLimitExceeded,
+)
+from .cache import (
+    TTLCache,
+    get_cache,
+)
+from .resilience import (
+    FailureReason,
+    ToolResult,
+    safe_execute,
+    classify_failure,
+    suggest_alternatives,
+)
+
+__all__ = [
+    "retry",
+    "CircuitBreaker",
+    "CircuitOpenError",
+    "get_circuit",
+    "_is_retryable",
+    "TokenBucket",
+    "get_bucket",
+    "RateLimitExceeded",
+    "TTLCache",
+    "get_cache",
+    "FailureReason",
+    "ToolResult",
+    "safe_execute",
+    "classify_failure",
+    "suggest_alternatives",
+]

src/watson/infra/cache.py

@@ -0,0 +1,171 @@
+"""TTL cache with size-bounded LRU eviction and stats.
+
+Also exposes a small registry/helpers for aggregating cache statistics.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+from typing import Any, Dict, Optional
+
+
+class _Entry:
+    __slots__ = ("value", "expiry", "hits", "inserted_at")
+
+    def __init__(self, value: Any, expiry: float, now: float):
+        self.value = value
+        self.expiry = expiry
+        self.hits = 0
+        self.inserted_at = now
+
+
+class TTLCache:
+    """In-memory TTL cache with max-size eviction.
+
+    Note: :meth:`set` takes ``(value, key)`` — value first, key second —
+    matching the project's OSINT-collection ergonomics.
+    """
+
+    def __init__(self, name: str, ttl: float = 60.0, max_size: int = 1000):
+        self.name = name
+        self.ttl = float(ttl)
+        self.max_size = int(max_size)
+        self._data: Dict[str, _Entry] = {}
+        self._lock = threading.Lock()
+        self.stats: Dict[str, Any] = {
+            "hits": 0,
+            "misses": 0,
+            "evictions": 0,
+            "size": 0,
+        }
+
+    # -- core operations ------------------------------------------
+    def set(self, value: Any, key: str) -> None:
+        """Store *value* under *key* with the configured TTL."""
+        now = time.monotonic()
+        with self._lock:
+            if key in self._data:
+                # Update in place, keep hit count
+                entry = self._data[key]
+                entry.value = value
+                entry.expiry = now + self.ttl
+            else:
+                self._data[key] = _Entry(value, now + self.ttl, now)
+                if len(self._data) > self.max_size:
+                    self._evict()
+            self.stats["size"] = len(self._data)
+
+    def get(self, key: str) -> Any:
+        """Return the value for *key* or ``None`` on miss/expiry."""
+        now = time.monotonic()
+        with self._lock:
+            entry = self._data.get(key)
+            if entry is None:
+                self.stats["misses"] += 1
+                return None
+            if now >= entry.expiry:
+                # Expired — remove and count as miss
+                del self._data[key]
+                self.stats["misses"] += 1
+                self.stats["size"] = len(self._data)
+                return None
+            entry.hits += 1
+            self.stats["hits"] += 1
+            return entry.value
+
+    def invalidate(self, key: str) -> None:
+        """Remove *key* from the cache if present."""
+        with self._lock:
+            if key in self._data:
+                del self._data[key]
+                self.stats["size"] = len(self._data)
+
+    def clear(self) -> None:
+        with self._lock:
+            self._data.clear()
+            self.stats["size"] = 0
+
+    # -- eviction --------------------------------------------------
+    def _evict(self) -> None:
+        """Evict the least-hit (oldest on tie) entry. Caller holds lock."""
+        if not self._data:
+            return
+        # Find entry with fewest hits; ties broken by insertion order (oldest first)
+        victim_key = None
+        victim_hits = None
+        for k, entry in self._data.items():
+            if victim_key is None or entry.hits < victim_hits:
+                victim_key = k
+                victim_hits = entry.hits
+        if victim_key is not None:
+            del self._data[victim_key]
+            self.stats["evictions"] += 1
+
+    # -- derived metrics ------------------------------------------
+    @property
+    def hit_rate(self) -> float:
+        total = self.stats["hits"] + self.stats["misses"]
+        if total == 0:
+            return 0.0
+        return self.stats["hits"] / total
+
+
+# ─────────────────────────────────────────────────────────────────
+# Pre-configured caches
+# ─────────────────────────────────────────────────────────────────
+
+_PRECONFIGURED: Dict[str, Dict[str, Any]] = {
+    "dns": {"ttl": 300, "max_size": 500},
+    "whois": {"ttl": 3600, "max_size": 500},
+    "cert": {"ttl": 600, "max_size": 1000},
+    "default": {"ttl": 60, "max_size": 1000},
+}
+
+_caches: Dict[str, TTLCache] = {}
+_caches_lock = threading.Lock()
+
+
+def get_cache(name: str) -> TTLCache:
+    """Get-or-create a singleton :class:`TTLCache` for *name*."""
+    with _caches_lock:
+        c = _caches.get(name)
+        if c is None:
+            cfg = _PRECONFIGURED.get(name, _PRECONFIGURED["default"])
+            c = TTLCache(
+                name=name,
+                ttl=cfg.get("ttl", 60),
+                max_size=cfg.get("max_size", 1000),
+            )
+            _caches[name] = c
+        return c
+
+
+# ─────────────────────────────────────────────────────────────────
+# Stats aggregation (legacy helpers)
+# ─────────────────────────────────────────────────────────────────
+
+_cache_registry: Dict[str, Any] = {}
+
+
+def register_cache(name: str, cache_obj: Any) -> None:
+    """Register a cache for stats collection."""
+    _cache_registry[name] = cache_obj
+
+
+def all_cache_stats() -> Dict[str, Dict[str, Any]]:
+    """Return stats for all registered caches."""
+    stats: Dict[str, Dict[str, Any]] = {}
+    for name, cache in _cache_registry.items():
+        s = getattr(cache, "stats", None)
+        if isinstance(s, dict):
+            stats[name] = dict(s)
+        elif callable(s):
+            stats[name] = s()
+        elif hasattr(cache, "get_stats"):
+            stats[name] = cache.get_stats()
+        else:
+            stats[name] = {"size": 0, "hits": 0, "misses": 0}
+    if not stats:
+        stats["default"] = {"size": 0, "hits": 0, "misses": 0}
+    return stats

src/watson/infra/ratelimit.py

@@ -0,0 +1,125 @@
+"""Token-bucket rate limiting."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+import time
+from typing import Any, Dict, Optional
+
+
+class RateLimitExceeded(Exception):
+    """Raised when a rate-limited operation cannot acquire a token in time."""
+
+    def __init__(self, name: str = "", message: str = ""):
+        self.name = name
+        super().__init__(message or f"rate limit exceeded for '{name}'")
+
+
+class TokenBucket:
+    """Classic token-bucket rate limiter.
+
+    Tokens accumulate at *refill_rate* per second up to *capacity*. Each
+    :meth:`acquire` consumes one token. ``acquire`` blocks up to *timeout*
+    seconds waiting for a token to become available.
+    """
+
+    def __init__(
+        self,
+        name: str = "",
+        capacity: float = 5,
+        refill_rate: float = 1.0,
+    ):
+        self.name = name
+        self.capacity = float(capacity)
+        self.refill_rate = float(refill_rate)
+        self._tokens = float(capacity)
+        self._last_refill = time.monotonic()
+        self._lock = threading.Lock()
+
+    # -- internals -------------------------------------------------
+    def _refill(self) -> None:
+        """Add tokens based on elapsed wall-clock time. Caller holds lock."""
+        now = time.monotonic()
+        elapsed = now - self._last_refill
+        if elapsed > 0 and self.refill_rate > 0:
+            self._tokens = min(self.capacity, self._tokens + elapsed * self.refill_rate)
+        self._last_refill = now
+
+    def _try_consume(self) -> bool:
+        """Non-blocking consume. Caller must hold lock."""
+        self._refill()
+        if self._tokens >= 1.0:
+            self._tokens -= 1.0
+            return True
+        return False
+
+    # -- public API ------------------------------------------------
+    def acquire(self, timeout: float = 0.0) -> bool:
+        """Block up to *timeout* seconds for a token. Return True if acquired.
+
+        A *timeout* of 0 means non-blocking: try once and return immediately.
+        """
+        deadline = time.monotonic() + timeout
+        while True:
+            with self._lock:
+                if self._try_consume():
+                    return True
+            if time.monotonic() >= deadline:
+                return False
+            # Sleep a short slice before retrying so we don't busy-spin hard.
+            remaining = deadline - time.monotonic()
+            time.sleep(max(0.0, min(0.005, remaining)))
+
+    async def async_acquire(self, timeout: float = 0.0) -> bool:
+        """Async variant of :meth:`acquire`."""
+        deadline = time.monotonic() + timeout
+        while True:
+            with self._lock:
+                if self._try_consume():
+                    return True
+            if time.monotonic() >= deadline:
+                return False
+            remaining = deadline - time.monotonic()
+            await asyncio.sleep(max(0.0, min(0.005, remaining)))
+
+    # -- context manager ------------------------------------------
+    def __enter__(self) -> "TokenBucket":
+        if not self.acquire(timeout=0.0):
+            raise RateLimitExceeded(self.name)
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> bool:
+        return False  # never suppress
+
+
+# ─────────────────────────────────────────────────────────────────
+# Pre-configured buckets
+# ─────────────────────────────────────────────────────────────────
+
+_PRECONFIGURED: Dict[str, Dict[str, Any]] = {
+    "crt.sh": {"capacity": 3, "refill_rate": 1.0},
+    "virustotal": {"capacity": 4, "refill_rate": 1.0},
+    "shodan": {"capacity": 1, "refill_rate": 1.0},
+}
+
+_DEFAULT_CAPACITY = 5
+_DEFAULT_REFILL_RATE = 1.0
+
+_buckets: Dict[str, TokenBucket] = {}
+_buckets_lock = threading.Lock()
+
+
+def get_bucket(name: str) -> TokenBucket:
+    """Get-or-create a singleton :class:`TokenBucket` for *name*."""
+    with _buckets_lock:
+        b = _buckets.get(name)
+        if b is None:
+            cfg = _PRECONFIGURED.get(name, {})
+            b = TokenBucket(
+                name=name,
+                capacity=cfg.get("capacity", _DEFAULT_CAPACITY),
+                refill_rate=cfg.get("refill_rate", _DEFAULT_REFILL_RATE),
+            )
+            _buckets[name] = b
+        return b

src/watson/infra/resilience.py

@@ -0,0 +1,239 @@
+"""Resilience layer — failure classification, safe execution, alternatives.
+
+Wraps tool invocations so the orchestrator gets a uniform :class:`ToolResult`
+regardless of whether a tool succeeded, returned nothing, or raised.
+"""
+
+from __future__ import annotations
+
+import enum
+import time
+from typing import Any, Callable, List, Optional
+
+
+class FailureReason(enum.Enum):
+    """Why a tool invocation failed (or yielded no usable data)."""
+
+    NO_RESULTS = "no_results"
+    TIMEOUT = "timeout"
+    CLOUDFLARE_PROTECTED = "cloudflare_protected"
+    RATE_LIMITED = "rate_limited"
+    WHOIS_REDACTED = "whois_redacted"
+    BLOCKED = "blocked"
+    UNKNOWN = "unknown"
+
+
+# Reasons that indicate the target is actively blocking us — these are
+# intelligence signals, not mere failures.
+_BLOCK_REASONS = {
+    FailureReason.CLOUDFLARE_PROTECTED,
+    FailureReason.RATE_LIMITED,
+    FailureReason.WHOIS_REDACTED,
+    FailureReason.BLOCKED,
+}
+
+
+class ToolResult:
+    """Uniform result wrapper for tool invocations."""
+
+    __slots__ = (
+        "tool",
+        "success",
+        "data",
+        "failure_reason",
+        "latency_ms",
+        "retry_count",
+        "context",
+    )
+
+    def __init__(
+        self,
+        success: bool,
+        data: Any = None,
+        failure_reason: Optional[FailureReason] = None,
+        latency_ms: float = 0.0,
+        retry_count: int = 0,
+        tool: str = "",
+        context: Any = None,
+    ):
+        self.tool = tool
+        self.success = success
+        self.data = data
+        self.failure_reason = failure_reason
+        self.latency_ms = latency_ms
+        self.retry_count = retry_count
+        self.context = context
+
+    # -- derived properties ---------------------------------------
+    @property
+    def blocked(self) -> bool:
+        """True when the failure reason indicates active blocking."""
+        return self.failure_reason in _BLOCK_REASONS
+
+    @property
+    def is_intelligence(self) -> bool:
+        """True when the outcome itself is an intelligence signal."""
+        return self.failure_reason in _BLOCK_REASONS
+
+    # -- serialization --------------------------------------------
+    def to_dict(self) -> dict:
+        return {
+            "tool": self.tool,
+            "success": self.success,
+            "data": self.data,
+            "failure_reason": self.failure_reason.value if self.failure_reason else None,
+            "latency_ms": self.latency_ms,
+            "retry_count": self.retry_count,
+            "context": self.context,
+        }
+
+    def __repr__(self) -> str:  # pragma: no cover
+        return (
+            f"ToolResult(success={self.success}, tool={self.tool!r}, "
+            f"failure_reason={self.failure_reason}, latency_ms={self.latency_ms})"
+        )
+
+
+# ─────────────────────────────────────────────────────────────────
+# Failure classification
+# ─────────────────────────────────────────────────────────────────
+
+
+def classify_failure(exc: BaseException) -> FailureReason:
+    """Classify an exception into a :class:`FailureReason`."""
+    # 1. Timeout
+    if isinstance(exc, TimeoutError):
+        return FailureReason.TIMEOUT
+
+    # 2. HTTP-style rate limiting (objects exposing a status code)
+    code = getattr(exc, "code", None)
+    if code is None:
+        code = getattr(exc, "status_code", None)
+    if code == 429:
+        return FailureReason.RATE_LIMITED
+
+    msg = str(exc).lower()
+
+    # 3. Cloudflare / WAF protection
+    if "cloudflare" in msg or "cf-ray" in msg or "attention required" in msg:
+        return FailureReason.CLOUDFLARE_PROTECTED
+
+    # 4. WHOIS privacy / GDPR redaction
+    if "redacted" in msg or "gdpr" in msg or "private registration" in msg:
+        return FailureReason.WHOIS_REDACTED
+
+    # 5. Generic network errors → could be transient; classify as unknown
+    return FailureReason.UNKNOWN
+
+
+# ─────────────────────────────────────────────────────────────────
+# Safe execution
+# ─────────────────────────────────────────────────────────────────
+
+
+def safe_execute(
+    tool: str,
+    func: Callable[[], Any],
+    context: Any = None,
+    retry_count: int = 0,
+) -> ToolResult:
+    """Run *func* and return a :class:`ToolResult` capturing the outcome.
+
+    A return value of ``None`` or an empty collection is treated as
+    ``NO_RESULTS`` rather than success.
+    """
+    start = time.monotonic()
+    try:
+        data = func()
+    except BaseException as exc:  # noqa: BLE001
+        latency_ms = (time.monotonic() - start) * 1000.0
+        reason = classify_failure(exc)
+        return ToolResult(
+            success=False,
+            failure_reason=reason,
+            latency_ms=latency_ms,
+            retry_count=retry_count,
+            tool=tool,
+            context=context,
+        )
+
+    latency_ms = (time.monotonic() - start) * 1000.0
+
+    # Treat None / empty containers as "no results"
+    if data is None:
+        return ToolResult(
+            success=False,
+            failure_reason=FailureReason.NO_RESULTS,
+            latency_ms=latency_ms,
+            retry_count=retry_count,
+            tool=tool,
+            context=context,
+        )
+    if isinstance(data, (list, tuple, set, dict, str)) and len(data) == 0:
+        return ToolResult(
+            success=False,
+            failure_reason=FailureReason.NO_RESULTS,
+            latency_ms=latency_ms,
+            retry_count=retry_count,
+            tool=tool,
+            context=context,
+        )
+
+    return ToolResult(
+        success=True,
+        data=data,
+        latency_ms=latency_ms,
+        retry_count=retry_count,
+        tool=tool,
+        context=context,
+    )
+
+
+# ─────────────────────────────────────────────────────────────────
+# Alternative-source suggestions
+# ─────────────────────────────────────────────────────────────────
+
+_ALTERNATIVES: dict = {
+    FailureReason.WHOIS_REDACTED: [
+        "crt.sh certificate transparency search",
+        "DNS historical records (SecurityTrails)",
+        "Search engine cached WHOIS",
+        "Passive DNS replication",
+    ],
+    FailureReason.CLOUDFLARE_PROTECTED: [
+        "Historical DNS records (SecurityTrails)",
+        "crt.sh certificate transparency search",
+        "Search engine cache / archive.org",
+        "Shodan / Censys internet scan",
+    ],
+    FailureReason.RATE_LIMITED: [
+        "Wait and retry with backoff",
+        "Use an alternate data source",
+        "Reduce request frequency",
+    ],
+    FailureReason.TIMEOUT: [
+        "Retry with longer timeout",
+        "Use an alternate data source",
+        "crt.sh certificate transparency search",
+    ],
+    FailureReason.NO_RESULTS: [
+        "Try alternate query / domain variant",
+        "Use an alternate data source",
+        "crt.sh certificate transparency search",
+    ],
+    FailureReason.BLOCKED: [
+        "Historical DNS records",
+        "crt.sh certificate transparency search",
+        "archive.org",
+    ],
+    FailureReason.UNKNOWN: [
+        "Retry the operation",
+        "Use an alternate data source",
+        "crt.sh certificate transparency search",
+    ],
+}
+
+
+def suggest_alternatives(reason: FailureReason) -> List[str]:
+    """Return a list of suggested alternative approaches for *reason*."""
+    return list(_ALTERNATIVES.get(reason, _ALTERNATIVES[FailureReason.UNKNOWN]))