agent-mcp 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- agent_mcp/__init__.py +66 -12
- agent_mcp/a2a_protocol.py +316 -0
- agent_mcp/agent_lightning_library.py +214 -0
- agent_mcp/claude_mcp_adapter.py +195 -0
- agent_mcp/google_ai_mcp_adapter.py +183 -0
- agent_mcp/llamaindex_mcp_adapter.py +410 -0
- agent_mcp/microsoft_agent_framework.py +591 -0
- agent_mcp/missing_frameworks.py +435 -0
- agent_mcp/openapi_protocol.py +616 -0
- agent_mcp/payments.py +804 -0
- agent_mcp/pydantic_ai_mcp_adapter.py +628 -0
- agent_mcp/registry.py +768 -0
- agent_mcp/security.py +864 -0
- {agent_mcp-0.1.4.dist-info → agent_mcp-0.1.6.dist-info}/METADATA +182 -55
- {agent_mcp-0.1.4.dist-info → agent_mcp-0.1.6.dist-info}/RECORD +19 -6
- {agent_mcp-0.1.4.dist-info → agent_mcp-0.1.6.dist-info}/WHEEL +1 -1
- agent_mcp-0.1.6.dist-info/entry_points.txt +4 -0
- demos/comprehensive_framework_demo.py +202 -0
- agent_mcp-0.1.4.dist-info/entry_points.txt +0 -2
- {agent_mcp-0.1.4.dist-info → agent_mcp-0.1.6.dist-info}/top_level.txt +0 -0
agent_mcp/security.py
ADDED
|
@@ -0,0 +1,864 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Zero Trust Security Layer for AgentMCP
|
|
3
|
+
Comprehensive security implementation with identity, authentication, authorization, and audit
|
|
4
|
+
|
|
5
|
+
This module provides enterprise-grade security for AI agents including:
|
|
6
|
+
- Decentralized Identity (DID)
|
|
7
|
+
- Dynamic Authorization (ABAC)
|
|
8
|
+
- Audit Logging
|
|
9
|
+
- Rate Limiting
|
|
10
|
+
- Capability-based Access Control
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
import json
|
|
14
|
+
import uuid
|
|
15
|
+
import hashlib
|
|
16
|
+
import secrets
|
|
17
|
+
import asyncio
|
|
18
|
+
from datetime import datetime, timezone, timedelta
|
|
19
|
+
from typing import Dict, Any, List, Optional, Callable, Union, Set, Tuple
|
|
20
|
+
from dataclasses import dataclass, asdict
|
|
21
|
+
from enum import Enum
|
|
22
|
+
import jwt
|
|
23
|
+
import bcrypt
|
|
24
|
+
import logging
|
|
25
|
+
from collections import defaultdict, deque
|
|
26
|
+
|
|
27
|
+
logger = logging.getLogger(__name__)
|
|
28
|
+
|
|
29
|
+
class Permission(Enum):
|
|
30
|
+
"""Agent permissions for zero-trust authorization"""
|
|
31
|
+
READ_CONTEXT = "context:read"
|
|
32
|
+
WRITE_CONTEXT = "context:write"
|
|
33
|
+
EXECUTE_TASK = "task:execute"
|
|
34
|
+
SEND_MESSAGE = "message:send"
|
|
35
|
+
RECEIVE_MESSAGE = "message:receive"
|
|
36
|
+
REGISTER_TOOL = "tool:register"
|
|
37
|
+
CALL_TOOL = "tool:call"
|
|
38
|
+
MANAGE_PAYMENT = "payment:manage"
|
|
39
|
+
AGENT_DISCOVERY = "agent:discover"
|
|
40
|
+
SYSTEM_ADMIN = "system:admin"
|
|
41
|
+
|
|
42
|
+
class SecurityLevel(Enum):
|
|
43
|
+
"""Security levels for agents"""
|
|
44
|
+
LOW = "low"
|
|
45
|
+
MEDIUM = "medium"
|
|
46
|
+
HIGH = "high"
|
|
47
|
+
CRITICAL = "critical"
|
|
48
|
+
|
|
49
|
+
@dataclass
|
|
50
|
+
class AgentIdentity:
|
|
51
|
+
"""Decentralized identity for an AI agent"""
|
|
52
|
+
agent_id: str
|
|
53
|
+
did: str # Decentralized Identifier
|
|
54
|
+
public_key: str
|
|
55
|
+
capabilities: List[str]
|
|
56
|
+
owner: Optional[str] # Human owner's identity
|
|
57
|
+
created_at: str
|
|
58
|
+
security_level: SecurityLevel = SecurityLevel.MEDIUM
|
|
59
|
+
reputation_score: float = 0.0
|
|
60
|
+
last_active: Optional[str] = None
|
|
61
|
+
metadata: Dict[str, Any] = None
|
|
62
|
+
|
|
63
|
+
def __post_init__(self):
|
|
64
|
+
if self.metadata is None:
|
|
65
|
+
self.metadata = {}
|
|
66
|
+
if self.created_at is None:
|
|
67
|
+
self.created_at = datetime.now(timezone.utc).isoformat()
|
|
68
|
+
|
|
69
|
+
@dataclass
|
|
70
|
+
class AccessPolicy:
|
|
71
|
+
"""Access control policy for zero-trust authorization"""
|
|
72
|
+
name: str
|
|
73
|
+
required_capabilities: List[str]
|
|
74
|
+
required_permissions: List[Permission]
|
|
75
|
+
max_task_value: Optional[float] = None
|
|
76
|
+
allowed_targets: Optional[List[str]] = None
|
|
77
|
+
denied_targets: Optional[List[str]] = None
|
|
78
|
+
time_restrictions: Optional[Dict[str, Any]] = None
|
|
79
|
+
security_level_required: SecurityLevel = SecurityLevel.MEDIUM
|
|
80
|
+
rate_limits: Optional[Dict[str, Dict[str, int]]] = None
|
|
81
|
+
|
|
82
|
+
@dataclass
|
|
83
|
+
class AuditLogEntry:
|
|
84
|
+
"""Immutable audit log entry"""
|
|
85
|
+
timestamp: str
|
|
86
|
+
agent_id: str
|
|
87
|
+
action: str
|
|
88
|
+
target: str
|
|
89
|
+
resource: str
|
|
90
|
+
result: str
|
|
91
|
+
permission_used: Optional[str] = None
|
|
92
|
+
task_id: Optional[str] = None
|
|
93
|
+
payment_id: Optional[str] = None
|
|
94
|
+
ip_address: Optional[str] = None
|
|
95
|
+
user_agent: Optional[str] = None
|
|
96
|
+
risk_score: float = 0.0
|
|
97
|
+
metadata: Dict[str, Any] = None
|
|
98
|
+
|
|
99
|
+
def __post_init__(self):
|
|
100
|
+
if self.metadata is None:
|
|
101
|
+
self.metadata = {}
|
|
102
|
+
if self.timestamp is None:
|
|
103
|
+
self.timestamp = datetime.now(timezone.utc).isoformat()
|
|
104
|
+
|
|
105
|
+
def compute_hash(self) -> str:
|
|
106
|
+
"""Compute cryptographic hash for integrity"""
|
|
107
|
+
entry_data = {
|
|
108
|
+
"timestamp": self.timestamp,
|
|
109
|
+
"agent_id": self.agent_id,
|
|
110
|
+
"action": self.action,
|
|
111
|
+
"target": self.target,
|
|
112
|
+
"result": self.result
|
|
113
|
+
}
|
|
114
|
+
entry_str = json.dumps(entry_data, sort_keys=True)
|
|
115
|
+
return hashlib.sha256(entry_str.encode()).hexdigest()
|
|
116
|
+
|
|
117
|
+
class DecentralizedIdentityManager:
|
|
118
|
+
"""Manager for decentralized agent identities using DIDs"""
|
|
119
|
+
|
|
120
|
+
def __init__(self, registry_url: str = None, secret_key: str = None):
|
|
121
|
+
self.registry_url = registry_url or "https://did-registry.agentmcp.com"
|
|
122
|
+
self.secret_key = secret_key or secrets.token_urlsafe(32)
|
|
123
|
+
self.identities = {}
|
|
124
|
+
self.revoked = set()
|
|
125
|
+
|
|
126
|
+
def create_did(self, agent_id: str, public_key: str) -> str:
|
|
127
|
+
"""Create a DID for an agent"""
|
|
128
|
+
# did:mcp:agent:<hash>
|
|
129
|
+
hash_input = f"{agent_id}:{public_key}:{datetime.now(timezone.utc).isoformat()}"
|
|
130
|
+
did_hash = hashlib.sha256(hash_input.encode()).hexdigest()[:32]
|
|
131
|
+
return f"did:mcp:agent:{did_hash}"
|
|
132
|
+
|
|
133
|
+
def create_identity(
|
|
134
|
+
self,
|
|
135
|
+
agent_id: str,
|
|
136
|
+
capabilities: List[str],
|
|
137
|
+
owner: str = None,
|
|
138
|
+
security_level: SecurityLevel = SecurityLevel.MEDIUM
|
|
139
|
+
) -> AgentIdentity:
|
|
140
|
+
"""Create a new agent identity"""
|
|
141
|
+
# Generate key pair
|
|
142
|
+
private_key = secrets.token_urlsafe(32)
|
|
143
|
+
public_key = hashlib.sha256(private_key.encode()).hexdigest()
|
|
144
|
+
|
|
145
|
+
# Create DID
|
|
146
|
+
did = self.create_did(agent_id, public_key)
|
|
147
|
+
|
|
148
|
+
# Create identity
|
|
149
|
+
identity = AgentIdentity(
|
|
150
|
+
agent_id=agent_id,
|
|
151
|
+
did=did,
|
|
152
|
+
public_key=public_key,
|
|
153
|
+
capabilities=capabilities,
|
|
154
|
+
owner=owner,
|
|
155
|
+
security_level=security_level
|
|
156
|
+
)
|
|
157
|
+
|
|
158
|
+
# Store locally
|
|
159
|
+
self.identities[agent_id] = identity
|
|
160
|
+
|
|
161
|
+
logger.info(f"Created identity for agent {agent_id}: {did}")
|
|
162
|
+
return identity
|
|
163
|
+
|
|
164
|
+
def verify_identity(self, identity: AgentIdentity, signature: str, data: str) -> bool:
|
|
165
|
+
"""Verify an agent's identity using cryptographic signature"""
|
|
166
|
+
if identity.agent_id in self.revoked:
|
|
167
|
+
return False
|
|
168
|
+
|
|
169
|
+
expected_hash = hashlib.sha256(f"{data}{identity.public_key}".encode()).hexdigest()
|
|
170
|
+
return secrets.compare_digest(signature, expected_hash)
|
|
171
|
+
|
|
172
|
+
def create_verifiable_credential(
|
|
173
|
+
self,
|
|
174
|
+
identity: AgentIdentity,
|
|
175
|
+
capabilities: List[str],
|
|
176
|
+
issuer: str = "AgentMCP"
|
|
177
|
+
) -> Dict[str, Any]:
|
|
178
|
+
"""Create a verifiable credential for an agent"""
|
|
179
|
+
credential = {
|
|
180
|
+
"@context": ["https://www.w3.org/2018/credentials/v1"],
|
|
181
|
+
"type": ["VerifiableCredential", "AgentCapabilityCredential"],
|
|
182
|
+
"issuer": issuer,
|
|
183
|
+
"issuanceDate": datetime.now(timezone.utc).isoformat(),
|
|
184
|
+
"credentialSubject": {
|
|
185
|
+
"id": identity.did,
|
|
186
|
+
"agentId": identity.agent_id,
|
|
187
|
+
"capabilities": capabilities,
|
|
188
|
+
"securityLevel": identity.security_level.value
|
|
189
|
+
},
|
|
190
|
+
"proof": {
|
|
191
|
+
"type": "Ed25519Signature2018",
|
|
192
|
+
"created": datetime.now(timezone.utc).isoformat(),
|
|
193
|
+
"proofPurpose": "assertionMethod",
|
|
194
|
+
"verificationMethod": f"{identity.did}#keys-1",
|
|
195
|
+
"jws": self._sign_credential(identity, capabilities)
|
|
196
|
+
}
|
|
197
|
+
}
|
|
198
|
+
return credential
|
|
199
|
+
|
|
200
|
+
def _sign_credential(self, identity: AgentIdentity, capabilities: List[str]) -> str:
|
|
201
|
+
"""Sign a credential (simplified implementation)"""
|
|
202
|
+
data = f"{identity.did}:{json.dumps(capabilities, sort_keys=True)}"
|
|
203
|
+
return hashlib.sha256(f"{data}{self.secret_key}".encode()).hexdigest()
|
|
204
|
+
|
|
205
|
+
class ZeroTrustAuthorizer:
|
|
206
|
+
"""Zero Trust authorization engine with ABAC (Attribute-Based Access Control)"""
|
|
207
|
+
|
|
208
|
+
def __init__(self, secret_key: str):
|
|
209
|
+
self.secret_key = secret_key
|
|
210
|
+
self.policies: Dict[str, AccessPolicy] = {}
|
|
211
|
+
self.active_sessions: Dict[str, Dict[str, Any]] = {}
|
|
212
|
+
self.decision_cache = {}
|
|
213
|
+
|
|
214
|
+
# Default policies
|
|
215
|
+
self._setup_default_policies()
|
|
216
|
+
|
|
217
|
+
def _setup_default_policies(self):
|
|
218
|
+
"""Setup default zero-trust policies"""
|
|
219
|
+
# Tool access policy
|
|
220
|
+
self.policies["tool_access"] = AccessPolicy(
|
|
221
|
+
name="tool_access",
|
|
222
|
+
required_capabilities=["tool_usage"],
|
|
223
|
+
required_permissions=[Permission.CALL_TOOL],
|
|
224
|
+
max_task_value=100.0,
|
|
225
|
+
security_level_required=SecurityLevel.MEDIUM,
|
|
226
|
+
rate_limits={
|
|
227
|
+
"tool:call": {"requests": 100, "window": 60} # 100 calls per minute
|
|
228
|
+
}
|
|
229
|
+
)
|
|
230
|
+
|
|
231
|
+
# Agent communication policy
|
|
232
|
+
self.policies["agent_communication"] = AccessPolicy(
|
|
233
|
+
name="agent_communication",
|
|
234
|
+
required_capabilities=["messaging"],
|
|
235
|
+
required_permissions=[Permission.SEND_MESSAGE, Permission.RECEIVE_MESSAGE],
|
|
236
|
+
security_level_required=SecurityLevel.LOW,
|
|
237
|
+
rate_limits={
|
|
238
|
+
"message:send": {"requests": 50, "window": 60} # 50 messages per minute
|
|
239
|
+
}
|
|
240
|
+
)
|
|
241
|
+
|
|
242
|
+
# Payment management policy
|
|
243
|
+
self.policies["payment_management"] = AccessPolicy(
|
|
244
|
+
name="payment_management",
|
|
245
|
+
required_capabilities=["payment_processing"],
|
|
246
|
+
required_permissions=[Permission.MANAGE_PAYMENT],
|
|
247
|
+
max_task_value=10000.0,
|
|
248
|
+
security_level_required=SecurityLevel.HIGH,
|
|
249
|
+
rate_limits={
|
|
250
|
+
"payment:manage": {"requests": 10, "window": 60} # 10 payment ops per minute
|
|
251
|
+
}
|
|
252
|
+
)
|
|
253
|
+
|
|
254
|
+
# System administration policy
|
|
255
|
+
self.policies["system_admin"] = AccessPolicy(
|
|
256
|
+
name="system_admin",
|
|
257
|
+
required_capabilities=["admin_access"],
|
|
258
|
+
required_permissions=[Permission.SYSTEM_ADMIN],
|
|
259
|
+
security_level_required=SecurityLevel.CRITICAL,
|
|
260
|
+
rate_limits={
|
|
261
|
+
"system:admin": {"requests": 5, "window": 60} # 5 admin ops per minute
|
|
262
|
+
}
|
|
263
|
+
)
|
|
264
|
+
|
|
265
|
+
def check_authorization(
|
|
266
|
+
self,
|
|
267
|
+
identity: AgentIdentity,
|
|
268
|
+
action: str,
|
|
269
|
+
context: Dict[str, Any],
|
|
270
|
+
resource: str = None
|
|
271
|
+
) -> Dict[str, Any]:
|
|
272
|
+
"""Perform zero-trust authorization check"""
|
|
273
|
+
# Get relevant policy
|
|
274
|
+
policy_name = self._determine_policy(action, resource)
|
|
275
|
+
policy = self.policies.get(policy_name)
|
|
276
|
+
|
|
277
|
+
if not policy:
|
|
278
|
+
return {
|
|
279
|
+
"allowed": False,
|
|
280
|
+
"reason": f"No policy found for action: {action}",
|
|
281
|
+
"policy_used": None
|
|
282
|
+
}
|
|
283
|
+
|
|
284
|
+
# Check security level requirement
|
|
285
|
+
if self._compare_security_levels(identity.security_level, policy.security_level_required) < 0:
|
|
286
|
+
return {
|
|
287
|
+
"allowed": False,
|
|
288
|
+
"reason": f"Security level {identity.security_level.value} insufficient, requires {policy.security_level_required.value}",
|
|
289
|
+
"policy_used": policy_name
|
|
290
|
+
}
|
|
291
|
+
|
|
292
|
+
# Check required capabilities
|
|
293
|
+
missing_caps = set(policy.required_capabilities) - set(identity.capabilities)
|
|
294
|
+
if missing_caps:
|
|
295
|
+
return {
|
|
296
|
+
"allowed": False,
|
|
297
|
+
"reason": f"Missing capabilities: {list(missing_caps)}",
|
|
298
|
+
"policy_used": policy_name
|
|
299
|
+
}
|
|
300
|
+
|
|
301
|
+
# Check target restrictions
|
|
302
|
+
target = context.get("target")
|
|
303
|
+
if policy.denied_targets and target in policy.denied_targets:
|
|
304
|
+
return {
|
|
305
|
+
"allowed": False,
|
|
306
|
+
"reason": f"Target {target} is explicitly denied",
|
|
307
|
+
"policy_used": policy_name
|
|
308
|
+
}
|
|
309
|
+
|
|
310
|
+
if policy.allowed_targets and target not in policy.allowed_targets:
|
|
311
|
+
return {
|
|
312
|
+
"allowed": False,
|
|
313
|
+
"reason": f"Target {target} not in allowed list",
|
|
314
|
+
"policy_used": policy_name
|
|
315
|
+
}
|
|
316
|
+
|
|
317
|
+
# Check task value limits
|
|
318
|
+
task_value = context.get("value", 0)
|
|
319
|
+
if policy.max_task_value and task_value > policy.max_task_value:
|
|
320
|
+
return {
|
|
321
|
+
"allowed": False,
|
|
322
|
+
"reason": f"Task value {task_value} exceeds limit {policy.max_task_value}",
|
|
323
|
+
"policy_used": policy_name
|
|
324
|
+
}
|
|
325
|
+
|
|
326
|
+
# Check time restrictions
|
|
327
|
+
if policy.time_restrictions:
|
|
328
|
+
current_time = datetime.now(timezone.utc)
|
|
329
|
+
if not self._check_time_restrictions(current_time, policy.time_restrictions):
|
|
330
|
+
return {
|
|
331
|
+
"allowed": False,
|
|
332
|
+
"reason": "Access outside allowed time window",
|
|
333
|
+
"policy_used": policy_name
|
|
334
|
+
}
|
|
335
|
+
|
|
336
|
+
# All checks passed
|
|
337
|
+
return {
|
|
338
|
+
"allowed": True,
|
|
339
|
+
"reason": "All authorization checks passed",
|
|
340
|
+
"policy_used": policy_name
|
|
341
|
+
}
|
|
342
|
+
|
|
343
|
+
def _determine_policy(self, action: str, resource: str) -> str:
|
|
344
|
+
"""Determine which policy applies to an action"""
|
|
345
|
+
if action.startswith("tool:"):
|
|
346
|
+
return "tool_access"
|
|
347
|
+
elif action.startswith("message:"):
|
|
348
|
+
return "agent_communication"
|
|
349
|
+
elif action.startswith("payment:"):
|
|
350
|
+
return "payment_management"
|
|
351
|
+
elif action.startswith("system:"):
|
|
352
|
+
return "system_admin"
|
|
353
|
+
else:
|
|
354
|
+
return "tool_access" # Default
|
|
355
|
+
|
|
356
|
+
def _compare_security_levels(self, current: SecurityLevel, required: SecurityLevel) -> int:
|
|
357
|
+
"""Compare security levels (-1: insufficient, 0: equal, 1: sufficient)"""
|
|
358
|
+
levels = {
|
|
359
|
+
SecurityLevel.LOW: 1,
|
|
360
|
+
SecurityLevel.MEDIUM: 2,
|
|
361
|
+
SecurityLevel.HIGH: 3,
|
|
362
|
+
SecurityLevel.CRITICAL: 4
|
|
363
|
+
}
|
|
364
|
+
return levels[current] - levels[required]
|
|
365
|
+
|
|
366
|
+
def _check_time_restrictions(self, current_time: datetime, restrictions: Dict[str, Any]) -> bool:
|
|
367
|
+
"""Check if current time is within allowed restrictions"""
|
|
368
|
+
# Simple implementation - can be extended
|
|
369
|
+
if "allowed_hours" in restrictions:
|
|
370
|
+
current_hour = current_time.hour
|
|
371
|
+
allowed_start = restrictions["allowed_hours"]["start"]
|
|
372
|
+
allowed_end = restrictions["allowed_hours"]["end"]
|
|
373
|
+
|
|
374
|
+
if not (allowed_start <= current_hour <= allowed_end):
|
|
375
|
+
return False
|
|
376
|
+
|
|
377
|
+
if "allowed_days" in restrictions:
|
|
378
|
+
current_day = current_time.weekday() # 0 = Monday
|
|
379
|
+
if current_day not in restrictions["allowed_days"]:
|
|
380
|
+
return False
|
|
381
|
+
|
|
382
|
+
return True
|
|
383
|
+
|
|
384
|
+
def create_scoped_token(
|
|
385
|
+
self,
|
|
386
|
+
identity: AgentIdentity,
|
|
387
|
+
permissions: List[Permission],
|
|
388
|
+
ttl_minutes: int = 15,
|
|
389
|
+
context: Dict[str, Any] = None
|
|
390
|
+
) -> str:
|
|
391
|
+
"""Create a short-lived scoped JWT token"""
|
|
392
|
+
now = datetime.utcnow()
|
|
393
|
+
|
|
394
|
+
payload = {
|
|
395
|
+
"sub": identity.did,
|
|
396
|
+
"agent_id": identity.agent_id,
|
|
397
|
+
"permissions": [p.value for p in permissions],
|
|
398
|
+
"security_level": identity.security_level.value,
|
|
399
|
+
"iat": int(now.timestamp()),
|
|
400
|
+
"exp": int((now + timedelta(minutes=ttl_minutes)).timestamp()),
|
|
401
|
+
"jti": str(uuid.uuid4()), # Unique token ID
|
|
402
|
+
"context": context or {}
|
|
403
|
+
}
|
|
404
|
+
|
|
405
|
+
# Add capability-based claims
|
|
406
|
+
if identity.capabilities:
|
|
407
|
+
payload["capabilities"] = identity.capabilities
|
|
408
|
+
|
|
409
|
+
token = jwt.encode(payload, self.secret_key, algorithm="HS256")
|
|
410
|
+
|
|
411
|
+
# Store session info
|
|
412
|
+
session_id = payload["jti"]
|
|
413
|
+
self.active_sessions[session_id] = {
|
|
414
|
+
"agent_id": identity.agent_id,
|
|
415
|
+
"permissions": permissions,
|
|
416
|
+
"created_at": now.isoformat(),
|
|
417
|
+
"last_used": now.isoformat(),
|
|
418
|
+
"usage_count": 0
|
|
419
|
+
}
|
|
420
|
+
|
|
421
|
+
return token
|
|
422
|
+
|
|
423
|
+
def verify_token(self, token: str) -> Optional[Dict[str, Any]]:
|
|
424
|
+
"""Verify and decode a JWT token"""
|
|
425
|
+
try:
|
|
426
|
+
payload = jwt.decode(token, self.secret_key, algorithms=["HS256"])
|
|
427
|
+
|
|
428
|
+
# Check if token is expired
|
|
429
|
+
if datetime.utcnow() > datetime.fromtimestamp(payload["exp"]):
|
|
430
|
+
return None
|
|
431
|
+
|
|
432
|
+
# Check if session is still active
|
|
433
|
+
session_id = payload.get("jti")
|
|
434
|
+
if session_id and session_id not in self.active_sessions:
|
|
435
|
+
return None
|
|
436
|
+
|
|
437
|
+
return payload
|
|
438
|
+
|
|
439
|
+
except jwt.InvalidTokenError:
|
|
440
|
+
return None
|
|
441
|
+
except Exception as e:
|
|
442
|
+
logger.error(f"Error verifying token: {e}")
|
|
443
|
+
return None
|
|
444
|
+
|
|
445
|
+
def revoke_token(self, token: str) -> bool:
|
|
446
|
+
"""Revoke a JWT token"""
|
|
447
|
+
try:
|
|
448
|
+
payload = jwt.decode(token, self.secret_key, algorithms=["HS256"])
|
|
449
|
+
session_id = payload.get("jti")
|
|
450
|
+
|
|
451
|
+
if session_id and session_id in self.active_sessions:
|
|
452
|
+
del self.active_sessions[session_id]
|
|
453
|
+
logger.info(f"Revoked token session: {session_id}")
|
|
454
|
+
return True
|
|
455
|
+
except Exception as e:
|
|
456
|
+
logger.error(f"Error revoking token: {e}")
|
|
457
|
+
|
|
458
|
+
return False
|
|
459
|
+
|
|
460
|
+
class RateLimiter:
|
|
461
|
+
"""Advanced rate limiting for agent operations"""
|
|
462
|
+
|
|
463
|
+
def __init__(self):
|
|
464
|
+
self.requests = defaultdict(lambda: defaultdict(deque))
|
|
465
|
+
self.limits = {}
|
|
466
|
+
self.cleanup_interval = 300 # 5 minutes
|
|
467
|
+
self.last_cleanup = datetime.now(timezone.utc)
|
|
468
|
+
|
|
469
|
+
def set_rate_limit(self, key: str, limit: int, window: int):
|
|
470
|
+
"""Set rate limit for a specific key"""
|
|
471
|
+
self.limits[key] = {"limit": limit, "window": window}
|
|
472
|
+
|
|
473
|
+
def check_rate_limit(self, agent_id: str, action: str) -> Dict[str, Any]:
|
|
474
|
+
"""Check if agent is within rate limits"""
|
|
475
|
+
# Cleanup old entries
|
|
476
|
+
self._cleanup_old_entries()
|
|
477
|
+
|
|
478
|
+
key = f"{agent_id}:{action}"
|
|
479
|
+
limit_info = self.limits.get(action, {"limit": 100, "window": 60})
|
|
480
|
+
|
|
481
|
+
# Get current requests in window
|
|
482
|
+
now = datetime.now(timezone.utc)
|
|
483
|
+
cutoff_time = now - timedelta(seconds=limit_info["window"])
|
|
484
|
+
|
|
485
|
+
# Remove old requests
|
|
486
|
+
while (self.requests[agent_id][action] and
|
|
487
|
+
self.requests[agent_id][action][0] < cutoff_time):
|
|
488
|
+
self.requests[agent_id][action].popleft()
|
|
489
|
+
|
|
490
|
+
# Check limit
|
|
491
|
+
request_count = len(self.requests[agent_id][action])
|
|
492
|
+
|
|
493
|
+
if request_count >= limit_info["limit"]:
|
|
494
|
+
return {
|
|
495
|
+
"allowed": False,
|
|
496
|
+
"remaining": 0,
|
|
497
|
+
"reset_time": (self.requests[agent_id][action][0] + timedelta(seconds=limit_info["window"])).isoformat() if self.requests[agent_id][action] else (now + timedelta(seconds=limit_info["window"])).isoformat(),
|
|
498
|
+
"limit": limit_info["limit"],
|
|
499
|
+
"window": limit_info["window"]
|
|
500
|
+
}
|
|
501
|
+
|
|
502
|
+
# Add current request
|
|
503
|
+
self.requests[agent_id][action].append(now)
|
|
504
|
+
|
|
505
|
+
return {
|
|
506
|
+
"allowed": True,
|
|
507
|
+
"remaining": limit_info["limit"] - (request_count + 1),
|
|
508
|
+
"reset_time": (now + timedelta(seconds=limit_info["window"])).isoformat(),
|
|
509
|
+
"limit": limit_info["limit"],
|
|
510
|
+
"window": limit_info["window"]
|
|
511
|
+
}
|
|
512
|
+
|
|
513
|
+
def _cleanup_old_entries(self):
|
|
514
|
+
"""Cleanup old rate limit entries"""
|
|
515
|
+
now = datetime.now(timezone.utc)
|
|
516
|
+
|
|
517
|
+
if (now - self.last_cleanup).seconds < self.cleanup_interval:
|
|
518
|
+
return
|
|
519
|
+
|
|
520
|
+
for agent_id in self.requests:
|
|
521
|
+
for action in self.requests[agent_id]:
|
|
522
|
+
cutoff_time = now - timedelta(seconds=300) # 5 minutes
|
|
523
|
+
while (self.requests[agent_id][action] and
|
|
524
|
+
self.requests[agent_id][action][0] < cutoff_time):
|
|
525
|
+
self.requests[agent_id][action].popleft()
|
|
526
|
+
|
|
527
|
+
self.last_cleanup = now
|
|
528
|
+
|
|
529
|
+
class AuditLogger:
|
|
530
|
+
"""Immutable audit logging for compliance"""
|
|
531
|
+
|
|
532
|
+
def __init__(self, storage_backend):
|
|
533
|
+
self.storage = storage_backend # Could be Firestore, PostgreSQL, etc.
|
|
534
|
+
self.logger = logging.getLogger("audit")
|
|
535
|
+
|
|
536
|
+
async def log_action(
|
|
537
|
+
self,
|
|
538
|
+
agent_id: str,
|
|
539
|
+
action: str,
|
|
540
|
+
target: str = None,
|
|
541
|
+
resource: str = None,
|
|
542
|
+
result: str = "success",
|
|
543
|
+
permission_used: str = None,
|
|
544
|
+
task_id: str = None,
|
|
545
|
+
payment_id: str = None,
|
|
546
|
+
context: Dict[str, Any] = None
|
|
547
|
+
) -> bool:
|
|
548
|
+
"""Log an agent action for audit trail"""
|
|
549
|
+
try:
|
|
550
|
+
# Calculate risk score
|
|
551
|
+
risk_score = self._calculate_risk_score(action, result, context)
|
|
552
|
+
|
|
553
|
+
entry = AuditLogEntry(
|
|
554
|
+
timestamp=datetime.now(timezone.utc).isoformat(),
|
|
555
|
+
agent_id=agent_id,
|
|
556
|
+
action=action,
|
|
557
|
+
target=target or "unknown",
|
|
558
|
+
resource=resource or "unknown",
|
|
559
|
+
result=result,
|
|
560
|
+
permission_used=permission_used,
|
|
561
|
+
task_id=task_id,
|
|
562
|
+
payment_id=payment_id,
|
|
563
|
+
risk_score=risk_score,
|
|
564
|
+
ip_address=context.get("ip_address") if context else None,
|
|
565
|
+
user_agent=context.get("user_agent") if context else None,
|
|
566
|
+
metadata=context or {}
|
|
567
|
+
)
|
|
568
|
+
|
|
569
|
+
# Store with integrity hash
|
|
570
|
+
entry_hash = entry.compute_hash()
|
|
571
|
+
entry_data = asdict(entry)
|
|
572
|
+
entry_data["hash"] = entry_hash
|
|
573
|
+
|
|
574
|
+
# Store in backend
|
|
575
|
+
await self.storage.write("audit_logs", entry_data)
|
|
576
|
+
|
|
577
|
+
# Also log to standard logger
|
|
578
|
+
self.logger.info(f"Audit: {agent_id} performed {action} on {target} - {result}")
|
|
579
|
+
|
|
580
|
+
return True
|
|
581
|
+
|
|
582
|
+
except Exception as e:
|
|
583
|
+
self.logger.error(f"Failed to log audit entry: {e}")
|
|
584
|
+
return False
|
|
585
|
+
|
|
586
|
+
def _calculate_risk_score(self, action: str, result: str, context: Dict[str, Any]) -> float:
|
|
587
|
+
"""Calculate risk score for the action"""
|
|
588
|
+
base_score = 0.0
|
|
589
|
+
|
|
590
|
+
# Higher risk for sensitive actions
|
|
591
|
+
if action.startswith("payment:"):
|
|
592
|
+
base_score += 0.7
|
|
593
|
+
elif action.startswith("system:admin"):
|
|
594
|
+
base_score += 0.9
|
|
595
|
+
elif action.startswith("tool:"):
|
|
596
|
+
base_score += 0.3
|
|
597
|
+
|
|
598
|
+
# Higher risk for failures
|
|
599
|
+
if result == "error":
|
|
600
|
+
base_score += 0.2
|
|
601
|
+
|
|
602
|
+
# Contextual factors
|
|
603
|
+
if context:
|
|
604
|
+
if context.get("is_external_request", False):
|
|
605
|
+
base_score += 0.1
|
|
606
|
+
if context.get("late_night_access", False):
|
|
607
|
+
base_score += 0.15
|
|
608
|
+
|
|
609
|
+
return min(base_score, 1.0)
|
|
610
|
+
|
|
611
|
+
async def verify_audit_chain(self, agent_id: str, start_time: str = None, end_time: str = None) -> Dict[str, Any]:
|
|
612
|
+
"""Verify integrity of audit log chain"""
|
|
613
|
+
try:
|
|
614
|
+
# Get audit logs for the agent
|
|
615
|
+
logs = await self.storage.query(
|
|
616
|
+
"audit_logs",
|
|
617
|
+
filters={"agent_id": agent_id, "timestamp_gte": start_time, "timestamp_lte": end_time}
|
|
618
|
+
)
|
|
619
|
+
|
|
620
|
+
if not logs:
|
|
621
|
+
return {"verified": False, "reason": "No audit logs found"}
|
|
622
|
+
|
|
623
|
+
# Verify hash chain
|
|
624
|
+
verification_errors = []
|
|
625
|
+
for i, log in enumerate(logs[1:], 1):
|
|
626
|
+
if "hash" not in log:
|
|
627
|
+
verification_errors.append(f"Missing hash for log at index {i}")
|
|
628
|
+
continue
|
|
629
|
+
|
|
630
|
+
# Recalculate hash of previous log
|
|
631
|
+
prev_log = logs[i-1]
|
|
632
|
+
prev_data = {
|
|
633
|
+
"timestamp": prev_log["timestamp"],
|
|
634
|
+
"agent_id": prev_log["agent_id"],
|
|
635
|
+
"action": prev_log["action"],
|
|
636
|
+
"target": prev_log["target"],
|
|
637
|
+
"result": prev_log["result"]
|
|
638
|
+
}
|
|
639
|
+
expected_hash = hashlib.sha256(json.dumps(prev_data, sort_keys=True).encode()).hexdigest()
|
|
640
|
+
|
|
641
|
+
if log["hash"] != expected_hash:
|
|
642
|
+
verification_errors.append(f"Hash mismatch at index {i}")
|
|
643
|
+
|
|
644
|
+
return {
|
|
645
|
+
"verified": len(verification_errors) == 0,
|
|
646
|
+
"total_logs": len(logs),
|
|
647
|
+
"verification_errors": verification_errors
|
|
648
|
+
}
|
|
649
|
+
|
|
650
|
+
except Exception as e:
|
|
651
|
+
logger.error(f"Error verifying audit chain: {e}")
|
|
652
|
+
return {"verified": False, "reason": str(e)}
|
|
653
|
+
|
|
654
|
+
class ZeroTrustSecurityLayer:
|
|
655
|
+
"""Main security coordinator combining all security components"""
|
|
656
|
+
|
|
657
|
+
def __init__(
|
|
658
|
+
self,
|
|
659
|
+
secret_key: str = None,
|
|
660
|
+
storage_backend = None,
|
|
661
|
+
did_registry_url: str = None
|
|
662
|
+
):
|
|
663
|
+
self.secret_key = secret_key or secrets.token_urlsafe(32)
|
|
664
|
+
|
|
665
|
+
# Initialize security components
|
|
666
|
+
self.identity_manager = DecentralizedIdentityManager(did_registry_url, self.secret_key)
|
|
667
|
+
self.authorizer = ZeroTrustAuthorizer(self.secret_key)
|
|
668
|
+
self.rate_limiter = RateLimiter()
|
|
669
|
+
self.audit_logger = AuditLogger(storage_backend)
|
|
670
|
+
|
|
671
|
+
# Setup rate limits
|
|
672
|
+
self._setup_rate_limits()
|
|
673
|
+
|
|
674
|
+
def _setup_rate_limits(self):
|
|
675
|
+
"""Setup default rate limits"""
|
|
676
|
+
self.rate_limiter.set_rate_limit("tool:call", 200, 60) # 200 tool calls per minute
|
|
677
|
+
self.rate_limiter.set_rate_limit("message:send", 100, 60) # 100 messages per minute
|
|
678
|
+
self.rate_limiter.set_rate_limit("payment:manage", 10, 60) # 10 payment ops per minute
|
|
679
|
+
self.rate_limiter.set_rate_limit("system:admin", 5, 60) # 5 admin ops per minute
|
|
680
|
+
|
|
681
|
+
async def create_agent_identity(
|
|
682
|
+
self,
|
|
683
|
+
agent_id: str,
|
|
684
|
+
capabilities: List[str],
|
|
685
|
+
owner: str = None,
|
|
686
|
+
security_level: SecurityLevel = SecurityLevel.MEDIUM
|
|
687
|
+
) -> Tuple[AgentIdentity, str]:
|
|
688
|
+
"""Create agent identity and return authentication token"""
|
|
689
|
+
# Create identity
|
|
690
|
+
identity = self.identity_manager.create_identity(
|
|
691
|
+
agent_id=agent_id,
|
|
692
|
+
capabilities=capabilities,
|
|
693
|
+
owner=owner,
|
|
694
|
+
security_level=security_level
|
|
695
|
+
)
|
|
696
|
+
|
|
697
|
+
# Create initial token
|
|
698
|
+
token = self.authorizer.create_scoped_token(
|
|
699
|
+
identity=identity,
|
|
700
|
+
permissions=[
|
|
701
|
+
Permission.READ_CONTEXT,
|
|
702
|
+
Permission.WRITE_CONTEXT,
|
|
703
|
+
Permission.EXECUTE_TASK,
|
|
704
|
+
Permission.CALL_TOOL
|
|
705
|
+
],
|
|
706
|
+
ttl_minutes=60 # Longer initial token
|
|
707
|
+
)
|
|
708
|
+
|
|
709
|
+
await self.audit_logger.log_action(
|
|
710
|
+
agent_id=agent_id,
|
|
711
|
+
action="identity:created",
|
|
712
|
+
resource="agent_identity",
|
|
713
|
+
context={"security_level": security_level.value}
|
|
714
|
+
)
|
|
715
|
+
|
|
716
|
+
return identity, token
|
|
717
|
+
|
|
718
|
+
async def authenticate_agent(self, token: str, action: str, context: Dict[str, Any] = None) -> Optional[Dict[str, Any]]:
|
|
719
|
+
"""Authenticate an agent token for a specific action"""
|
|
720
|
+
# Verify token
|
|
721
|
+
token_data = self.authorizer.verify_token(token)
|
|
722
|
+
if not token_data:
|
|
723
|
+
await self.audit_logger.log_action(
|
|
724
|
+
agent_id="unknown",
|
|
725
|
+
action="auth:failed",
|
|
726
|
+
result="error",
|
|
727
|
+
context={"reason": "invalid_token"}
|
|
728
|
+
)
|
|
729
|
+
return None
|
|
730
|
+
|
|
731
|
+
# Get agent identity
|
|
732
|
+
agent_id = token_data["agent_id"]
|
|
733
|
+
identity = self.identity_manager.identities.get(agent_id)
|
|
734
|
+
if not identity:
|
|
735
|
+
await self.audit_logger.log_action(
|
|
736
|
+
agent_id=agent_id,
|
|
737
|
+
action="auth:failed",
|
|
738
|
+
result="error",
|
|
739
|
+
context={"reason": "identity_not_found"}
|
|
740
|
+
)
|
|
741
|
+
return None
|
|
742
|
+
|
|
743
|
+
# Check rate limits
|
|
744
|
+
rate_check = self.rate_limiter.check_rate_limit(agent_id, action)
|
|
745
|
+
if not rate_check["allowed"]:
|
|
746
|
+
await self.audit_logger.log_action(
|
|
747
|
+
agent_id=agent_id,
|
|
748
|
+
action="auth:failed",
|
|
749
|
+
result="error",
|
|
750
|
+
context={"reason": "rate_limit_exceeded", "rate_info": rate_check}
|
|
751
|
+
)
|
|
752
|
+
return None
|
|
753
|
+
|
|
754
|
+
# Authorization check
|
|
755
|
+
authz_result = self.authorizer.check_authorization(identity, action, context or {})
|
|
756
|
+
|
|
757
|
+
if not authz_result["allowed"]:
|
|
758
|
+
await self.audit_logger.log_action(
|
|
759
|
+
agent_id=agent_id,
|
|
760
|
+
action="auth:denied",
|
|
761
|
+
result="denied",
|
|
762
|
+
permission_used=action,
|
|
763
|
+
context={"reason": authz_result["reason"], "policy": authz_result["policy_used"]}
|
|
764
|
+
)
|
|
765
|
+
return None
|
|
766
|
+
|
|
767
|
+
# Update session usage
|
|
768
|
+
session_id = token_data.get("jti")
|
|
769
|
+
if session_id and session_id in self.authorizer.active_sessions:
|
|
770
|
+
self.authorizer.active_sessions[session_id]["last_used"] = datetime.utcnow().isoformat()
|
|
771
|
+
self.authorizer.active_sessions[session_id]["usage_count"] += 1
|
|
772
|
+
|
|
773
|
+
# Log successful authentication
|
|
774
|
+
await self.audit_logger.log_action(
|
|
775
|
+
agent_id=agent_id,
|
|
776
|
+
action="auth:success",
|
|
777
|
+
result="success",
|
|
778
|
+
permission_used=action,
|
|
779
|
+
context={"policy": authz_result["policy_used"]}
|
|
780
|
+
)
|
|
781
|
+
|
|
782
|
+
return {
|
|
783
|
+
"identity": identity,
|
|
784
|
+
"token_data": token_data,
|
|
785
|
+
"authorization": authz_result
|
|
786
|
+
}
|
|
787
|
+
|
|
788
|
+
async def create_capability_token(
|
|
789
|
+
self,
|
|
790
|
+
agent_id: str,
|
|
791
|
+
permissions: List[Permission],
|
|
792
|
+
ttl_minutes: int = 15,
|
|
793
|
+
context: Dict[str, Any] = None
|
|
794
|
+
) -> Optional[str]:
|
|
795
|
+
"""Create a scoped capability token"""
|
|
796
|
+
identity = self.identity_manager.identities.get(agent_id)
|
|
797
|
+
if not identity:
|
|
798
|
+
return None
|
|
799
|
+
|
|
800
|
+
token = self.authorizer.create_scoped_token(identity, permissions, ttl_minutes, context)
|
|
801
|
+
|
|
802
|
+
await self.audit_logger.log_action(
|
|
803
|
+
agent_id=agent_id,
|
|
804
|
+
action="token:created",
|
|
805
|
+
resource="capability_token",
|
|
806
|
+
context={
|
|
807
|
+
"permissions": [p.value for p in permissions],
|
|
808
|
+
"ttl_minutes": ttl_minutes
|
|
809
|
+
}
|
|
810
|
+
)
|
|
811
|
+
|
|
812
|
+
return token
|
|
813
|
+
|
|
814
|
+
async def revoke_agent_access(self, agent_id: str, reason: str = None) -> bool:
|
|
815
|
+
"""Revoke an agent's access"""
|
|
816
|
+
# Remove from identity manager
|
|
817
|
+
if agent_id in self.identity_manager.identities:
|
|
818
|
+
del self.identity_manager.identities[agent_id]
|
|
819
|
+
|
|
820
|
+
# Add to revoked list
|
|
821
|
+
self.identity_manager.revoked.add(agent_id)
|
|
822
|
+
|
|
823
|
+
# Clear active sessions
|
|
824
|
+
sessions_to_remove = [
|
|
825
|
+
session_id for session_id, session in self.authorizer.active_sessions.items()
|
|
826
|
+
if session.get("agent_id") == agent_id
|
|
827
|
+
]
|
|
828
|
+
|
|
829
|
+
for session_id in sessions_to_remove:
|
|
830
|
+
del self.authorizer.active_sessions[session_id]
|
|
831
|
+
|
|
832
|
+
await self.audit_logger.log_action(
|
|
833
|
+
agent_id=agent_id,
|
|
834
|
+
action="access:revoked",
|
|
835
|
+
result="revoked",
|
|
836
|
+
context={"reason": reason or "administrative_action"}
|
|
837
|
+
)
|
|
838
|
+
|
|
839
|
+
return True
|
|
840
|
+
|
|
841
|
+
def get_security_status(self) -> Dict[str, Any]:
|
|
842
|
+
"""Get overall security status"""
|
|
843
|
+
return {
|
|
844
|
+
"active_identities": len(self.identity_manager.identities),
|
|
845
|
+
"revoked_identities": len(self.identity_manager.revoked),
|
|
846
|
+
"active_sessions": len(self.authorizer.active_sessions),
|
|
847
|
+
"security_policies": list(self.authorizer.policies.keys()),
|
|
848
|
+
"rate_limits": self.rate_limiter.limits,
|
|
849
|
+
"timestamp": datetime.now(timezone.utc).isoformat()
|
|
850
|
+
}
|
|
851
|
+
|
|
852
|
+
# Export for easy importing
|
|
853
|
+
__all__ = [
|
|
854
|
+
'Permission',
|
|
855
|
+
'SecurityLevel',
|
|
856
|
+
'AgentIdentity',
|
|
857
|
+
'AccessPolicy',
|
|
858
|
+
'AuditLogEntry',
|
|
859
|
+
'DecentralizedIdentityManager',
|
|
860
|
+
'ZeroTrustAuthorizer',
|
|
861
|
+
'RateLimiter',
|
|
862
|
+
'AuditLogger',
|
|
863
|
+
'ZeroTrustSecurityLayer'
|
|
864
|
+
]
|