shieldops-sdk 0.1.0__py3-none-any.whl

shieldops_sdk/__init__.py

@@ -0,0 +1,50 @@
+"""ShieldOps SDK -- AI Security Control Plane SDK.
+
+Intercept and govern AI agent tool calls with one line of code.
+
+Quick start::
+
+    from shieldops_sdk import ShieldOpsClient, ShieldOpsInterceptor, ShieldOpsConfig
+
+    # API client for the ShieldOps platform
+    client = ShieldOpsClient(api_key="sk-...")
+
+    # Framework-agnostic tool call interceptor
+    config = ShieldOpsConfig(api_key="sk-...", mode="enforce")
+    interceptor = ShieldOpsInterceptor(config)
+    decision = interceptor.check("my_tool", {"arg": "value"})
+"""
+
+from __future__ import annotations
+
+from shieldops_sdk.async_client import AsyncShieldOpsClient
+from shieldops_sdk.client import ShieldOpsClient
+from shieldops_sdk.config import SDKMode, ShieldOpsConfig
+from shieldops_sdk.exceptions import (
+    AuthenticationError,
+    NotFoundError,
+    RateLimitError,
+    ShieldOpsConnectionError,
+    ShieldOpsDeniedError,
+    ShieldOpsError,
+    ValidationError,
+)
+from shieldops_sdk.interceptor import Decision, ShieldOpsInterceptor, ToolCall
+
+__all__ = [
+    "AsyncShieldOpsClient",
+    "AuthenticationError",
+    "Decision",
+    "NotFoundError",
+    "RateLimitError",
+    "SDKMode",
+    "ShieldOpsClient",
+    "ShieldOpsConfig",
+    "ShieldOpsConnectionError",
+    "ShieldOpsDeniedError",
+    "ShieldOpsError",
+    "ShieldOpsInterceptor",
+    "ToolCall",
+    "ValidationError",
+]
+__version__ = "0.1.0"

shieldops_sdk/_policy/__init__.py

