secure-ai-guard 0.2.0__tar.gz
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.
- secure_ai_guard-0.2.0/.gitignore +8 -0
- secure_ai_guard-0.2.0/LICENSE +21 -0
- secure_ai_guard-0.2.0/PKG-INFO +112 -0
- secure_ai_guard-0.2.0/README.md +99 -0
- secure_ai_guard-0.2.0/pyproject.toml +28 -0
- secure_ai_guard-0.2.0/secure_ai/__init__.py +512 -0
- secure_ai_guard-0.2.0/tests/test_secure_ai.py +445 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Secure AI
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,112 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: secure-ai-guard
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Data loss prevention for AI agents. Inspect what an agent is about to do, before it does it.
|
|
5
|
+
Project-URL: Homepage, https://secureai.one/developers
|
|
6
|
+
License: MIT
|
|
7
|
+
License-File: LICENSE
|
|
8
|
+
Keywords: ai-agents,audit,dlp,mcp,redaction,security
|
|
9
|
+
Requires-Python: >=3.9
|
|
10
|
+
Provides-Extra: dev
|
|
11
|
+
Requires-Dist: pytest>=7; extra == 'dev'
|
|
12
|
+
Description-Content-Type: text/markdown
|
|
13
|
+
|
|
14
|
+
# secure-ai
|
|
15
|
+
|
|
16
|
+
Data loss prevention for AI agents. Put a check in front of every action an
|
|
17
|
+
agent takes, before it takes it.
|
|
18
|
+
|
|
19
|
+
```bash
|
|
20
|
+
pip install secure-ai-guard
|
|
21
|
+
```
|
|
22
|
+
|
|
23
|
+
> **Status:** written, not yet executed — there was no Python interpreter on
|
|
24
|
+
> the machine this was built on. Run `pytest` before trusting it. The
|
|
25
|
+
> TypeScript SDK it mirrors does pass its suite.
|
|
26
|
+
|
|
27
|
+
## The one thing worth knowing
|
|
28
|
+
|
|
29
|
+
`guard` wraps a function your agent already calls, and calls it with the
|
|
30
|
+
**rewritten** arguments:
|
|
31
|
+
|
|
32
|
+
```python
|
|
33
|
+
from secure_ai import SecureAI
|
|
34
|
+
|
|
35
|
+
sai = SecureAI(api_key=os.environ["SECURE_AI_KEY"], agent="support-bot")
|
|
36
|
+
|
|
37
|
+
send_email = sai.guard("email.send", mailer.send)
|
|
38
|
+
send_email({"to": "ana@clientfirm.com", "body": "About invoice 4471…"})
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
The agent wrote a real address. `mailer.send` receives a stand-in. Nothing in
|
|
42
|
+
the agent's code had to check a decision — which is the point. Protection that
|
|
43
|
+
depends on remembering to check it ends at the fourteenth call site.
|
|
44
|
+
|
|
45
|
+
If the policy refuses, the wrapped function is never called and `ActionBlocked`
|
|
46
|
+
is raised.
|
|
47
|
+
|
|
48
|
+
As a decorator:
|
|
49
|
+
|
|
50
|
+
```python
|
|
51
|
+
@sai.guarded("email.send")
|
|
52
|
+
def send(payload): ...
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
## A whole toolbelt at once
|
|
56
|
+
|
|
57
|
+
```python
|
|
58
|
+
tools = sai.guard_tools(agent.tools)
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
Wraps whichever attribute holds the callable — `func`, `_run`, `run` — so
|
|
62
|
+
LangChain and similar frameworks are governed in one line. Shapes it does not
|
|
63
|
+
recognise are returned untouched rather than raising.
|
|
64
|
+
|
|
65
|
+
## Rules
|
|
66
|
+
|
|
67
|
+
Rules live on the account, not in the request. An agent cannot argue with them.
|
|
68
|
+
|
|
69
|
+
```python
|
|
70
|
+
sai.set_policy({
|
|
71
|
+
"fallback": "redact",
|
|
72
|
+
"rules": [
|
|
73
|
+
{"kind": "secret", "decision": "block"},
|
|
74
|
+
{"kind": "card", "decision": "block", "direction": "outbound"},
|
|
75
|
+
{"kind": "email", "decision": "allow", "tools": ["crm.*"]},
|
|
76
|
+
],
|
|
77
|
+
"denyTools": ["shell.*"],
|
|
78
|
+
})
|
|
79
|
+
|
|
80
|
+
sai.allow_value("@ourcompany.com") # stop flagging your own domain
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
When an action contains several findings, the **most severe** decision wins.
|
|
84
|
+
|
|
85
|
+
## The trail
|
|
86
|
+
|
|
87
|
+
```python
|
|
88
|
+
sai.summary() # counts over a recent window
|
|
89
|
+
sai.audit(limit=100) # the records themselves
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
Records hold the **kind and location** of what was found — `card` at
|
|
93
|
+
`body.payment.number` — and never the value. `Finding` has no field for one.
|
|
94
|
+
|
|
95
|
+
## When Secure AI is unreachable
|
|
96
|
+
|
|
97
|
+
Default is `on_unreachable="closed"`: the action does not happen. Correct for a
|
|
98
|
+
security control, and it does mean an outage here stops agents.
|
|
99
|
+
|
|
100
|
+
```python
|
|
101
|
+
SecureAI(api_key=..., on_unreachable="open") # availability over control
|
|
102
|
+
```
|
|
103
|
+
|
|
104
|
+
Failing open never waves a real refusal through — a 401, 402 or quota error is
|
|
105
|
+
an answer, not an outage, and still raises.
|
|
106
|
+
|
|
107
|
+
## No dependencies
|
|
108
|
+
|
|
109
|
+
Standard library only. This runs beside your model client and your framework,
|
|
110
|
+
and every dependency it adds is a version conflict it can cause.
|
|
111
|
+
|
|
112
|
+
Full reference: <https://secureai.one/developers>
|
|
@@ -0,0 +1,99 @@
|
|
|
1
|
+
# secure-ai
|
|
2
|
+
|
|
3
|
+
Data loss prevention for AI agents. Put a check in front of every action an
|
|
4
|
+
agent takes, before it takes it.
|
|
5
|
+
|
|
6
|
+
```bash
|
|
7
|
+
pip install secure-ai-guard
|
|
8
|
+
```
|
|
9
|
+
|
|
10
|
+
> **Status:** written, not yet executed — there was no Python interpreter on
|
|
11
|
+
> the machine this was built on. Run `pytest` before trusting it. The
|
|
12
|
+
> TypeScript SDK it mirrors does pass its suite.
|
|
13
|
+
|
|
14
|
+
## The one thing worth knowing
|
|
15
|
+
|
|
16
|
+
`guard` wraps a function your agent already calls, and calls it with the
|
|
17
|
+
**rewritten** arguments:
|
|
18
|
+
|
|
19
|
+
```python
|
|
20
|
+
from secure_ai import SecureAI
|
|
21
|
+
|
|
22
|
+
sai = SecureAI(api_key=os.environ["SECURE_AI_KEY"], agent="support-bot")
|
|
23
|
+
|
|
24
|
+
send_email = sai.guard("email.send", mailer.send)
|
|
25
|
+
send_email({"to": "ana@clientfirm.com", "body": "About invoice 4471…"})
|
|
26
|
+
```
|
|
27
|
+
|
|
28
|
+
The agent wrote a real address. `mailer.send` receives a stand-in. Nothing in
|
|
29
|
+
the agent's code had to check a decision — which is the point. Protection that
|
|
30
|
+
depends on remembering to check it ends at the fourteenth call site.
|
|
31
|
+
|
|
32
|
+
If the policy refuses, the wrapped function is never called and `ActionBlocked`
|
|
33
|
+
is raised.
|
|
34
|
+
|
|
35
|
+
As a decorator:
|
|
36
|
+
|
|
37
|
+
```python
|
|
38
|
+
@sai.guarded("email.send")
|
|
39
|
+
def send(payload): ...
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
## A whole toolbelt at once
|
|
43
|
+
|
|
44
|
+
```python
|
|
45
|
+
tools = sai.guard_tools(agent.tools)
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
Wraps whichever attribute holds the callable — `func`, `_run`, `run` — so
|
|
49
|
+
LangChain and similar frameworks are governed in one line. Shapes it does not
|
|
50
|
+
recognise are returned untouched rather than raising.
|
|
51
|
+
|
|
52
|
+
## Rules
|
|
53
|
+
|
|
54
|
+
Rules live on the account, not in the request. An agent cannot argue with them.
|
|
55
|
+
|
|
56
|
+
```python
|
|
57
|
+
sai.set_policy({
|
|
58
|
+
"fallback": "redact",
|
|
59
|
+
"rules": [
|
|
60
|
+
{"kind": "secret", "decision": "block"},
|
|
61
|
+
{"kind": "card", "decision": "block", "direction": "outbound"},
|
|
62
|
+
{"kind": "email", "decision": "allow", "tools": ["crm.*"]},
|
|
63
|
+
],
|
|
64
|
+
"denyTools": ["shell.*"],
|
|
65
|
+
})
|
|
66
|
+
|
|
67
|
+
sai.allow_value("@ourcompany.com") # stop flagging your own domain
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
When an action contains several findings, the **most severe** decision wins.
|
|
71
|
+
|
|
72
|
+
## The trail
|
|
73
|
+
|
|
74
|
+
```python
|
|
75
|
+
sai.summary() # counts over a recent window
|
|
76
|
+
sai.audit(limit=100) # the records themselves
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
Records hold the **kind and location** of what was found — `card` at
|
|
80
|
+
`body.payment.number` — and never the value. `Finding` has no field for one.
|
|
81
|
+
|
|
82
|
+
## When Secure AI is unreachable
|
|
83
|
+
|
|
84
|
+
Default is `on_unreachable="closed"`: the action does not happen. Correct for a
|
|
85
|
+
security control, and it does mean an outage here stops agents.
|
|
86
|
+
|
|
87
|
+
```python
|
|
88
|
+
SecureAI(api_key=..., on_unreachable="open") # availability over control
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
Failing open never waves a real refusal through — a 401, 402 or quota error is
|
|
92
|
+
an answer, not an outage, and still raises.
|
|
93
|
+
|
|
94
|
+
## No dependencies
|
|
95
|
+
|
|
96
|
+
Standard library only. This runs beside your model client and your framework,
|
|
97
|
+
and every dependency it adds is a version conflict it can cause.
|
|
98
|
+
|
|
99
|
+
Full reference: <https://secureai.one/developers>
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "secure-ai-guard"
|
|
7
|
+
version = "0.2.0"
|
|
8
|
+
description = "Data loss prevention for AI agents. Inspect what an agent is about to do, before it does it."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
keywords = ["dlp", "ai-agents", "security", "redaction", "audit", "mcp"]
|
|
13
|
+
# No dependencies, deliberately. This runs inside somebody's agent beside
|
|
14
|
+
# their model client and framework; every dependency it adds is a version
|
|
15
|
+
# conflict it can cause in a process already carrying too many.
|
|
16
|
+
dependencies = []
|
|
17
|
+
|
|
18
|
+
[project.urls]
|
|
19
|
+
Homepage = "https://secureai.one/developers"
|
|
20
|
+
|
|
21
|
+
[project.optional-dependencies]
|
|
22
|
+
dev = ["pytest>=7"]
|
|
23
|
+
|
|
24
|
+
[tool.hatch.build.targets.wheel]
|
|
25
|
+
packages = ["secure_ai"]
|
|
26
|
+
|
|
27
|
+
[tool.pytest.ini_options]
|
|
28
|
+
testpaths = ["tests"]
|
|
@@ -0,0 +1,512 @@
|
|
|
1
|
+
"""Secure AI — data loss prevention for AI agents.
|
|
2
|
+
|
|
3
|
+
The API is plain HTTP and anybody can call it with ``requests``. This exists
|
|
4
|
+
because the shape that makes the product work is not "call an endpoint", it is
|
|
5
|
+
"put a check in front of every action", and the difference between those two is
|
|
6
|
+
whether somebody remembers to do it at the fourteenth call site.
|
|
7
|
+
|
|
8
|
+
So the centre of this package is :meth:`SecureAI.guard`, which wraps a function
|
|
9
|
+
an agent already calls and invokes it with the *rewritten* arguments::
|
|
10
|
+
|
|
11
|
+
send = sai.guard("email.send", raw_send)
|
|
12
|
+
send({"to": "ana@clientfirm.com"}) # raw_send receives a stand-in
|
|
13
|
+
|
|
14
|
+
Python before TypeScript would have been the better order — LangChain, CrewAI,
|
|
15
|
+
LlamaIndex and most agent code is Python — and this is the correction.
|
|
16
|
+
|
|
17
|
+
Why no dependencies
|
|
18
|
+
-------------------
|
|
19
|
+
|
|
20
|
+
This runs inside somebody's agent, beside their model client, their framework
|
|
21
|
+
and their vendor SDKs. Every dependency it adds is a version conflict it can
|
|
22
|
+
cause in a process already carrying too many, and a security tool that is
|
|
23
|
+
awkward to install is one that gets removed. So: ``urllib`` from the standard
|
|
24
|
+
library, nothing else. It is slower per call than ``httpx`` and that is not the
|
|
25
|
+
constraint — the constraint is that ``pip install secure-ai`` never fails.
|
|
26
|
+
"""
|
|
27
|
+
|
|
28
|
+
from __future__ import annotations
|
|
29
|
+
|
|
30
|
+
import json
|
|
31
|
+
import time
|
|
32
|
+
import urllib.error
|
|
33
|
+
import urllib.parse
|
|
34
|
+
import urllib.request
|
|
35
|
+
from dataclasses import dataclass, field, replace
|
|
36
|
+
from typing import Any, Callable, Iterable, Literal, TypeVar
|
|
37
|
+
from urllib.parse import quote, urlencode
|
|
38
|
+
|
|
39
|
+
__all__ = [
|
|
40
|
+
"SecureAI",
|
|
41
|
+
"ActionBlocked",
|
|
42
|
+
"ApprovalRefused",
|
|
43
|
+
"SecureAIError",
|
|
44
|
+
"Approval",
|
|
45
|
+
"Finding",
|
|
46
|
+
"Inspection",
|
|
47
|
+
"KINDS",
|
|
48
|
+
]
|
|
49
|
+
|
|
50
|
+
__version__ = "0.2.0"
|
|
51
|
+
|
|
52
|
+
Decision = Literal["allow", "redact", "approve", "block"]
|
|
53
|
+
Direction = Literal["outbound", "inbound"]
|
|
54
|
+
ApprovalStatus = Literal["pending", "approved", "denied", "expired"]
|
|
55
|
+
|
|
56
|
+
#: Every kind the scanner reports, in the order a policy editor should list
|
|
57
|
+
#: them: the ones that end careers first. Mirrors KINDS in the Worker.
|
|
58
|
+
KINDS = (
|
|
59
|
+
"secret", "card", "iban", "ssn", "govid",
|
|
60
|
+
"email", "phone", "address", "postcode", "name", "host",
|
|
61
|
+
)
|
|
62
|
+
|
|
63
|
+
T = TypeVar("T")
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
@dataclass(frozen=True)
|
|
67
|
+
class Finding:
|
|
68
|
+
"""Something the scanner located: what it was and where it sat.
|
|
69
|
+
|
|
70
|
+
Deliberately no ``value``. The trail never carries one and neither does
|
|
71
|
+
this, so a caller logging a finding cannot accidentally log the thing the
|
|
72
|
+
product exists to protect.
|
|
73
|
+
"""
|
|
74
|
+
|
|
75
|
+
kind: str
|
|
76
|
+
path: str
|
|
77
|
+
decision: Decision
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
@dataclass(frozen=True)
|
|
81
|
+
class Inspection:
|
|
82
|
+
decision: Decision
|
|
83
|
+
#: The action, ready to send. ``None`` when blocked — a refusal carries
|
|
84
|
+
#: nothing sendable, so a caller cannot reach past the decision.
|
|
85
|
+
input: Any
|
|
86
|
+
#: ``{stand_in: real}``. Keep it: it is the only way back, and it is not
|
|
87
|
+
#: stored on our side.
|
|
88
|
+
map: dict[str, str]
|
|
89
|
+
findings: list[Finding]
|
|
90
|
+
tool_denied: bool
|
|
91
|
+
policy_source: str
|
|
92
|
+
audit_id: str
|
|
93
|
+
#: Set when the decision is "approve": the id to come back with once a
|
|
94
|
+
#: person has decided.
|
|
95
|
+
approval_id: str | None = None
|
|
96
|
+
#: When waiting stops being worth it. Milliseconds since epoch.
|
|
97
|
+
expires_at: int | None = None
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
def _sendable(tool: str, verdict: "Inspection", original: Any) -> Any:
|
|
101
|
+
"""What to hand the wrapped function.
|
|
102
|
+
|
|
103
|
+
The API omits ``input`` only on a block, so on any other decision it is
|
|
104
|
+
there. The question is what to do if it is not, and the answer is not
|
|
105
|
+
"send what the caller had".
|
|
106
|
+
|
|
107
|
+
On a redact the caller's input is the one thing that must not go: it still
|
|
108
|
+
holds the values the decision just said to replace. Falling back to it
|
|
109
|
+
turns a missing field into the library doing the exact opposite of its
|
|
110
|
+
purpose, silently, with the trail recording a redaction that did not
|
|
111
|
+
happen.
|
|
112
|
+
|
|
113
|
+
On an allow nothing was rewritten, so the caller's own input is the right
|
|
114
|
+
thing to pass and the fallback belongs there.
|
|
115
|
+
"""
|
|
116
|
+
if verdict.input is not None:
|
|
117
|
+
return verdict.input
|
|
118
|
+
if verdict.decision == "redact":
|
|
119
|
+
raise SecureAIError(
|
|
120
|
+
f"Secure AI decided to redact {tool} but returned nothing to send. "
|
|
121
|
+
"The original was not sent: it still holds the values that decision was about.",
|
|
122
|
+
502,
|
|
123
|
+
"missing_rewritten_input",
|
|
124
|
+
)
|
|
125
|
+
return original
|
|
126
|
+
|
|
127
|
+
|
|
128
|
+
def _approval(raw: dict[str, Any]) -> "Approval":
|
|
129
|
+
return Approval(
|
|
130
|
+
id=raw.get("id", ""),
|
|
131
|
+
created_at=int(raw.get("createdAt", 0)),
|
|
132
|
+
expires_at=int(raw.get("expiresAt", 0)),
|
|
133
|
+
status=raw.get("status", "pending"),
|
|
134
|
+
agent=raw.get("agent"),
|
|
135
|
+
tool=raw.get("tool", ""),
|
|
136
|
+
key_id=raw.get("keyId", ""),
|
|
137
|
+
findings=[Finding(f["kind"], f["path"], f["decision"]) for f in raw.get("findings", [])],
|
|
138
|
+
decided_by=raw.get("decidedBy"),
|
|
139
|
+
decided_at=raw.get("decidedAt"),
|
|
140
|
+
note=raw.get("note"),
|
|
141
|
+
)
|
|
142
|
+
|
|
143
|
+
|
|
144
|
+
@dataclass(frozen=True)
|
|
145
|
+
class Approval:
|
|
146
|
+
"""An action a rule stopped and handed to a person.
|
|
147
|
+
|
|
148
|
+
Holds the shape of what the agent wanted to do — the tool and the kinds it
|
|
149
|
+
carried — and never the payload, for the same reason the trail does not.
|
|
150
|
+
"""
|
|
151
|
+
|
|
152
|
+
id: str
|
|
153
|
+
created_at: int
|
|
154
|
+
expires_at: int
|
|
155
|
+
status: ApprovalStatus
|
|
156
|
+
agent: str | None
|
|
157
|
+
tool: str
|
|
158
|
+
key_id: str
|
|
159
|
+
findings: list[Finding]
|
|
160
|
+
decided_by: str | None = None
|
|
161
|
+
decided_at: int | None = None
|
|
162
|
+
note: str | None = None
|
|
163
|
+
|
|
164
|
+
|
|
165
|
+
class SecureAIError(RuntimeError):
|
|
166
|
+
"""The API answered, and said no.
|
|
167
|
+
|
|
168
|
+
Distinct from :class:`ActionBlocked`: this is a problem with the call — no
|
|
169
|
+
key, no subscription, over quota — not a decision about the action.
|
|
170
|
+
"""
|
|
171
|
+
|
|
172
|
+
def __init__(self, message: str, status: int, code: str | None) -> None:
|
|
173
|
+
super().__init__(message)
|
|
174
|
+
self.status = status
|
|
175
|
+
self.code = code
|
|
176
|
+
|
|
177
|
+
|
|
178
|
+
class ActionBlocked(Exception):
|
|
179
|
+
"""The policy refused this action.
|
|
180
|
+
|
|
181
|
+
Raised rather than returned so a caller cannot ignore it by not reading a
|
|
182
|
+
field. A guarded function that returned a verdict would be protection you
|
|
183
|
+
have to remember to check.
|
|
184
|
+
"""
|
|
185
|
+
|
|
186
|
+
def __init__(self, tool: str, result: Inspection) -> None:
|
|
187
|
+
if result.tool_denied:
|
|
188
|
+
what = "the tool itself is not permitted"
|
|
189
|
+
else:
|
|
190
|
+
blocked = [f for f in result.findings if f.decision == "block"]
|
|
191
|
+
what = ", ".join(f"{f.kind} at {f.path or 'the input'}" for f in blocked) or "policy"
|
|
192
|
+
super().__init__(f"Secure AI refused {tool}: {what}.")
|
|
193
|
+
self.tool = tool
|
|
194
|
+
self.findings = result.findings
|
|
195
|
+
self.tool_denied = result.tool_denied
|
|
196
|
+
self.audit_id = result.audit_id
|
|
197
|
+
|
|
198
|
+
|
|
199
|
+
class ApprovalRefused(Exception):
|
|
200
|
+
"""A person was asked, and the action did not go.
|
|
201
|
+
|
|
202
|
+
Three ways to arrive here and they are not the same: denied means somebody
|
|
203
|
+
looked and said no; expired means nobody looked in time, which is also a
|
|
204
|
+
no because an approval that runs out is a refusal rather than a release;
|
|
205
|
+
pending means the caller asked not to wait.
|
|
206
|
+
"""
|
|
207
|
+
|
|
208
|
+
def __init__(
|
|
209
|
+
self,
|
|
210
|
+
tool: str,
|
|
211
|
+
approval_id: str,
|
|
212
|
+
status: ApprovalStatus,
|
|
213
|
+
note: str | None = None,
|
|
214
|
+
) -> None:
|
|
215
|
+
if status == "pending":
|
|
216
|
+
what = "and it is still waiting for a person"
|
|
217
|
+
elif status == "expired":
|
|
218
|
+
what = "and nobody answered before it expired"
|
|
219
|
+
else:
|
|
220
|
+
what = "and it was refused" + (": " + note if note else "")
|
|
221
|
+
super().__init__(f"Secure AI held {tool} for approval {what}.")
|
|
222
|
+
self.tool = tool
|
|
223
|
+
self.approval_id = approval_id
|
|
224
|
+
self.status = status
|
|
225
|
+
self.note = note
|
|
226
|
+
|
|
227
|
+
|
|
228
|
+
@dataclass
|
|
229
|
+
class SecureAI:
|
|
230
|
+
"""A client.
|
|
231
|
+
|
|
232
|
+
:param api_key: from Settings → Developer.
|
|
233
|
+
:param agent: names this agent in the trail so its actions group together.
|
|
234
|
+
Worth setting — a trail where everything is unnamed is a trail nobody
|
|
235
|
+
can ask a question of.
|
|
236
|
+
:param on_unreachable: what to do when Secure AI itself cannot be reached.
|
|
237
|
+
``"closed"`` (the default) raises, so an action is not taken while the
|
|
238
|
+
thing governing it is down. That is correct for a security control and
|
|
239
|
+
it does mean an outage here stops agents. ``"open"`` lets the action
|
|
240
|
+
through unchecked; it exists because some workloads genuinely prefer
|
|
241
|
+
availability, and because somebody who wants it will otherwise write a
|
|
242
|
+
``try/except`` that also swallows real refusals.
|
|
243
|
+
"""
|
|
244
|
+
|
|
245
|
+
api_key: str
|
|
246
|
+
base_url: str = "https://api.secureai.one"
|
|
247
|
+
agent: str | None = None
|
|
248
|
+
timeout: float = 5.0
|
|
249
|
+
on_unreachable: Literal["closed", "open"] = "closed"
|
|
250
|
+
_opener: Any = field(default=None, repr=False, compare=False)
|
|
251
|
+
|
|
252
|
+
def __post_init__(self) -> None:
|
|
253
|
+
if not self.api_key:
|
|
254
|
+
raise ValueError("SecureAI needs an api_key.")
|
|
255
|
+
self.base_url = self.base_url.rstrip("/")
|
|
256
|
+
|
|
257
|
+
# ── plumbing ────────────────────────────────────────────────────────────
|
|
258
|
+
|
|
259
|
+
def _request(self, method: str, path: str, body: Any = None) -> Any:
|
|
260
|
+
data = None if body is None else json.dumps(body).encode("utf-8")
|
|
261
|
+
req = urllib.request.Request(
|
|
262
|
+
f"{self.base_url}{path}",
|
|
263
|
+
data=data,
|
|
264
|
+
method=method,
|
|
265
|
+
headers={
|
|
266
|
+
"Authorization": f"Bearer {self.api_key}",
|
|
267
|
+
**({"Content-Type": "application/json"} if data is not None else {}),
|
|
268
|
+
},
|
|
269
|
+
)
|
|
270
|
+
opener = self._opener or urllib.request.urlopen
|
|
271
|
+
try:
|
|
272
|
+
with opener(req, timeout=self.timeout) as res: # type: ignore[operator]
|
|
273
|
+
raw = res.read().decode("utf-8")
|
|
274
|
+
except urllib.error.HTTPError as exc:
|
|
275
|
+
raw = exc.read().decode("utf-8", errors="replace")
|
|
276
|
+
message, code = f"Secure AI returned {exc.code}.", None
|
|
277
|
+
try:
|
|
278
|
+
err = (json.loads(raw) or {}).get("error") or {}
|
|
279
|
+
message = err.get("message") or message
|
|
280
|
+
code = err.get("code")
|
|
281
|
+
except json.JSONDecodeError:
|
|
282
|
+
# Something in front of the API answered, not the API.
|
|
283
|
+
pass
|
|
284
|
+
raise SecureAIError(message, exc.code, code) from None
|
|
285
|
+
return json.loads(raw) if raw else None
|
|
286
|
+
|
|
287
|
+
# ── the API ─────────────────────────────────────────────────────────────
|
|
288
|
+
|
|
289
|
+
def inspect(
|
|
290
|
+
self,
|
|
291
|
+
tool: str,
|
|
292
|
+
action_input: Any,
|
|
293
|
+
*,
|
|
294
|
+
direction: Direction = "outbound",
|
|
295
|
+
agent: str | None = None,
|
|
296
|
+
approval_id: str | None = None,
|
|
297
|
+
) -> Inspection:
|
|
298
|
+
"""Judge an action without taking it.
|
|
299
|
+
|
|
300
|
+
Pass ``approval_id`` to come back with a yes a person has given.
|
|
301
|
+
The server re-checks the shape of what is being sent against what was
|
|
302
|
+
approved, so one yes cannot be spent on a different action.
|
|
303
|
+
"""
|
|
304
|
+
payload: dict[str, Any] = {
|
|
305
|
+
"tool": tool,
|
|
306
|
+
"input": action_input,
|
|
307
|
+
"direction": direction,
|
|
308
|
+
"agent": agent or self.agent,
|
|
309
|
+
}
|
|
310
|
+
if approval_id:
|
|
311
|
+
payload["approvalId"] = approval_id
|
|
312
|
+
body = self._request("POST", "/v1/inspect", payload)
|
|
313
|
+
return Inspection(
|
|
314
|
+
decision=body["decision"],
|
|
315
|
+
input=body.get("input"),
|
|
316
|
+
map=body.get("map") or {},
|
|
317
|
+
findings=[Finding(f["kind"], f["path"], f["decision"]) for f in body.get("findings", [])],
|
|
318
|
+
tool_denied=bool(body.get("toolDenied")),
|
|
319
|
+
policy_source=body.get("policySource", "default"),
|
|
320
|
+
audit_id=body.get("auditId", ""),
|
|
321
|
+
approval_id=body.get("approvalId"),
|
|
322
|
+
expires_at=body.get("expiresAt"),
|
|
323
|
+
)
|
|
324
|
+
|
|
325
|
+
def get_policy(self) -> dict[str, Any]:
|
|
326
|
+
"""The rules in force."""
|
|
327
|
+
return self._request("GET", "/v1/policy")
|
|
328
|
+
|
|
329
|
+
def set_policy(self, policy: dict[str, Any]) -> dict[str, Any]:
|
|
330
|
+
"""Replace them. Refused whole if any rule is malformed, naming it."""
|
|
331
|
+
return self._request("PUT", "/v1/policy", {"policy": policy})
|
|
332
|
+
|
|
333
|
+
def allow_value(self, value: str) -> dict[str, Any]:
|
|
334
|
+
"""Stop flagging a value — your own domain, a shared mailbox.
|
|
335
|
+
|
|
336
|
+
Applies to every agent on the account from the next action onward. A
|
|
337
|
+
leading ``@`` exempts a whole domain.
|
|
338
|
+
"""
|
|
339
|
+
return self._request("POST", "/v1/policy/allow", {"value": value})
|
|
340
|
+
|
|
341
|
+
def audit(self, *, limit: int | None = None, cursor: str | None = None) -> dict[str, Any]:
|
|
342
|
+
"""What agents did, newest first. Kinds and locations, never values."""
|
|
343
|
+
query = []
|
|
344
|
+
if limit:
|
|
345
|
+
query.append(f"limit={limit}")
|
|
346
|
+
if cursor:
|
|
347
|
+
query.append(f"cursor={urllib.parse.quote(cursor)}")
|
|
348
|
+
suffix = f"?{'&'.join(query)}" if query else ""
|
|
349
|
+
return self._request("GET", f"/v1/audit{suffix}")
|
|
350
|
+
|
|
351
|
+
def summary(self, *, limit: int | None = None) -> dict[str, Any]:
|
|
352
|
+
"""The counts, over a recent window."""
|
|
353
|
+
suffix = f"?limit={limit}" if limit else ""
|
|
354
|
+
return self._request("GET", f"/v1/audit/summary{suffix}")
|
|
355
|
+
|
|
356
|
+
def restore(self, text: str, mapping: dict[str, str]) -> str:
|
|
357
|
+
"""Put real values back into a reply, using an inspection's map."""
|
|
358
|
+
return self._request("POST", "/v1/restore", {"text": text, "map": mapping})["text"]
|
|
359
|
+
|
|
360
|
+
# ── the point of the library ────────────────────────────────────────────
|
|
361
|
+
|
|
362
|
+
def approval(self, approval_id: str) -> Approval:
|
|
363
|
+
"""Read one back."""
|
|
364
|
+
body = self._request("GET", f"/v1/approvals/{quote(approval_id, safe='')}")
|
|
365
|
+
return _approval(body["approval"])
|
|
366
|
+
|
|
367
|
+
def approvals(
|
|
368
|
+
self,
|
|
369
|
+
*,
|
|
370
|
+
status: ApprovalStatus | None = None,
|
|
371
|
+
limit: int | None = None,
|
|
372
|
+
) -> list[Approval]:
|
|
373
|
+
"""The queue, newest first."""
|
|
374
|
+
query: dict[str, Any] = {}
|
|
375
|
+
if status:
|
|
376
|
+
query["status"] = status
|
|
377
|
+
if limit:
|
|
378
|
+
query["limit"] = limit
|
|
379
|
+
path = "/v1/approvals" + ("?" + urlencode(query) if query else "")
|
|
380
|
+
return [_approval(a) for a in self._request("GET", path).get("approvals", [])]
|
|
381
|
+
|
|
382
|
+
def wait_for_approval(self, approval_id: str, *, poll_seconds: float = 2.0) -> Approval:
|
|
383
|
+
"""Block until somebody decides, or until the window closes.
|
|
384
|
+
|
|
385
|
+
Stops at the approval's own expiry rather than running forever: the
|
|
386
|
+
server denies it at that point regardless, so a caller polling past it
|
|
387
|
+
is waiting for an answer that has already been given.
|
|
388
|
+
"""
|
|
389
|
+
while True:
|
|
390
|
+
current = self.approval(approval_id)
|
|
391
|
+
if current.status != "pending":
|
|
392
|
+
return current
|
|
393
|
+
if time.time() * 1000 >= current.expires_at:
|
|
394
|
+
return replace(current, status="expired")
|
|
395
|
+
time.sleep(poll_seconds)
|
|
396
|
+
|
|
397
|
+
def guard(
|
|
398
|
+
self,
|
|
399
|
+
tool: str,
|
|
400
|
+
fn: Callable[[Any], T],
|
|
401
|
+
*,
|
|
402
|
+
direction: Direction = "outbound",
|
|
403
|
+
agent: str | None = None,
|
|
404
|
+
wait_for_approval: bool = True,
|
|
405
|
+
poll_seconds: float = 2.0,
|
|
406
|
+
) -> Callable[[Any], T]:
|
|
407
|
+
"""Wrap a function so it cannot run unchecked.
|
|
408
|
+
|
|
409
|
+
The returned callable inspects, then calls ``fn`` with the **rewritten**
|
|
410
|
+
arguments — so an agent that never looks at a decision still cannot send
|
|
411
|
+
a real card number — and raises :class:`ActionBlocked` when the policy
|
|
412
|
+
refuses.
|
|
413
|
+
|
|
414
|
+
Calling ``fn`` with the redacted input rather than the caller's is the
|
|
415
|
+
whole mechanism. Returning a verdict for the caller to check would make
|
|
416
|
+
protection opt-in at every site, which is what this exists to stop.
|
|
417
|
+
|
|
418
|
+
An action held for a person waits by default. Pass
|
|
419
|
+
``wait_for_approval=False`` to get :class:`ApprovalRefused` straight
|
|
420
|
+
away and do the waiting yourself.
|
|
421
|
+
"""
|
|
422
|
+
|
|
423
|
+
def guarded(action_input: Any) -> T:
|
|
424
|
+
try:
|
|
425
|
+
verdict = self.inspect(tool, action_input, direction=direction, agent=agent)
|
|
426
|
+
except SecureAIError:
|
|
427
|
+
# A refusal by the API is a real answer and must not be
|
|
428
|
+
# mistaken for an outage, whatever on_unreachable says.
|
|
429
|
+
raise
|
|
430
|
+
except Exception:
|
|
431
|
+
if self.on_unreachable == "open":
|
|
432
|
+
return fn(action_input)
|
|
433
|
+
raise
|
|
434
|
+
|
|
435
|
+
if verdict.decision == "block":
|
|
436
|
+
raise ActionBlocked(tool, verdict)
|
|
437
|
+
|
|
438
|
+
# Held for a person.
|
|
439
|
+
#
|
|
440
|
+
# This branch did not exist. "approve" did not match "block", so
|
|
441
|
+
# the action was sent immediately: on Python a policy saying an
|
|
442
|
+
# action must wait for a human was not weakened but bypassed, and
|
|
443
|
+
# the agent was told the check had passed.
|
|
444
|
+
if verdict.decision == "approve":
|
|
445
|
+
approval_id = verdict.approval_id or ""
|
|
446
|
+
if not approval_id:
|
|
447
|
+
raise ActionBlocked(tool, verdict)
|
|
448
|
+
if not wait_for_approval:
|
|
449
|
+
raise ApprovalRefused(tool, approval_id, "pending")
|
|
450
|
+
decided = self.wait_for_approval(approval_id, poll_seconds=poll_seconds)
|
|
451
|
+
if decided.status != "approved":
|
|
452
|
+
raise ApprovalRefused(tool, approval_id, decided.status, decided.note)
|
|
453
|
+
# Back with the id. The server re-checks the shape of what is
|
|
454
|
+
# being sent, so a yes cannot be spent on a different action.
|
|
455
|
+
after = self.inspect(
|
|
456
|
+
tool,
|
|
457
|
+
action_input,
|
|
458
|
+
direction=direction,
|
|
459
|
+
agent=agent,
|
|
460
|
+
approval_id=approval_id,
|
|
461
|
+
)
|
|
462
|
+
if after.decision == "block":
|
|
463
|
+
raise ActionBlocked(tool, after)
|
|
464
|
+
return fn(_sendable(tool, after, action_input))
|
|
465
|
+
|
|
466
|
+
return fn(_sendable(tool, verdict, action_input))
|
|
467
|
+
|
|
468
|
+
guarded.__name__ = getattr(fn, "__name__", "guarded")
|
|
469
|
+
guarded.__doc__ = getattr(fn, "__doc__", None)
|
|
470
|
+
return guarded
|
|
471
|
+
|
|
472
|
+
def guarded(self, tool: str, **kwargs: Any) -> Callable[[Callable[[Any], T]], Callable[[Any], T]]:
|
|
473
|
+
""":meth:`guard` as a decorator.
|
|
474
|
+
|
|
475
|
+
::
|
|
476
|
+
|
|
477
|
+
@sai.guarded("email.send")
|
|
478
|
+
def send(payload): ...
|
|
479
|
+
"""
|
|
480
|
+
|
|
481
|
+
def decorate(fn: Callable[[Any], T]) -> Callable[[Any], T]:
|
|
482
|
+
return self.guard(tool, fn, **kwargs)
|
|
483
|
+
|
|
484
|
+
return decorate
|
|
485
|
+
|
|
486
|
+
def guard_tools(self, tools: Iterable[Any], *, name_attr: str = "name") -> list[Any]:
|
|
487
|
+
"""Wrap a list of framework tool objects in place.
|
|
488
|
+
|
|
489
|
+
Agent frameworks hand around objects with a ``name`` and a callable —
|
|
490
|
+
LangChain's ``StructuredTool``, an OpenAI function spec, a plain
|
|
491
|
+
dataclass. This wraps whichever attribute holds the callable, so a
|
|
492
|
+
whole toolbelt is governed in one line rather than tool by tool.
|
|
493
|
+
|
|
494
|
+
Unknown shapes are returned untouched rather than raising: a helper
|
|
495
|
+
that refuses to start because one tool in a list is unfamiliar is a
|
|
496
|
+
helper nobody uses.
|
|
497
|
+
"""
|
|
498
|
+
out = []
|
|
499
|
+
for tool in tools:
|
|
500
|
+
name = getattr(tool, name_attr, None) or getattr(tool, "__name__", None)
|
|
501
|
+
for attr in ("func", "_run", "run", "fn", "callable"):
|
|
502
|
+
target = getattr(tool, attr, None)
|
|
503
|
+
if callable(target) and name:
|
|
504
|
+
try:
|
|
505
|
+
setattr(tool, attr, self.guard(str(name), target))
|
|
506
|
+
except (AttributeError, TypeError):
|
|
507
|
+
# Frozen or slotted objects cannot be patched. Left
|
|
508
|
+
# alone rather than failing the whole call.
|
|
509
|
+
pass
|
|
510
|
+
break
|
|
511
|
+
out.append(tool)
|
|
512
|
+
return out
|
|
@@ -0,0 +1,445 @@
|
|
|
1
|
+
"""The Python SDK.
|
|
2
|
+
|
|
3
|
+
NOT YET EXECUTED. There is no Python interpreter on the machine this was
|
|
4
|
+
written on, so every test below is unrun. They are written to the same shape as
|
|
5
|
+
the TypeScript suite — which does pass — and the behaviours they pin are the
|
|
6
|
+
ones that suite already proves at the API boundary, so the risk is in this
|
|
7
|
+
file's own syntax rather than in what it asserts. Run ``pytest`` once before
|
|
8
|
+
trusting any of it, and treat a first-run failure as a bug in the SDK, not as
|
|
9
|
+
a surprise.
|
|
10
|
+
|
|
11
|
+
The behaviour worth most of these: ``guard`` calls the wrapped function with
|
|
12
|
+
the *rewritten* input, not the caller's. That is what makes protection
|
|
13
|
+
automatic rather than opt-in, and a refactor could quietly undo it — every test
|
|
14
|
+
here would still pass if guard merely returned a verdict, unless one of them
|
|
15
|
+
checks what the underlying function actually received.
|
|
16
|
+
"""
|
|
17
|
+
|
|
18
|
+
from __future__ import annotations
|
|
19
|
+
|
|
20
|
+
import io
|
|
21
|
+
import json
|
|
22
|
+
import time
|
|
23
|
+
import urllib.error
|
|
24
|
+
|
|
25
|
+
import pytest
|
|
26
|
+
|
|
27
|
+
from secure_ai import (
|
|
28
|
+
ActionBlocked,
|
|
29
|
+
ApprovalRefused,
|
|
30
|
+
Finding,
|
|
31
|
+
SecureAI,
|
|
32
|
+
SecureAIError,
|
|
33
|
+
)
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
class FakeResponse:
|
|
37
|
+
def __init__(self, body: object, status: int = 200) -> None:
|
|
38
|
+
self._raw = json.dumps(body).encode("utf-8")
|
|
39
|
+
self.status = status
|
|
40
|
+
|
|
41
|
+
def read(self) -> bytes:
|
|
42
|
+
return self._raw
|
|
43
|
+
|
|
44
|
+
def __enter__(self) -> "FakeResponse":
|
|
45
|
+
return self
|
|
46
|
+
|
|
47
|
+
def __exit__(self, *_: object) -> None:
|
|
48
|
+
return None
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def client(handler, **kwargs):
|
|
52
|
+
"""A client whose opener records calls and returns what `handler` says."""
|
|
53
|
+
calls: list[dict] = []
|
|
54
|
+
|
|
55
|
+
def opener(req, timeout=None): # noqa: ARG001 - signature matches urlopen
|
|
56
|
+
body = json.loads(req.data.decode("utf-8")) if req.data else None
|
|
57
|
+
calls.append({
|
|
58
|
+
"url": req.full_url,
|
|
59
|
+
"method": req.get_method(),
|
|
60
|
+
"body": body,
|
|
61
|
+
"headers": {k.lower(): v for k, v in req.headers.items()},
|
|
62
|
+
})
|
|
63
|
+
result = handler(req)
|
|
64
|
+
if isinstance(result, urllib.error.HTTPError):
|
|
65
|
+
raise result
|
|
66
|
+
return result
|
|
67
|
+
|
|
68
|
+
sai = SecureAI(api_key="sai_test", base_url="https://api.test", _opener=opener, **kwargs)
|
|
69
|
+
return calls, sai
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def http_error(status: int, body: object) -> urllib.error.HTTPError:
|
|
73
|
+
payload = json.dumps(body).encode("utf-8")
|
|
74
|
+
return urllib.error.HTTPError(
|
|
75
|
+
"https://api.test", status, "err", {}, io.BytesIO(payload)
|
|
76
|
+
)
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
def allowed(value: object) -> FakeResponse:
|
|
80
|
+
return FakeResponse({
|
|
81
|
+
"decision": "allow", "input": value, "map": {}, "findings": [],
|
|
82
|
+
"toolDenied": False, "policySource": "default", "auditId": "a1",
|
|
83
|
+
})
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
def redacted(value: object, mapping: dict) -> FakeResponse:
|
|
87
|
+
return FakeResponse({
|
|
88
|
+
"decision": "redact", "input": value, "map": mapping, "findings": [],
|
|
89
|
+
"toolDenied": False, "policySource": "account", "auditId": "a2",
|
|
90
|
+
})
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
def blocked(findings: list | None = None, tool_denied: bool = False) -> FakeResponse:
|
|
94
|
+
return FakeResponse({
|
|
95
|
+
"decision": "block", "map": {},
|
|
96
|
+
"findings": findings if findings is not None else
|
|
97
|
+
[{"kind": "secret", "path": "headers.authorization", "decision": "block"}],
|
|
98
|
+
"toolDenied": tool_denied, "policySource": "account", "auditId": "a3",
|
|
99
|
+
})
|
|
100
|
+
|
|
101
|
+
|
|
102
|
+
class TestConstruction:
|
|
103
|
+
def test_insists_on_a_key(self) -> None:
|
|
104
|
+
with pytest.raises(ValueError, match="api_key"):
|
|
105
|
+
SecureAI(api_key="")
|
|
106
|
+
|
|
107
|
+
def test_trims_a_trailing_slash_rather_than_doubling_it(self) -> None:
|
|
108
|
+
calls, sai = client(lambda _r: allowed({}))
|
|
109
|
+
sai.base_url = "https://api.test/"
|
|
110
|
+
sai.__post_init__()
|
|
111
|
+
sai.inspect("t", {})
|
|
112
|
+
assert calls[0]["url"] == "https://api.test/v1/inspect"
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
class TestInspect:
|
|
116
|
+
def test_sends_the_action_with_outbound_as_the_default(self) -> None:
|
|
117
|
+
calls, sai = client(lambda _r: allowed({"a": 1}))
|
|
118
|
+
sai.inspect("http.post", {"a": 1})
|
|
119
|
+
assert calls[0]["body"]["tool"] == "http.post"
|
|
120
|
+
assert calls[0]["body"]["direction"] == "outbound"
|
|
121
|
+
|
|
122
|
+
def test_names_the_agent_from_the_client(self) -> None:
|
|
123
|
+
calls, sai = client(lambda _r: allowed({}), agent="nightly-sync")
|
|
124
|
+
sai.inspect("t", {})
|
|
125
|
+
assert calls[0]["body"]["agent"] == "nightly-sync"
|
|
126
|
+
|
|
127
|
+
def test_sends_the_key_as_a_bearer(self) -> None:
|
|
128
|
+
calls, sai = client(lambda _r: allowed({}))
|
|
129
|
+
sai.inspect("t", {})
|
|
130
|
+
assert calls[0]["headers"]["authorization"] == "Bearer sai_test"
|
|
131
|
+
|
|
132
|
+
def test_turns_a_refusal_into_an_error_carrying_the_code(self) -> None:
|
|
133
|
+
_calls, sai = client(lambda _r: http_error(402, {
|
|
134
|
+
"error": {"message": "No API access.", "code": "api_access_required"}
|
|
135
|
+
}))
|
|
136
|
+
with pytest.raises(SecureAIError) as caught:
|
|
137
|
+
sai.inspect("t", {})
|
|
138
|
+
assert caught.value.status == 402
|
|
139
|
+
assert caught.value.code == "api_access_required"
|
|
140
|
+
|
|
141
|
+
def test_survives_a_non_json_body_from_something_in_front_of_the_api(self) -> None:
|
|
142
|
+
def opener(req, timeout=None): # noqa: ARG001
|
|
143
|
+
raise urllib.error.HTTPError(req.full_url, 502, "bad", {}, io.BytesIO(b"<html>"))
|
|
144
|
+
|
|
145
|
+
sai = SecureAI(api_key="k", base_url="https://api.test", _opener=opener)
|
|
146
|
+
with pytest.raises(SecureAIError) as caught:
|
|
147
|
+
sai.inspect("t", {})
|
|
148
|
+
assert caught.value.status == 502
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
class TestGuard:
|
|
152
|
+
def test_calls_the_function_with_the_rewritten_input(self) -> None:
|
|
153
|
+
"""The whole mechanism.
|
|
154
|
+
|
|
155
|
+
If this passes while ``send`` receives the caller's original dict, the
|
|
156
|
+
library protects nothing.
|
|
157
|
+
"""
|
|
158
|
+
real = {"to": "ana@example.org"}
|
|
159
|
+
safe = {"to": "person1@example.com"}
|
|
160
|
+
_calls, sai = client(lambda _r: redacted(safe, {"person1@example.com": "ana@example.org"}))
|
|
161
|
+
|
|
162
|
+
seen: list[dict] = []
|
|
163
|
+
|
|
164
|
+
def send(payload: dict) -> str:
|
|
165
|
+
seen.append(payload)
|
|
166
|
+
return f"sent to {payload['to']}"
|
|
167
|
+
|
|
168
|
+
out = sai.guard("email.send", send)(real)
|
|
169
|
+
assert seen == [safe]
|
|
170
|
+
assert out == "sent to person1@example.com"
|
|
171
|
+
|
|
172
|
+
def test_raises_and_never_calls_the_function_when_refused(self) -> None:
|
|
173
|
+
_calls, sai = client(lambda _r: blocked())
|
|
174
|
+
called: list[object] = []
|
|
175
|
+
|
|
176
|
+
with pytest.raises(ActionBlocked):
|
|
177
|
+
sai.guard("http.post", lambda payload: called.append(payload))({})
|
|
178
|
+
assert called == []
|
|
179
|
+
|
|
180
|
+
def test_says_what_was_refused(self) -> None:
|
|
181
|
+
_calls, sai = client(lambda _r: blocked())
|
|
182
|
+
with pytest.raises(ActionBlocked, match=r"secret at headers\.authorization"):
|
|
183
|
+
sai.guard("http.post", lambda _p: "x")({})
|
|
184
|
+
|
|
185
|
+
def test_explains_a_denied_tool_differently(self) -> None:
|
|
186
|
+
_calls, sai = client(lambda _r: blocked(findings=[], tool_denied=True))
|
|
187
|
+
with pytest.raises(ActionBlocked, match="tool itself is not permitted"):
|
|
188
|
+
sai.guard("shell.exec", lambda _p: "x")({})
|
|
189
|
+
|
|
190
|
+
def test_passes_an_allowed_action_through(self) -> None:
|
|
191
|
+
payload = {"q": "orders shipped yesterday"}
|
|
192
|
+
_calls, sai = client(lambda _r: allowed(payload))
|
|
193
|
+
assert sai.guard("db.query", lambda p: p["q"])(payload) == "orders shipped yesterday"
|
|
194
|
+
|
|
195
|
+
def test_works_as_a_decorator(self) -> None:
|
|
196
|
+
_calls, sai = client(lambda _r: allowed({"a": 1}))
|
|
197
|
+
|
|
198
|
+
@sai.guarded("db.query")
|
|
199
|
+
def run(payload: dict) -> dict:
|
|
200
|
+
return payload
|
|
201
|
+
|
|
202
|
+
assert run({"a": 1}) == {"a": 1}
|
|
203
|
+
assert run.__name__ == "run"
|
|
204
|
+
|
|
205
|
+
|
|
206
|
+
class TestUnreachable:
|
|
207
|
+
@staticmethod
|
|
208
|
+
def _dead(req, timeout=None): # noqa: ARG004
|
|
209
|
+
raise OSError("network down")
|
|
210
|
+
|
|
211
|
+
def test_fails_closed_by_default(self) -> None:
|
|
212
|
+
sai = SecureAI(api_key="k", _opener=self._dead)
|
|
213
|
+
called: list[object] = []
|
|
214
|
+
with pytest.raises(OSError):
|
|
215
|
+
sai.guard("t", lambda p: called.append(p))({})
|
|
216
|
+
assert called == []
|
|
217
|
+
|
|
218
|
+
def test_fails_open_when_asked(self) -> None:
|
|
219
|
+
sai = SecureAI(api_key="k", _opener=self._dead, on_unreachable="open")
|
|
220
|
+
assert sai.guard("t", lambda _p: "done")({}) == "done"
|
|
221
|
+
|
|
222
|
+
def test_an_answered_refusal_is_not_an_outage(self) -> None:
|
|
223
|
+
"""Failing open must never wave a 401 through as a blinked network."""
|
|
224
|
+
_calls, sai = client(
|
|
225
|
+
lambda _r: http_error(401, {"error": {"message": "bad key", "code": "invalid_api_key"}}),
|
|
226
|
+
on_unreachable="open",
|
|
227
|
+
)
|
|
228
|
+
called: list[object] = []
|
|
229
|
+
with pytest.raises(SecureAIError):
|
|
230
|
+
sai.guard("t", lambda p: called.append(p))({})
|
|
231
|
+
assert called == []
|
|
232
|
+
|
|
233
|
+
|
|
234
|
+
class TestPolicyAndTrail:
|
|
235
|
+
def test_reads_and_replaces_the_policy(self) -> None:
|
|
236
|
+
calls, sai = client(lambda _r: FakeResponse({"policy": {"version": 1, "fallback": "block", "rules": []}}))
|
|
237
|
+
sai.set_policy({"fallback": "block", "rules": []})
|
|
238
|
+
assert calls[0]["method"] == "PUT"
|
|
239
|
+
assert calls[0]["body"] == {"policy": {"fallback": "block", "rules": []}}
|
|
240
|
+
|
|
241
|
+
def test_exempts_a_value(self) -> None:
|
|
242
|
+
calls, sai = client(lambda _r: FakeResponse({"policy": {}, "added": "@ours.com"}))
|
|
243
|
+
sai.allow_value("@ours.com")
|
|
244
|
+
assert calls[0]["url"] == "https://api.test/v1/policy/allow"
|
|
245
|
+
assert calls[0]["body"] == {"value": "@ours.com"}
|
|
246
|
+
|
|
247
|
+
def test_passes_paging_on_the_query_string(self) -> None:
|
|
248
|
+
calls, sai = client(lambda _r: FakeResponse({"events": [], "cursor": None}))
|
|
249
|
+
sai.audit(limit=10, cursor="abc")
|
|
250
|
+
assert calls[0]["url"] == "https://api.test/v1/audit?limit=10&cursor=abc"
|
|
251
|
+
|
|
252
|
+
def test_asks_for_the_trail_plainly_when_unpaged(self) -> None:
|
|
253
|
+
calls, sai = client(lambda _r: FakeResponse({"events": [], "cursor": None}))
|
|
254
|
+
sai.audit()
|
|
255
|
+
assert calls[0]["url"] == "https://api.test/v1/audit"
|
|
256
|
+
|
|
257
|
+
def test_restores_a_reply(self) -> None:
|
|
258
|
+
_calls, sai = client(lambda _r: FakeResponse({"text": "call ana@example.org"}))
|
|
259
|
+
assert sai.restore("call person1@example.com", {"person1@example.com": "ana@example.org"}) \
|
|
260
|
+
== "call ana@example.org"
|
|
261
|
+
|
|
262
|
+
|
|
263
|
+
class TestFindings:
|
|
264
|
+
def test_a_finding_has_no_place_to_put_a_value(self) -> None:
|
|
265
|
+
"""The trail carries kinds and locations. So does this."""
|
|
266
|
+
assert set(Finding.__dataclass_fields__) == {"kind", "path", "decision"}
|
|
267
|
+
|
|
268
|
+
|
|
269
|
+
class TestGuardTools:
|
|
270
|
+
def test_wraps_a_framework_tool_in_place(self) -> None:
|
|
271
|
+
_calls, sai = client(lambda _r: allowed({"a": 1}))
|
|
272
|
+
|
|
273
|
+
class Tool:
|
|
274
|
+
name = "db.query"
|
|
275
|
+
|
|
276
|
+
def __init__(self) -> None:
|
|
277
|
+
self.func = lambda payload: payload
|
|
278
|
+
|
|
279
|
+
tool = Tool()
|
|
280
|
+
sai.guard_tools([tool])
|
|
281
|
+
assert tool.func({"a": 1}) == {"a": 1}
|
|
282
|
+
|
|
283
|
+
def test_leaves_an_unfamiliar_shape_alone_rather_than_raising(self) -> None:
|
|
284
|
+
_calls, sai = client(lambda _r: allowed({}))
|
|
285
|
+
odd = object()
|
|
286
|
+
assert sai.guard_tools([odd]) == [odd]
|
|
287
|
+
|
|
288
|
+
|
|
289
|
+
class TestHeldForAPerson:
|
|
290
|
+
"""The branch that did not exist.
|
|
291
|
+
|
|
292
|
+
``guard`` handled ``block`` and nothing else, so an ``approve`` decision
|
|
293
|
+
fell through to the send. On Python a policy saying an action must wait
|
|
294
|
+
for a human was not weakened but bypassed, and the agent was told the
|
|
295
|
+
check had passed.
|
|
296
|
+
"""
|
|
297
|
+
|
|
298
|
+
@staticmethod
|
|
299
|
+
def _held(approval_id="ap1"):
|
|
300
|
+
return {
|
|
301
|
+
"decision": "approve",
|
|
302
|
+
"approvalId": approval_id,
|
|
303
|
+
"expiresAt": int(time.time() * 1000) + 600_000,
|
|
304
|
+
"map": {},
|
|
305
|
+
"findings": [{"kind": "card", "path": "body.note", "decision": "approve"}],
|
|
306
|
+
"toolDenied": False,
|
|
307
|
+
"policySource": "account",
|
|
308
|
+
"auditId": "a1",
|
|
309
|
+
}
|
|
310
|
+
|
|
311
|
+
@staticmethod
|
|
312
|
+
def _approval(status, **extra):
|
|
313
|
+
base = {
|
|
314
|
+
"id": "ap1",
|
|
315
|
+
"createdAt": int(time.time() * 1000) - 1000,
|
|
316
|
+
"expiresAt": int(time.time() * 1000) + 600_000,
|
|
317
|
+
"status": status,
|
|
318
|
+
"agent": "billing",
|
|
319
|
+
"tool": "mail.send",
|
|
320
|
+
"keyId": "k1",
|
|
321
|
+
"findings": [],
|
|
322
|
+
}
|
|
323
|
+
base.update(extra)
|
|
324
|
+
return {"approval": base}
|
|
325
|
+
|
|
326
|
+
def test_waits_then_sends_once_a_person_says_yes(self):
|
|
327
|
+
seen = []
|
|
328
|
+
|
|
329
|
+
def handler(req):
|
|
330
|
+
if req.full_url.endswith("/v1/inspect"):
|
|
331
|
+
body = json.loads(req.data.decode("utf-8"))
|
|
332
|
+
if body.get("approvalId"):
|
|
333
|
+
return FakeResponse({
|
|
334
|
+
"decision": "allow",
|
|
335
|
+
"input": {"note": "go ahead"},
|
|
336
|
+
"map": {}, "findings": [], "toolDenied": False,
|
|
337
|
+
"policySource": "account", "auditId": "a2",
|
|
338
|
+
})
|
|
339
|
+
return FakeResponse(self._held())
|
|
340
|
+
return FakeResponse(self._approval("approved"))
|
|
341
|
+
|
|
342
|
+
_, sai = client(handler)
|
|
343
|
+
|
|
344
|
+
def send(payload):
|
|
345
|
+
seen.append(payload)
|
|
346
|
+
return "sent"
|
|
347
|
+
|
|
348
|
+
assert sai.guard("mail.send", send)({"note": "original"}) == "sent"
|
|
349
|
+
# It went, and it went with what came back from the second check.
|
|
350
|
+
assert seen == [{"note": "go ahead"}]
|
|
351
|
+
|
|
352
|
+
def test_does_not_send_when_the_person_says_no(self):
|
|
353
|
+
def handler(req):
|
|
354
|
+
if req.full_url.endswith("/v1/inspect"):
|
|
355
|
+
return FakeResponse(self._held())
|
|
356
|
+
return FakeResponse(self._approval("denied", note="not this customer"))
|
|
357
|
+
|
|
358
|
+
_, sai = client(handler)
|
|
359
|
+
sent = []
|
|
360
|
+
|
|
361
|
+
with pytest.raises(ApprovalRefused) as caught:
|
|
362
|
+
sai.guard("mail.send", lambda p: sent.append(p))({"note": "x"})
|
|
363
|
+
assert caught.value.status == "denied"
|
|
364
|
+
assert "not this customer" in str(caught.value)
|
|
365
|
+
assert sent == []
|
|
366
|
+
|
|
367
|
+
def test_an_expiry_is_a_refusal_not_a_release(self):
|
|
368
|
+
def handler(req):
|
|
369
|
+
if req.full_url.endswith("/v1/inspect"):
|
|
370
|
+
return FakeResponse(self._held())
|
|
371
|
+
return FakeResponse(self._approval("expired"))
|
|
372
|
+
|
|
373
|
+
_, sai = client(handler)
|
|
374
|
+
sent = []
|
|
375
|
+
with pytest.raises(ApprovalRefused) as caught:
|
|
376
|
+
sai.guard("mail.send", lambda p: sent.append(p))({"note": "x"})
|
|
377
|
+
assert caught.value.status == "expired"
|
|
378
|
+
assert sent == []
|
|
379
|
+
|
|
380
|
+
def test_can_refuse_to_wait_and_hand_the_id_back(self):
|
|
381
|
+
_, sai = client(lambda req: FakeResponse(self._held()))
|
|
382
|
+
sent = []
|
|
383
|
+
with pytest.raises(ApprovalRefused) as caught:
|
|
384
|
+
sai.guard("mail.send", lambda p: sent.append(p), wait_for_approval=False)({"a": 1})
|
|
385
|
+
assert caught.value.status == "pending"
|
|
386
|
+
assert caught.value.approval_id == "ap1"
|
|
387
|
+
assert sent == []
|
|
388
|
+
|
|
389
|
+
def test_a_held_action_with_no_id_is_refused_rather_than_sent(self):
|
|
390
|
+
held = self._held()
|
|
391
|
+
del held["approvalId"]
|
|
392
|
+
_, sai = client(lambda req: FakeResponse(held))
|
|
393
|
+
sent = []
|
|
394
|
+
with pytest.raises(ActionBlocked):
|
|
395
|
+
sai.guard("mail.send", lambda p: sent.append(p))({"a": 1})
|
|
396
|
+
assert sent == []
|
|
397
|
+
|
|
398
|
+
def test_reads_the_queue(self):
|
|
399
|
+
def handler(req):
|
|
400
|
+
assert "status=pending" in req.full_url
|
|
401
|
+
return FakeResponse({"approvals": [self._approval("pending")["approval"]]})
|
|
402
|
+
|
|
403
|
+
_, sai = client(handler)
|
|
404
|
+
queue = sai.approvals(status="pending")
|
|
405
|
+
assert len(queue) == 1
|
|
406
|
+
assert queue[0].tool == "mail.send"
|
|
407
|
+
assert queue[0].status == "pending"
|
|
408
|
+
|
|
409
|
+
def test_escapes_the_id_in_the_path(self):
|
|
410
|
+
calls, sai = client(lambda req: FakeResponse(self._approval("approved")))
|
|
411
|
+
sai.approval("a/b?c")
|
|
412
|
+
assert "a%2Fb%3Fc" in calls[0]["url"]
|
|
413
|
+
|
|
414
|
+
|
|
415
|
+
class TestARedactWithNothingToSend:
|
|
416
|
+
"""The fallback that braced open.
|
|
417
|
+
|
|
418
|
+
The API omits ``input`` only on a block, so on a redact it is always
|
|
419
|
+
there — until some day it is not. ``guard`` fell back to the caller's own
|
|
420
|
+
input, which on a redact is the one thing that must not go: it still
|
|
421
|
+
holds the values the decision had just said to replace.
|
|
422
|
+
"""
|
|
423
|
+
|
|
424
|
+
BROKEN = {
|
|
425
|
+
"decision": "redact",
|
|
426
|
+
"map": {},
|
|
427
|
+
"findings": [],
|
|
428
|
+
"toolDenied": False,
|
|
429
|
+
"policySource": "account",
|
|
430
|
+
"auditId": "a9",
|
|
431
|
+
}
|
|
432
|
+
|
|
433
|
+
def test_raises_rather_than_sending_the_original(self):
|
|
434
|
+
_, sai = client(lambda req: FakeResponse(self.BROKEN))
|
|
435
|
+
sent = []
|
|
436
|
+
with pytest.raises(SecureAIError) as caught:
|
|
437
|
+
sai.guard("mail.send", lambda p: sent.append(p))({"to": "ana@clientfirm.com"})
|
|
438
|
+
assert "nothing to send" in str(caught.value)
|
|
439
|
+
assert sent == []
|
|
440
|
+
|
|
441
|
+
def test_an_allow_still_passes_the_original_through(self):
|
|
442
|
+
allowed = dict(self.BROKEN, decision="allow")
|
|
443
|
+
_, sai = client(lambda req: FakeResponse(allowed))
|
|
444
|
+
original = {"note": "nothing sensitive"}
|
|
445
|
+
assert sai.guard("mail.send", lambda p: p)(original) == original
|