|
2 | 2 |
|
3 | 3 | from __future__ import annotations |
4 | 4 |
|
| 5 | +import contextlib |
5 | 6 | import json |
6 | 7 | import os |
7 | 8 | import time |
8 | 9 | from dataclasses import dataclass, field |
9 | 10 | from pathlib import Path |
10 | 11 | from typing import Any, Optional |
11 | 12 |
|
12 | | -from cryptography.fernet import Fernet, InvalidToken |
13 | | - |
14 | 13 | from teaagent.audit import AuditLogger |
15 | 14 | from teaagent.hooks import HookError, HookRegistry |
16 | 15 | from teaagent.tools import ToolRegistry |
17 | 16 |
|
| 17 | +Fernet: Any = None |
| 18 | +with contextlib.suppress(ImportError): |
| 19 | + from cryptography import fernet |
| 20 | + |
| 21 | + Fernet = fernet.Fernet |
| 22 | + |
18 | 23 |
|
19 | 24 | @dataclass |
20 | 25 | class MCPServerTrust: |
@@ -96,8 +101,14 @@ def trust_policy_path(root: str | Path) -> Path: |
96 | 101 | return Path(root).resolve() / '.teaagent' / 'mcp-trust.json' |
97 | 102 |
|
98 | 103 |
|
99 | | -def _get_trust_policy_fernet() -> Fernet: |
| 104 | +def _get_trust_policy_fernet() -> Any: |
100 | 105 | """Get Fernet instance for MCP trust policy encryption with validation.""" |
| 106 | + if Fernet is None: |
| 107 | + raise ValueError( |
| 108 | + "MCP trust policy encryption requires the 'cryptography' package. " |
| 109 | + "Install it using: pip install 'teaagent[crypto]'" |
| 110 | + ) |
| 111 | + |
101 | 112 | if 'TEAAGENT_MCP_TRUST_KEY' not in os.environ: |
102 | 113 | raise ValueError( |
103 | 114 | 'TEAAGENT_MCP_TRUST_KEY environment variable is required for MCP trust policy encryption. ' |
@@ -143,8 +154,12 @@ def load_mcp_trust_policy(root: str | Path) -> MCPTrustPolicy: |
143 | 154 | try: |
144 | 155 | raw_text = path.read_text(encoding='utf-8') |
145 | 156 | payload = _deserialize_policy(raw_text) |
146 | | - except (OSError, json.JSONDecodeError, InvalidToken, KeyError, ValueError): |
| 157 | + except (OSError, json.JSONDecodeError, KeyError, ValueError): |
147 | 158 | return MCPTrustPolicy() |
| 159 | + except Exception as exc: |
| 160 | + if exc.__class__.__name__ == 'InvalidToken': |
| 161 | + return MCPTrustPolicy() |
| 162 | + raise |
148 | 163 | return MCPTrustPolicy.from_dict(payload) |
149 | 164 |
|
150 | 165 |
|
|
0 commit comments