@@ -0,0 +1,41 @@
+"""Private policy module — default pattern catalogues and merge helpers.
+
+The leading underscore makes this package private. It is not part of the
+public SDK API and may be reorganised between minor releases. Public clients
+should configure policy via ``ShieldOpsConfig.extra_blocked_patterns`` and
+``ShieldOpsConfig.extra_high_risk_patterns`` rather than importing from here.
+
+The helpers centralise the "effective policy set" computation so future PRs
+(mode/telemetry split, callbacks) can reuse one merge implementation rather
+than duplicating ``defaults | extras`` in multiple call sites.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from shieldops_sdk._policy._defaults import (
+    DEFAULT_BLOCKED_PATTERNS,
+    DEFAULT_HIGH_RISK_PATTERNS,
+)
+
+if TYPE_CHECKING:
+    from shieldops_sdk.config import ShieldOpsConfig
+
+
+def effective_blocked_patterns(config: ShieldOpsConfig) -> set[str]:
+    """Return a fresh set of blocked patterns: defaults ∪ config.extra_blocked_patterns."""
+    return set(DEFAULT_BLOCKED_PATTERNS) | set(config.extra_blocked_patterns)
+
+
+def effective_high_risk_patterns(config: ShieldOpsConfig) -> set[str]:
+    """Return a fresh set of high-risk patterns: defaults ∪ config.extra_high_risk_patterns."""
+    return set(DEFAULT_HIGH_RISK_PATTERNS) | set(config.extra_high_risk_patterns)
+
+
+__all__ = [
+    "DEFAULT_BLOCKED_PATTERNS",
+    "DEFAULT_HIGH_RISK_PATTERNS",
+    "effective_blocked_patterns",
+    "effective_high_risk_patterns",
+]

shieldops_sdk/_policy/_defaults.py

@@ -0,0 +1,33 @@
+"""Default pattern catalogues for the ShieldOps interceptor.
+
+These are ``frozenset`` so the module-level defaults cannot be mutated by
+client code. The interceptor copies them into per-instance ``set`` attributes
+at construction time, preserving the existing pattern of
+``interceptor._blocked_tools.update(...)`` for advanced post-construct
+configuration.
+"""
+
+from __future__ import annotations
+
+DEFAULT_BLOCKED_PATTERNS: frozenset[str] = frozenset(
+    {
+        "delete_database",
+        "drop_table",
+        "modify_iam_root",
+        "rm_rf",
+        "format_disk",
+        "disable_firewall",
+        "delete_backup",
+    }
+)
+
+DEFAULT_HIGH_RISK_PATTERNS: frozenset[str] = frozenset(
+    {
+        "execute_command",
+        "run_shell",
+        "modify_security_group",
+        "change_iam_policy",
+        "create_user",
+        "rotate_credentials",
+    }
+)

shieldops_sdk/_response.py

@@ -0,0 +1,72 @@
+"""Shared HTTP response handling for sync and async clients."""
+
+from __future__ import annotations
+
+import httpx
+
+from shieldops_sdk.exceptions import (
+    AuthenticationError,
+    NotFoundError,
+    RateLimitError,
+    ServerError,
+    ShieldOpsError,
+    ValidationError,
+)
+
+
+def handle_response(response: httpx.Response) -> None:
+    """Raise the appropriate SDK exception for non-2xx responses.
+
+    This is intentionally a plain function so both sync and async
+    resource classes can reuse it without duplication.
+    """
+    if response.is_success:
+        return
+
+    status = response.status_code
+
+    # Try to extract a detail message from the JSON body
+    detail = ""
+    try:
+        body = response.json()
+        detail = body.get("detail", "") if isinstance(body, dict) else ""
+    except Exception:  # noqa: BLE001
+        detail = response.text[:200] if response.text else ""
+
+    if status in (401, 403):
+        raise AuthenticationError(detail or "Authentication failed")
+
+    if status == 404:
+        raise NotFoundError(detail or "Resource not found")
+
+    if status == 422:
+        raise ValidationError(detail or "Validation error")
+
+    if status == 429:
+        retry_after: int | None = None
+        raw = response.headers.get("Retry-After")
+        if raw is not None:
+            try:
+                retry_after = int(raw)
+            except ValueError:
+                pass
+        # Also try the JSON body which ShieldOps returns
+        if retry_after is None:
+            try:
+                body = response.json()
+                retry_after = body.get("retry_after")
+            except Exception:  # noqa: BLE001
+                pass
+        raise RateLimitError(
+            detail or "Rate limit exceeded",
+            retry_after=retry_after,
+        )
+
+    if status >= 500:
+        raise ServerError(
+            detail or "Internal server error",
+            status_code=status,
+        )
+
+    # Catch-all for other 4xx
+    raise ShieldOpsError(detail or f"Request failed with status {status}", status_code=status)

shieldops_sdk/async_client.py

@@ -0,0 +1,83 @@
+"""Asynchronous ShieldOps API client."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import httpx
+
+from shieldops_sdk._response import handle_response
+from shieldops_sdk.resources.agents import AsyncAgentsResource
+from shieldops_sdk.resources.investigations import AsyncInvestigationsResource
+from shieldops_sdk.resources.remediations import AsyncRemediationsResource
+from shieldops_sdk.resources.security import AsyncSecurityResource
+from shieldops_sdk.resources.vulnerabilities import AsyncVulnerabilitiesResource
+
+_DEFAULT_BASE_URL = "http://localhost:8000/api/v1"
+_DEFAULT_TIMEOUT = 30.0
+
+
+class AsyncShieldOpsClient:
+    """Async client for the ShieldOps REST API.
+
+    Usage::
+
+        async with AsyncShieldOpsClient(api_key="sk-...") as client:
+            invs = await client.investigations.list(limit=10)
+            for inv in invs.items:
+                print(inv.investigation_id, inv.status)
+    """
+
+    def __init__(
+        self,
+        *,
+        base_url: str = _DEFAULT_BASE_URL,
+        api_key: str | None = None,
+        token: str | None = None,
+        timeout: float = _DEFAULT_TIMEOUT,
+    ) -> None:
+        headers: dict[str, str] = {"User-Agent": "shieldops-sdk/1.0.0"}
+        if api_key:
+            headers["X-API-Key"] = api_key
+        elif token:
+            headers["Authorization"] = f"Bearer {token}"
+
+        self._http = httpx.AsyncClient(
+            base_url=base_url,
+            headers=headers,
+            timeout=timeout,
+        )
+
+        # Resource namespaces
+        self.investigations = AsyncInvestigationsResource(self._http)
+        self.remediations = AsyncRemediationsResource(self._http)
+        self.security = AsyncSecurityResource(self._http)
+        self.vulnerabilities = AsyncVulnerabilitiesResource(self._http)
+        self.agents = AsyncAgentsResource(self._http)
+
+    # -- Async context manager --------------------------------------------
+
+    async def __aenter__(self) -> AsyncShieldOpsClient:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: Any,
+    ) -> None:
+        await self.close()
+
+    async def close(self) -> None:
+        """Close the underlying HTTP transport."""
+        await self._http.aclose()
+
+    # -- Convenience methods ----------------------------------------------
+
+    async def health(self) -> dict[str, Any]:
+        """Check the API health endpoint."""
+        base = str(self._http.base_url)
+        root = base.rsplit("/api/", 1)[0] if "/api/" in base else base
+        resp = await self._http.get(f"{root}/health")
+        handle_response(resp)
+        return resp.json()

shieldops_sdk/client.py

@@ -0,0 +1,85 @@
+"""Synchronous ShieldOps API client."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import httpx
+
+from shieldops_sdk._response import handle_response
+from shieldops_sdk.resources.agents import AgentsResource
+from shieldops_sdk.resources.investigations import InvestigationsResource
+from shieldops_sdk.resources.remediations import RemediationsResource
+from shieldops_sdk.resources.security import SecurityResource
+from shieldops_sdk.resources.vulnerabilities import VulnerabilitiesResource
+
+_DEFAULT_BASE_URL = "http://localhost:8000/api/v1"
+_DEFAULT_TIMEOUT = 30.0
+
+
+class ShieldOpsClient:
+    """Synchronous client for the ShieldOps REST API.
+
+    Usage::
+
+        with ShieldOpsClient(api_key="sk-...") as client:
+            invs = client.investigations.list(limit=10)
+            for inv in invs.items:
+                print(inv.investigation_id, inv.status)
+    """
+
+    def __init__(
+        self,
+        *,
+        base_url: str = _DEFAULT_BASE_URL,
+        api_key: str | None = None,
+        token: str | None = None,
+        timeout: float = _DEFAULT_TIMEOUT,
+    ) -> None:
+        headers: dict[str, str] = {"User-Agent": "shieldops-sdk/1.0.0"}
+        if api_key:
+            headers["X-API-Key"] = api_key
+        elif token:
+            headers["Authorization"] = f"Bearer {token}"
+
+        self._http = httpx.Client(
+            base_url=base_url,
+            headers=headers,
+            timeout=timeout,
+        )
+
+        # Resource namespaces
+        self.investigations = InvestigationsResource(self._http)
+        self.remediations = RemediationsResource(self._http)
+        self.security = SecurityResource(self._http)
+        self.vulnerabilities = VulnerabilitiesResource(self._http)
+        self.agents = AgentsResource(self._http)
+
+    # -- Context manager --------------------------------------------------
+
+    def __enter__(self) -> ShieldOpsClient:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: Any,
+    ) -> None:
+        self.close()
+
+    def close(self) -> None:
+        """Close the underlying HTTP transport."""
+        self._http.close()
+
+    # -- Convenience methods ----------------------------------------------
+
+    def health(self) -> dict[str, Any]:
+        """Check the API health endpoint."""
+        # Health lives outside the /api/v1 prefix, so use an absolute URL.
+        base = str(self._http.base_url)
+        # Strip the API version path to reach the root.
+        root = base.rsplit("/api/", 1)[0] if "/api/" in base else base
+        resp = self._http.get(f"{root}/health")
+        handle_response(resp)
+        return resp.json()

shieldops_sdk/config.py

@@ -0,0 +1,95 @@
+"""ShieldOps SDK configuration — loaded from constructor args or environment variables."""
+
+from __future__ import annotations
+
+import os
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+
+class SDKMode(str, Enum):
+    """SDK enforcement mode."""
+
+    AUDIT = "audit"
+    ENFORCE = "enforce"
+
+
+class SDKTelemetry(str, Enum):
+    """SDK telemetry destination.
+
+    Separate from :class:`SDKMode` — mode controls block-vs-audit on policy
+    violations; telemetry controls *where* records of those decisions go.
+
+    - ``LOCAL``: keep all events in-process, no network at all.
+    - ``REMOTE``: send to the ShieldOps backend (requires ``api_key``).
+    - ``OTLP``: send to the user's OpenTelemetry collector.
+    """
+
+    LOCAL = "local"
+    REMOTE = "remote"
+    OTLP = "otlp"
+
+
+class ShieldOpsConfig(BaseModel):
+    """Configuration for the ShieldOps SDK.
+
+    Values can be provided directly or read from environment variables:
+    - ``SHIELDOPS_API_KEY``
+    - ``SHIELDOPS_ENDPOINT``
+    - ``SHIELDOPS_MODE``
+
+    Attributes:
+        api_key: ShieldOps API key for authentication.
+        endpoint: ShieldOps API base URL.
+        mode: Operating mode -- ``audit`` logs without blocking, ``enforce`` blocks risky calls.
+        timeout: HTTP request timeout in seconds.
+    """
+
+    api_key: str = Field(default="")
+    endpoint: str = Field(default="https://api.shieldops.io")
+    mode: SDKMode = SDKMode.AUDIT
+    telemetry: SDKTelemetry = Field(
+        default=SDKTelemetry.LOCAL,
+        description=(
+            "Where intercept records are sent. LOCAL = in-process only "
+            "(default, no network). REMOTE = POST to ShieldOps backend "
+            "(requires api_key). OTLP = export via OpenTelemetry collector."
+        ),
+    )
+    timeout: float = Field(default=5.0, ge=0.1)
+    extra_blocked_patterns: set[str] = Field(
+        default_factory=set,
+        description=(
+            "Additional tool-name patterns to deny on top of the SDK defaults. "
+            "Merged with shieldops_sdk._policy defaults at interceptor construction."
+        ),
+    )
+    extra_high_risk_patterns: set[str] = Field(
+        default_factory=set,
+        description=(
+            "Additional tool-name patterns to flag as high-risk on top of the "
+            "SDK defaults. Merged with shieldops_sdk._policy defaults at "
+            "interceptor construction."
+        ),
+    )
+
+    def model_post_init(self, __context: object) -> None:
+        """Populate unset fields from environment variables."""
+        if not self.api_key:
+            self.api_key = os.environ.get("SHIELDOPS_API_KEY", "")
+        if self.endpoint == "https://api.shieldops.io":
+            env_endpoint = os.environ.get("SHIELDOPS_ENDPOINT", "")
+            if env_endpoint:
+                self.endpoint = env_endpoint
+        env_mode = os.environ.get("SHIELDOPS_MODE", "")
+        if env_mode and env_mode.lower() in ("audit", "enforce"):
+            self.mode = SDKMode(env_mode.lower())
+
+    @property
+    def is_enforce(self) -> bool:
+        return self.mode == SDKMode.ENFORCE
+
+    @property
+    def is_audit(self) -> bool:
+        return self.mode == SDKMode.AUDIT

shieldops_sdk/exceptions.py

@@ -0,0 +1,85 @@
+"""ShieldOps SDK exceptions."""
+
+from __future__ import annotations
+
+
+class ShieldOpsError(Exception):
+    """Base exception for all SDK errors."""
+
+    def __init__(self, message: str, status_code: int | None = None) -> None:
+        super().__init__(message)
+        self.message = message
+        self.status_code = status_code
+
+
+class AuthenticationError(ShieldOpsError):
+    """Invalid or expired API credentials (401/403)."""
+
+    def __init__(self, message: str = "Authentication failed") -> None:
+        super().__init__(message, status_code=401)
+
+
+class NotFoundError(ShieldOpsError):
+    """Resource not found (404)."""
+
+    def __init__(self, message: str = "Resource not found") -> None:
+        super().__init__(message, status_code=404)
+
+
+class RateLimitError(ShieldOpsError):
+    """Rate limit exceeded (429)."""
+
+    def __init__(
+        self,
+        message: str = "Rate limit exceeded",
+        retry_after: int | None = None,
+    ) -> None:
+        super().__init__(message, status_code=429)
+        self.retry_after = retry_after
+
+
+class ValidationError(ShieldOpsError):
+    """Request validation failed (422)."""
+
+    def __init__(self, message: str = "Validation error") -> None:
+        super().__init__(message, status_code=422)
+
+
+class ServerError(ShieldOpsError):
+    """Server-side error (5xx)."""
+
+    def __init__(
+        self,
+        message: str = "Internal server error",
+        status_code: int = 500,
+    ) -> None:
+        super().__init__(message, status_code=status_code)
+
+
+class ShieldOpsDeniedError(ShieldOpsError):
+    """Raised in enforce mode when a tool call is denied by policy.
+
+    Attributes:
+        tool_name: The tool that was denied.
+        reasons: List of policy violation reasons.
+        risk_score: The computed risk score for the tool call.
+    """
+
+    def __init__(
+        self,
+        tool_name: str = "",
+        reasons: list[str] | None = None,
+        risk_score: float = 0.0,
+    ) -> None:
+        self.tool_name = tool_name
+        self.reasons = reasons or []
+        self.risk_score = risk_score
+        detail = f"Tool '{tool_name}' denied: {', '.join(self.reasons)}"
+        super().__init__(detail, status_code=403)
+
+
+class ShieldOpsConnectionError(ShieldOpsError):
+    """Raised when the SDK cannot reach the ShieldOps API."""
+
+    def __init__(self, message: str = "Failed to connect to ShieldOps API") -> None:
+        super().__init__(message, status_code=None)

shieldops_sdk/experimental/__init__.py

@@ -0,0 +1,22 @@
+"""ShieldOps SDK experimental namespace.
+
+Modules under ``shieldops_sdk.experimental`` provide integrations for AI agent
+frameworks whose APIs are still in flux. Their public surface may change in any
+minor SDK release without a deprecation cycle — pin a specific SDK version if
+you depend on these.
+
+Stable integrations live under ``shieldops_sdk.integrations``.
+"""
+
+from __future__ import annotations
+
+import warnings
+
+warnings.warn(
+    "shieldops_sdk.experimental contains unstable integrations whose surface "
+    "may change without notice between minor releases. Pin your SDK version "
+    "if you depend on these. Stable integrations live under "
+    "shieldops_sdk.integrations.",
+    UserWarning,
+    stacklevel=2,
+)

shieldops_sdk/experimental/autogen.py

@@ -0,0 +1,114 @@
+"""ShieldOps Microsoft AutoGen integration (experimental).
+
+Wraps an AutoGen agent so its ``execute_function`` calls flow through the
+ShieldOps interceptor. In ``enforce`` mode a denied call raises
+``ShieldOpsDeniedError`` from inside the wrapped function — the AutoGen agent
+sees the raise the same as any other tool error.
+
+This module is **experimental**. The legacy version exposed ``record_tool_result``,
+``on_message``, and ``get_audit_report`` for telemetry; those depended on
+internal interceptor methods that the public SDK does not expose. Use
+``shieldops_sdk.telemetry`` for result/latency reporting instead.
+
+Usage::
+
+    from shieldops_sdk.experimental.autogen import ShieldOpsAutoGenWrapper
+
+    wrapper = ShieldOpsAutoGenWrapper(api_key="sk-...", mode="enforce")
+    secured_agent = wrapper.wrap_agent(my_autogen_agent)
+"""
+
+from __future__ import annotations
+
+import functools
+import logging
+from typing import Any
+
+from shieldops_sdk.config import SDKMode, ShieldOpsConfig
+from shieldops_sdk.interceptor import ShieldOpsInterceptor
+
+logger = logging.getLogger("shieldops_sdk.experimental.autogen")
+
+
+class ShieldOpsAutoGenWrapper:
+    """Wraps Microsoft AutoGen agents with ShieldOps firewall interception."""
+
+    def __init__(
+        self,
+        api_key: str = "",
+        endpoint: str = "https://api.shieldops.io",
+        mode: str = "audit",
+        agent_id: str = "autogen-agent",
+    ) -> None:
+        config = ShieldOpsConfig(
+            api_key=api_key,
+            endpoint=endpoint,
+            mode=SDKMode(mode),
+        )
+        self._config = config
+        self._interceptor = ShieldOpsInterceptor(config)
+        self._agent_id = agent_id
+        logger.info(
+            "shieldops.autogen.initialized agent_id=%s mode=%s",
+            agent_id,
+            mode,
+        )
+
+    def wrap_tool_execution(
+        self,
+        tool_name: str,
+        tool_args: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Evaluate a tool call before execution.
+
+        Returns ``{"action", "risk_score"}``. In ``enforce`` mode a denied call
+        raises ``ShieldOpsDeniedError`` instead of returning.
+        """
+        decision = self._interceptor.check(
+            tool_name,
+            tool_args or {},
+            agent_id=self._agent_id,
+        )
+        return {"action": decision.action, "risk_score": decision.risk_score}
+
+    def wrap_agent(self, agent: Any) -> Any:
+        """Monkey-patch ``agent.execute_function`` to route through the interceptor.
+
+        Returns the same agent for chainability. If the agent has no
+        ``execute_function`` attribute, this is a no-op.
+        """
+        if not hasattr(agent, "execute_function"):
+            logger.info(
+                "shieldops.autogen.wrap_agent_noop agent=%s reason=no_execute_function",
+                getattr(agent, "name", agent),
+            )
+            return agent
+
+        interceptor = self._interceptor
+        agent_id = self._agent_id
+        original_exec = agent.execute_function
+
+        @functools.wraps(original_exec)
+        def wrapped_exec(func_call: dict[str, Any], **kwargs: Any) -> Any:
+            tool_name = func_call.get("name", "unknown_function")
+            tool_args = func_call.get("arguments", {})
+            # check() raises ShieldOpsDeniedError in enforce mode for denied calls.
+            interceptor.check(tool_name, tool_args, agent_id=agent_id)
+            return original_exec(func_call, **kwargs)
+
+        agent.execute_function = wrapped_exec
+        logger.info(
+            "shieldops.autogen.agent_wrapped agent=%s",
+            getattr(agent, "name", agent),
+        )
+        return agent
+
+    @property
+    def interceptor(self) -> ShieldOpsInterceptor:
+        """Underlying interceptor (for advanced telemetry/stats access)."""
+        return self._interceptor
+
+    @property
+    def stats(self) -> dict[str, Any]:
+        """Interception statistics."""
+        return self._interceptor.stats