iris-security-vertexai 0.1.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.
- iris_security_vertexai-0.1.0/PKG-INFO +16 -0
- iris_security_vertexai-0.1.0/iris_security_vertexai.egg-info/PKG-INFO +16 -0
- iris_security_vertexai-0.1.0/iris_security_vertexai.egg-info/SOURCES.txt +11 -0
- iris_security_vertexai-0.1.0/iris_security_vertexai.egg-info/dependency_links.txt +1 -0
- iris_security_vertexai-0.1.0/iris_security_vertexai.egg-info/requires.txt +9 -0
- iris_security_vertexai-0.1.0/iris_security_vertexai.egg-info/top_level.txt +1 -0
- iris_security_vertexai-0.1.0/iris_vertexai/__init__.py +29 -0
- iris_security_vertexai-0.1.0/iris_vertexai/client.py +314 -0
- iris_security_vertexai-0.1.0/iris_vertexai/fedramp.py +89 -0
- iris_security_vertexai-0.1.0/iris_vertexai/init.py +5 -0
- iris_security_vertexai-0.1.0/pyproject.toml +36 -0
- iris_security_vertexai-0.1.0/setup.cfg +4 -0
- iris_security_vertexai-0.1.0/tests/test_vertexai_integration.py +167 -0
|
@@ -0,0 +1,16 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: iris-security-vertexai
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: IRIS governance for Vertex AI via google-cloud-aiplatform
|
|
5
|
+
Author-email: IRIS Platform <sdk@iris.ai>
|
|
6
|
+
License: Apache-2.0
|
|
7
|
+
Project-URL: Homepage, https://github.com/gimartinb/iris-sdk
|
|
8
|
+
Project-URL: Repository, https://github.com/gimartinb/iris-sdk
|
|
9
|
+
Requires-Python: >=3.10
|
|
10
|
+
Requires-Dist: iris-security-core>=0.1.0
|
|
11
|
+
Requires-Dist: iris-security-sdk>=0.1.0
|
|
12
|
+
Provides-Extra: vertexai
|
|
13
|
+
Requires-Dist: google-cloud-aiplatform>=1.60; extra == "vertexai"
|
|
14
|
+
Provides-Extra: dev
|
|
15
|
+
Requires-Dist: pytest>=8.0; extra == "dev"
|
|
16
|
+
Requires-Dist: ruff>=0.4; extra == "dev"
|
|
@@ -0,0 +1,16 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: iris-security-vertexai
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: IRIS governance for Vertex AI via google-cloud-aiplatform
|
|
5
|
+
Author-email: IRIS Platform <sdk@iris.ai>
|
|
6
|
+
License: Apache-2.0
|
|
7
|
+
Project-URL: Homepage, https://github.com/gimartinb/iris-sdk
|
|
8
|
+
Project-URL: Repository, https://github.com/gimartinb/iris-sdk
|
|
9
|
+
Requires-Python: >=3.10
|
|
10
|
+
Requires-Dist: iris-security-core>=0.1.0
|
|
11
|
+
Requires-Dist: iris-security-sdk>=0.1.0
|
|
12
|
+
Provides-Extra: vertexai
|
|
13
|
+
Requires-Dist: google-cloud-aiplatform>=1.60; extra == "vertexai"
|
|
14
|
+
Provides-Extra: dev
|
|
15
|
+
Requires-Dist: pytest>=8.0; extra == "dev"
|
|
16
|
+
Requires-Dist: ruff>=0.4; extra == "dev"
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
pyproject.toml
|
|
2
|
+
iris_security_vertexai.egg-info/PKG-INFO
|
|
3
|
+
iris_security_vertexai.egg-info/SOURCES.txt
|
|
4
|
+
iris_security_vertexai.egg-info/dependency_links.txt
|
|
5
|
+
iris_security_vertexai.egg-info/requires.txt
|
|
6
|
+
iris_security_vertexai.egg-info/top_level.txt
|
|
7
|
+
iris_vertexai/__init__.py
|
|
8
|
+
iris_vertexai/client.py
|
|
9
|
+
iris_vertexai/fedramp.py
|
|
10
|
+
iris_vertexai/init.py
|
|
11
|
+
tests/test_vertexai_integration.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
iris_vertexai
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
"""
|
|
2
|
+
IRIS Vertex AI integration - governed wrapper for Vertex AI GenerativeModel.
|
|
3
|
+
"""
|
|
4
|
+
|
|
5
|
+
from __future__ import annotations
|
|
6
|
+
|
|
7
|
+
from iris import IrisViolationError
|
|
8
|
+
from iris_core.models.passport import (
|
|
9
|
+
AgentPassport,
|
|
10
|
+
ComplianceTag,
|
|
11
|
+
DataClassification,
|
|
12
|
+
Environment,
|
|
13
|
+
)
|
|
14
|
+
from iris_core.models.policy import Violation
|
|
15
|
+
|
|
16
|
+
from iris_vertexai.client import IrisGenerativeModel, IrisVertexAI
|
|
17
|
+
|
|
18
|
+
__version__ = "0.1.0"
|
|
19
|
+
|
|
20
|
+
__all__ = [
|
|
21
|
+
"IrisVertexAI",
|
|
22
|
+
"IrisGenerativeModel",
|
|
23
|
+
"IrisViolationError",
|
|
24
|
+
"AgentPassport",
|
|
25
|
+
"ComplianceTag",
|
|
26
|
+
"DataClassification",
|
|
27
|
+
"Environment",
|
|
28
|
+
"Violation",
|
|
29
|
+
]
|
|
@@ -0,0 +1,314 @@
|
|
|
1
|
+
"""Governed Vertex AI wrapper with IRIS policy enforcement."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import importlib
|
|
6
|
+
import json
|
|
7
|
+
import os
|
|
8
|
+
import uuid
|
|
9
|
+
from datetime import datetime
|
|
10
|
+
from typing import Any, Optional
|
|
11
|
+
|
|
12
|
+
from iris import IrisViolationError
|
|
13
|
+
from iris_core.dlp import DLPScanner
|
|
14
|
+
from iris_core.dlp.enforcement import (
|
|
15
|
+
enforce_prompt_dlp,
|
|
16
|
+
extract_gemini_response_text,
|
|
17
|
+
handle_response_dlp,
|
|
18
|
+
)
|
|
19
|
+
from iris_core.engine.cedar import CedarEngine, EvaluationContext
|
|
20
|
+
from iris_core.rbac.context import UserContext
|
|
21
|
+
from iris_core.evidence.vault import EvidenceVault
|
|
22
|
+
from iris_core.models.passport import AgentPassport
|
|
23
|
+
from iris_core.models.policy import PolicyResult, Severity, Violation
|
|
24
|
+
|
|
25
|
+
from iris_gemini.guardrails import _extract_text
|
|
26
|
+
|
|
27
|
+
from iris_vertexai.fedramp import check_fedramp_location, check_fedramp_model
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def _lazy_vertexai():
|
|
31
|
+
try:
|
|
32
|
+
return importlib.import_module("vertexai")
|
|
33
|
+
except ModuleNotFoundError as exc:
|
|
34
|
+
raise ImportError(
|
|
35
|
+
"google-cloud-aiplatform is required for IrisVertexAI. "
|
|
36
|
+
"Install with: pip install google-cloud-aiplatform"
|
|
37
|
+
) from exc
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def _lazy_vertexai_generative_models():
|
|
41
|
+
try:
|
|
42
|
+
return importlib.import_module("vertexai.generative_models")
|
|
43
|
+
except ModuleNotFoundError as exc:
|
|
44
|
+
raise ImportError(
|
|
45
|
+
"vertexai.generative_models is unavailable. "
|
|
46
|
+
"Install or upgrade: pip install google-cloud-aiplatform>=1.60"
|
|
47
|
+
) from exc
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
def _is_fedramp_enabled(passport: AgentPassport) -> bool:
|
|
51
|
+
for tag in passport.compliance_tags:
|
|
52
|
+
value = getattr(tag, "value", str(tag)).lower()
|
|
53
|
+
if value == "fedramp":
|
|
54
|
+
return True
|
|
55
|
+
return False
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def _load_passport_policy(engine: CedarEngine, passport: AgentPassport) -> None:
|
|
59
|
+
if not passport.policy_ref:
|
|
60
|
+
return
|
|
61
|
+
from pathlib import Path
|
|
62
|
+
|
|
63
|
+
policy_path = Path(passport.policy_ref)
|
|
64
|
+
if not policy_path.is_absolute():
|
|
65
|
+
policy_path = Path.cwd() / policy_path
|
|
66
|
+
if policy_path.exists():
|
|
67
|
+
engine.load_policy_file(passport.agent_id, policy_path)
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def _current_environment():
|
|
71
|
+
from iris_core.models.passport import Environment
|
|
72
|
+
|
|
73
|
+
return Environment(os.environ.get("IRIS_ENV", "dev"))
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def _enforce_result(result: PolicyResult) -> None:
|
|
77
|
+
if result.decision == "DENY":
|
|
78
|
+
raise IrisViolationError(result)
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
class IrisVertexAI:
|
|
82
|
+
"""Top-level Vertex AI client wrapper with governance-aware model creation."""
|
|
83
|
+
|
|
84
|
+
def __init__(
|
|
85
|
+
self,
|
|
86
|
+
passport: AgentPassport,
|
|
87
|
+
project: Optional[str] = None,
|
|
88
|
+
location: Optional[str] = None,
|
|
89
|
+
user_email: Optional[str] = None,
|
|
90
|
+
user_role: Optional[str] = None,
|
|
91
|
+
**kwargs: Any,
|
|
92
|
+
):
|
|
93
|
+
from iris_core.dev_trust import print_dev_trust_message
|
|
94
|
+
|
|
95
|
+
print_dev_trust_message()
|
|
96
|
+
evidence_vault_dir = kwargs.pop("evidence_vault_dir", None)
|
|
97
|
+
vertexai = _lazy_vertexai()
|
|
98
|
+
vertexai.init(project=project, location=location, **kwargs)
|
|
99
|
+
self._passport = passport
|
|
100
|
+
self._project = project
|
|
101
|
+
self._location = location
|
|
102
|
+
self._user_email = user_email
|
|
103
|
+
self._user_role = user_role
|
|
104
|
+
self._engine = CedarEngine()
|
|
105
|
+
self._vault = EvidenceVault(agent_id=passport.agent_id, vault_dir=evidence_vault_dir)
|
|
106
|
+
self._dlp = DLPScanner(passport)
|
|
107
|
+
_load_passport_policy(self._engine, passport)
|
|
108
|
+
|
|
109
|
+
def get_model(self, model_name: str) -> "IrisGenerativeModel":
|
|
110
|
+
return IrisGenerativeModel(
|
|
111
|
+
model_name=model_name,
|
|
112
|
+
passport=self._passport,
|
|
113
|
+
location=self._location,
|
|
114
|
+
project=self._project,
|
|
115
|
+
engine=self._engine,
|
|
116
|
+
vault=self._vault,
|
|
117
|
+
dlp=self._dlp,
|
|
118
|
+
user_email=self._user_email,
|
|
119
|
+
user_role=self._user_role,
|
|
120
|
+
)
|
|
121
|
+
|
|
122
|
+
|
|
123
|
+
class IrisGenerativeModel:
|
|
124
|
+
"""Governed wrapper around vertexai.generative_models.GenerativeModel."""
|
|
125
|
+
|
|
126
|
+
def __init__(
|
|
127
|
+
self,
|
|
128
|
+
model_name: str,
|
|
129
|
+
passport: AgentPassport,
|
|
130
|
+
location: Optional[str] = None,
|
|
131
|
+
project: Optional[str] = None,
|
|
132
|
+
engine: Optional[CedarEngine] = None,
|
|
133
|
+
vault: Optional[EvidenceVault] = None,
|
|
134
|
+
dlp: Optional[DLPScanner] = None,
|
|
135
|
+
user_email: Optional[str] = None,
|
|
136
|
+
user_role: Optional[str] = None,
|
|
137
|
+
):
|
|
138
|
+
generative_models = _lazy_vertexai_generative_models()
|
|
139
|
+
self._model_name = model_name
|
|
140
|
+
self._passport = passport
|
|
141
|
+
self._location = location
|
|
142
|
+
self._project = project
|
|
143
|
+
self._engine = engine or CedarEngine()
|
|
144
|
+
self._vault = vault or EvidenceVault(agent_id=passport.agent_id)
|
|
145
|
+
self._dlp = dlp or DLPScanner(passport)
|
|
146
|
+
self._user_email = user_email
|
|
147
|
+
self._user_role = user_role
|
|
148
|
+
self._model = generative_models.GenerativeModel(model_name)
|
|
149
|
+
|
|
150
|
+
def _record_with_metadata(
|
|
151
|
+
self, ctx: EvaluationContext, result: PolicyResult, location_violations: list[Violation]
|
|
152
|
+
) -> None:
|
|
153
|
+
entry = {
|
|
154
|
+
"event_id": str(uuid.uuid4()),
|
|
155
|
+
"timestamp": datetime.utcnow().isoformat(),
|
|
156
|
+
"agent_id": self._passport.agent_id,
|
|
157
|
+
"action": ctx.action,
|
|
158
|
+
"resource": ctx.resource,
|
|
159
|
+
"environment": ctx.environment.value,
|
|
160
|
+
"decision": result.decision,
|
|
161
|
+
"gcp_project": self._project,
|
|
162
|
+
"gcp_location": self._location,
|
|
163
|
+
"additional": {
|
|
164
|
+
**ctx.additional,
|
|
165
|
+
"gcp_project": self._project,
|
|
166
|
+
"gcp_location": self._location,
|
|
167
|
+
},
|
|
168
|
+
"violations": [
|
|
169
|
+
{
|
|
170
|
+
"rule_id": v.rule_id,
|
|
171
|
+
"severity": v.severity.value,
|
|
172
|
+
"message": v.message,
|
|
173
|
+
"compliance_refs": v.compliance_refs,
|
|
174
|
+
}
|
|
175
|
+
for v in (list(result.violations) + list(location_violations))
|
|
176
|
+
],
|
|
177
|
+
}
|
|
178
|
+
with open(self._vault._log_file, "a") as handle:
|
|
179
|
+
handle.write(json.dumps(entry) + "\n")
|
|
180
|
+
|
|
181
|
+
def _govern(self, action: str, contents: Any = None) -> None:
|
|
182
|
+
env = _current_environment()
|
|
183
|
+
prompt_text = "\n".join(_extract_text(contents))
|
|
184
|
+
dlp_result = enforce_prompt_dlp(
|
|
185
|
+
self._dlp,
|
|
186
|
+
self._vault,
|
|
187
|
+
self._passport,
|
|
188
|
+
env,
|
|
189
|
+
prompt_text,
|
|
190
|
+
resource=f"vertexai/{self._model_name}",
|
|
191
|
+
)
|
|
192
|
+
violations: list[Violation] = []
|
|
193
|
+
|
|
194
|
+
if self._location:
|
|
195
|
+
if self._passport.allowed_regions and self._location not in self._passport.allowed_regions:
|
|
196
|
+
violations.append(
|
|
197
|
+
Violation(
|
|
198
|
+
rule_id="IRIS-REGION-001",
|
|
199
|
+
severity=Severity.CRITICAL,
|
|
200
|
+
message=(
|
|
201
|
+
f"Location '{self._location}' is not in passport.allowed_regions."
|
|
202
|
+
),
|
|
203
|
+
compliance_refs=["iris:region-policy"],
|
|
204
|
+
remediation=(
|
|
205
|
+
"Use an allowed location or update passport.allowed_regions "
|
|
206
|
+
"with security approval."
|
|
207
|
+
),
|
|
208
|
+
)
|
|
209
|
+
)
|
|
210
|
+
if self._location.startswith("us-gov-"):
|
|
211
|
+
violations.append(
|
|
212
|
+
Violation(
|
|
213
|
+
rule_id="IRIS-REGION-002",
|
|
214
|
+
severity=Severity.HIGH,
|
|
215
|
+
message=(
|
|
216
|
+
f"Location '{self._location}' is a restricted government region."
|
|
217
|
+
),
|
|
218
|
+
compliance_refs=["fedramp:high:region-review"],
|
|
219
|
+
remediation=(
|
|
220
|
+
"Ensure FedRAMP High authorization is documented for this location."
|
|
221
|
+
),
|
|
222
|
+
)
|
|
223
|
+
)
|
|
224
|
+
else:
|
|
225
|
+
violations.append(
|
|
226
|
+
Violation(
|
|
227
|
+
rule_id="IRIS-REGION-003",
|
|
228
|
+
severity=Severity.CRITICAL,
|
|
229
|
+
message="Vertex AI location is required but missing.",
|
|
230
|
+
compliance_refs=["iris:region-policy"],
|
|
231
|
+
remediation="Provide a valid Vertex AI location before calling the model.",
|
|
232
|
+
)
|
|
233
|
+
)
|
|
234
|
+
|
|
235
|
+
if _is_fedramp_enabled(self._passport):
|
|
236
|
+
fedramp_location_violation = check_fedramp_location(self._location or "")
|
|
237
|
+
if fedramp_location_violation:
|
|
238
|
+
violations.append(fedramp_location_violation)
|
|
239
|
+
fedramp_model_violation = check_fedramp_model(self._model_name)
|
|
240
|
+
if fedramp_model_violation:
|
|
241
|
+
violations.append(fedramp_model_violation)
|
|
242
|
+
|
|
243
|
+
user_ctx = UserContext.from_params(self._user_email, self._user_role)
|
|
244
|
+
ctx = EvaluationContext(
|
|
245
|
+
agent_id=self._passport.agent_id,
|
|
246
|
+
action=action,
|
|
247
|
+
resource=f"vertexai/{self._model_name}",
|
|
248
|
+
resource_type="api",
|
|
249
|
+
environment=env,
|
|
250
|
+
data_region=self._location,
|
|
251
|
+
data_classification=self._passport.data_classification.value,
|
|
252
|
+
dlp_prompt_findings=dlp_result.findings,
|
|
253
|
+
additional={
|
|
254
|
+
"model": self._model_name,
|
|
255
|
+
"gcp_project": self._project,
|
|
256
|
+
"gcp_location": self._location,
|
|
257
|
+
},
|
|
258
|
+
**user_ctx.evaluation_fields(),
|
|
259
|
+
)
|
|
260
|
+
result = self._engine.evaluate(self._passport, ctx)
|
|
261
|
+
if violations:
|
|
262
|
+
result = PolicyResult(
|
|
263
|
+
decision="DENY",
|
|
264
|
+
violations=list(result.violations) + violations,
|
|
265
|
+
agent_id=result.agent_id,
|
|
266
|
+
action=result.action,
|
|
267
|
+
resource=result.resource,
|
|
268
|
+
environment=result.environment,
|
|
269
|
+
)
|
|
270
|
+
|
|
271
|
+
self._record_with_metadata(ctx, result, [])
|
|
272
|
+
_enforce_result(result)
|
|
273
|
+
|
|
274
|
+
def _scan_response(self, response: Any) -> Any:
|
|
275
|
+
env = _current_environment()
|
|
276
|
+
response_text = extract_gemini_response_text(response)
|
|
277
|
+
blocked, _ = handle_response_dlp(
|
|
278
|
+
self._dlp,
|
|
279
|
+
self._vault,
|
|
280
|
+
self._passport,
|
|
281
|
+
env,
|
|
282
|
+
response_text,
|
|
283
|
+
response,
|
|
284
|
+
resource=f"vertexai/{self._model_name}",
|
|
285
|
+
)
|
|
286
|
+
return blocked
|
|
287
|
+
|
|
288
|
+
def generate_content(self, contents: Any, **kwargs: Any) -> Any:
|
|
289
|
+
self._govern(action="generate_content", contents=contents)
|
|
290
|
+
response = self._model.generate_content(contents, **kwargs)
|
|
291
|
+
return self._scan_response(response)
|
|
292
|
+
|
|
293
|
+
def generate_content_stream(self, **kwargs: Any) -> Any:
|
|
294
|
+
self._govern(action="generate_content_stream", contents=kwargs.get("contents"))
|
|
295
|
+
return self._model.generate_content_stream(**kwargs)
|
|
296
|
+
|
|
297
|
+
def start_chat(self) -> "IrisChatSession":
|
|
298
|
+
return IrisChatSession(
|
|
299
|
+
chat_session=self._model.start_chat(),
|
|
300
|
+
model=self,
|
|
301
|
+
)
|
|
302
|
+
|
|
303
|
+
|
|
304
|
+
class IrisChatSession:
|
|
305
|
+
"""Governed wrapper around Vertex AI chat sessions."""
|
|
306
|
+
|
|
307
|
+
def __init__(self, chat_session: Any, model: IrisGenerativeModel):
|
|
308
|
+
self._chat_session = chat_session
|
|
309
|
+
self._model = model
|
|
310
|
+
|
|
311
|
+
def send_message(self, content: Any, **kwargs: Any) -> Any:
|
|
312
|
+
self._model._govern(action="chat.send_message", contents=content)
|
|
313
|
+
response = self._chat_session.send_message(content, **kwargs)
|
|
314
|
+
return self._model._scan_response(response)
|
|
@@ -0,0 +1,89 @@
|
|
|
1
|
+
"""FedRAMP-specific governance checks for Vertex AI integrations."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from typing import Optional
|
|
6
|
+
|
|
7
|
+
from iris_core.models.policy import Severity, Violation
|
|
8
|
+
|
|
9
|
+
FEDRAMP_AUTHORIZED_VERTEX_REGIONS = {
|
|
10
|
+
"us-central1",
|
|
11
|
+
"us-east1",
|
|
12
|
+
"us-east4",
|
|
13
|
+
"us-west1",
|
|
14
|
+
"us-west2",
|
|
15
|
+
"us-west3",
|
|
16
|
+
"us-west4",
|
|
17
|
+
"us-south1",
|
|
18
|
+
}
|
|
19
|
+
|
|
20
|
+
FEDRAMP_AUTHORIZED_VERTEX_MODELS = {
|
|
21
|
+
"gemini-1.5-pro",
|
|
22
|
+
"gemini-1.5-flash",
|
|
23
|
+
}
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def check_fedramp_location(location: str) -> Optional[Violation]:
|
|
27
|
+
"""Return violation when Vertex AI location is not FedRAMP-authorized."""
|
|
28
|
+
normalized = (location or "").strip().lower()
|
|
29
|
+
if not normalized:
|
|
30
|
+
return Violation(
|
|
31
|
+
rule_id="FEDRAMP-001",
|
|
32
|
+
severity=Severity.CRITICAL,
|
|
33
|
+
message="FedRAMP location is missing for this Vertex AI request.",
|
|
34
|
+
compliance_refs=["fedramp:moderate:residency"],
|
|
35
|
+
remediation=(
|
|
36
|
+
"Set a FedRAMP-authorized Vertex AI region and retry the request."
|
|
37
|
+
),
|
|
38
|
+
)
|
|
39
|
+
|
|
40
|
+
if normalized in FEDRAMP_AUTHORIZED_VERTEX_REGIONS:
|
|
41
|
+
return None
|
|
42
|
+
|
|
43
|
+
if normalized.startswith("us-gov-"):
|
|
44
|
+
return Violation(
|
|
45
|
+
rule_id="FEDRAMP-001",
|
|
46
|
+
severity=Severity.CRITICAL,
|
|
47
|
+
message=(
|
|
48
|
+
f"Location '{location}' requires additional FedRAMP High authorization."
|
|
49
|
+
),
|
|
50
|
+
compliance_refs=["fedramp:high:authorization-boundary"],
|
|
51
|
+
remediation=(
|
|
52
|
+
"Obtain documented FedRAMP High authorization for this region "
|
|
53
|
+
"or use an authorized FedRAMP Moderate region."
|
|
54
|
+
),
|
|
55
|
+
)
|
|
56
|
+
|
|
57
|
+
return Violation(
|
|
58
|
+
rule_id="FEDRAMP-001",
|
|
59
|
+
severity=Severity.CRITICAL,
|
|
60
|
+
message=(
|
|
61
|
+
f"Location '{location}' is not in the FedRAMP-authorized Vertex AI region list."
|
|
62
|
+
),
|
|
63
|
+
compliance_refs=["fedramp:moderate:residency"],
|
|
64
|
+
remediation=(
|
|
65
|
+
"Use one of the authorized regions: "
|
|
66
|
+
+ ", ".join(sorted(FEDRAMP_AUTHORIZED_VERTEX_REGIONS))
|
|
67
|
+
+ "."
|
|
68
|
+
),
|
|
69
|
+
)
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def check_fedramp_model(model_name: str) -> Optional[Violation]:
|
|
73
|
+
"""Return violation when Vertex AI model authorization cannot be verified."""
|
|
74
|
+
normalized = (model_name or "").strip().lower()
|
|
75
|
+
if normalized in FEDRAMP_AUTHORIZED_VERTEX_MODELS:
|
|
76
|
+
return None
|
|
77
|
+
|
|
78
|
+
return Violation(
|
|
79
|
+
rule_id="FEDRAMP-002",
|
|
80
|
+
severity=Severity.HIGH,
|
|
81
|
+
message=(
|
|
82
|
+
f"Model '{model_name}' is not on the current FedRAMP-authorized Vertex AI list."
|
|
83
|
+
),
|
|
84
|
+
compliance_refs=["fedramp:ai-workload:authorization"],
|
|
85
|
+
remediation=(
|
|
86
|
+
"Use a known authorized model (gemini-1.5-pro or gemini-1.5-flash) "
|
|
87
|
+
"or complete an authorization review for this model."
|
|
88
|
+
),
|
|
89
|
+
)
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=68", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "iris-security-vertexai"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "IRIS governance for Vertex AI via google-cloud-aiplatform"
|
|
9
|
+
license = { text = "Apache-2.0" }
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
authors = [{ name = "IRIS Platform", email = "sdk@iris.ai" }]
|
|
12
|
+
dependencies = [
|
|
13
|
+
"iris-security-core>=0.1.0",
|
|
14
|
+
"iris-security-sdk>=0.1.0",
|
|
15
|
+
]
|
|
16
|
+
|
|
17
|
+
[project.optional-dependencies]
|
|
18
|
+
vertexai = ["google-cloud-aiplatform>=1.60"]
|
|
19
|
+
dev = [
|
|
20
|
+
"pytest>=8.0",
|
|
21
|
+
"ruff>=0.4",
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
[project.urls]
|
|
25
|
+
Homepage = "https://github.com/gimartinb/iris-sdk"
|
|
26
|
+
Repository = "https://github.com/gimartinb/iris-sdk"
|
|
27
|
+
|
|
28
|
+
[tool.setuptools]
|
|
29
|
+
packages = ["iris_vertexai"]
|
|
30
|
+
|
|
31
|
+
[tool.ruff]
|
|
32
|
+
line-length = 100
|
|
33
|
+
target-version = "py310"
|
|
34
|
+
|
|
35
|
+
[tool.pytest.ini_options]
|
|
36
|
+
testpaths = ["tests"]
|
|
@@ -0,0 +1,167 @@
|
|
|
1
|
+
"""Integration tests for iris-vertexai with mocked Vertex AI SDK."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import json
|
|
6
|
+
import types
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
from unittest.mock import MagicMock, patch
|
|
9
|
+
|
|
10
|
+
import pytest
|
|
11
|
+
|
|
12
|
+
from iris import IrisViolationError
|
|
13
|
+
from iris_core.engine.cedar import CedarEngine
|
|
14
|
+
from iris_core.models.passport import AgentPassport, ComplianceTag, DataClassification, Environment
|
|
15
|
+
from iris_vertexai import IrisVertexAI
|
|
16
|
+
from iris_vertexai.fedramp import check_fedramp_location, check_fedramp_model
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def _mock_vertex_modules():
|
|
20
|
+
vertexai_module = types.ModuleType("vertexai")
|
|
21
|
+
vertexai_module.init = MagicMock()
|
|
22
|
+
|
|
23
|
+
gm_module = types.ModuleType("vertexai.generative_models")
|
|
24
|
+
response = MagicMock()
|
|
25
|
+
response.text = "Vertex response"
|
|
26
|
+
chat_response = MagicMock()
|
|
27
|
+
chat_response.text = "chat response"
|
|
28
|
+
chat_session = MagicMock()
|
|
29
|
+
chat_session.send_message.return_value = chat_response
|
|
30
|
+
|
|
31
|
+
model_instance = MagicMock()
|
|
32
|
+
model_instance.generate_content.return_value = response
|
|
33
|
+
model_instance.generate_content_stream.return_value = iter([response])
|
|
34
|
+
model_instance.start_chat.return_value = chat_session
|
|
35
|
+
|
|
36
|
+
gm_module.GenerativeModel = MagicMock(return_value=model_instance)
|
|
37
|
+
return vertexai_module, gm_module, model_instance, chat_session
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
class _FedrampTag:
|
|
41
|
+
value = "fedramp"
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
@pytest.fixture
|
|
45
|
+
def passport():
|
|
46
|
+
return AgentPassport(
|
|
47
|
+
name="vertex-agent",
|
|
48
|
+
owner="team@gov-agency.gov",
|
|
49
|
+
data_classification=DataClassification.INTERNAL,
|
|
50
|
+
compliance_tags=[ComplianceTag.COLORADO_AI_ACT],
|
|
51
|
+
environments=[Environment.DEV, Environment.PRODUCTION],
|
|
52
|
+
allowed_regions=["us-central1", "us-east1"],
|
|
53
|
+
)
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def _permit_engine(passport: AgentPassport) -> CedarEngine:
|
|
57
|
+
engine = CedarEngine()
|
|
58
|
+
engine.load_policy(passport.agent_id, "permit(principal, action, resource);")
|
|
59
|
+
return engine
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def test_authorized_location_permitted(passport, tmp_path, monkeypatch):
|
|
63
|
+
monkeypatch.setenv("IRIS_ENV", "dev")
|
|
64
|
+
vertexai_module, gm_module, model_instance, _ = _mock_vertex_modules()
|
|
65
|
+
|
|
66
|
+
with patch.dict(
|
|
67
|
+
"sys.modules",
|
|
68
|
+
{"vertexai": vertexai_module, "vertexai.generative_models": gm_module},
|
|
69
|
+
):
|
|
70
|
+
client = IrisVertexAI(
|
|
71
|
+
passport=passport,
|
|
72
|
+
project="proj-1",
|
|
73
|
+
location="us-central1",
|
|
74
|
+
evidence_vault_dir=tmp_path,
|
|
75
|
+
)
|
|
76
|
+
client._engine = _permit_engine(passport)
|
|
77
|
+
model = client.get_model("gemini-1.5-pro")
|
|
78
|
+
result = model.generate_content("hello")
|
|
79
|
+
|
|
80
|
+
assert result.text == "Vertex response"
|
|
81
|
+
model_instance.generate_content.assert_called_once()
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def test_unauthorized_location_blocked(passport, tmp_path, monkeypatch):
|
|
85
|
+
monkeypatch.setenv("IRIS_ENV", "production")
|
|
86
|
+
vertexai_module, gm_module, model_instance, _ = _mock_vertex_modules()
|
|
87
|
+
|
|
88
|
+
with patch.dict(
|
|
89
|
+
"sys.modules",
|
|
90
|
+
{"vertexai": vertexai_module, "vertexai.generative_models": gm_module},
|
|
91
|
+
):
|
|
92
|
+
client = IrisVertexAI(
|
|
93
|
+
passport=passport,
|
|
94
|
+
project="proj-1",
|
|
95
|
+
location="europe-west1",
|
|
96
|
+
evidence_vault_dir=tmp_path,
|
|
97
|
+
)
|
|
98
|
+
client._engine = _permit_engine(passport)
|
|
99
|
+
model = client.get_model("gemini-1.5-pro")
|
|
100
|
+
with pytest.raises(IrisViolationError):
|
|
101
|
+
model.generate_content("hello")
|
|
102
|
+
|
|
103
|
+
model_instance.generate_content.assert_not_called()
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def test_fedramp_location_check():
|
|
107
|
+
assert check_fedramp_location("us-central1") is None
|
|
108
|
+
violation = check_fedramp_location("europe-west1")
|
|
109
|
+
assert violation is not None
|
|
110
|
+
assert violation.rule_id == "FEDRAMP-001"
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
def test_fedramp_model_check():
|
|
114
|
+
assert check_fedramp_model("gemini-1.5-pro") is None
|
|
115
|
+
violation = check_fedramp_model("gemini-2.0-flash")
|
|
116
|
+
assert violation is not None
|
|
117
|
+
assert violation.rule_id == "FEDRAMP-002"
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def test_chat_session_intercept(passport, tmp_path, monkeypatch):
|
|
121
|
+
monkeypatch.setenv("IRIS_ENV", "dev")
|
|
122
|
+
vertexai_module, gm_module, _, chat_session = _mock_vertex_modules()
|
|
123
|
+
|
|
124
|
+
with patch.dict(
|
|
125
|
+
"sys.modules",
|
|
126
|
+
{"vertexai": vertexai_module, "vertexai.generative_models": gm_module},
|
|
127
|
+
):
|
|
128
|
+
client = IrisVertexAI(
|
|
129
|
+
passport=passport,
|
|
130
|
+
project="proj-1",
|
|
131
|
+
location="us-central1",
|
|
132
|
+
evidence_vault_dir=tmp_path,
|
|
133
|
+
)
|
|
134
|
+
client._engine = _permit_engine(passport)
|
|
135
|
+
model = client.get_model("gemini-1.5-pro")
|
|
136
|
+
chat = model.start_chat()
|
|
137
|
+
chat.send_message("hello")
|
|
138
|
+
|
|
139
|
+
chat_session.send_message.assert_called_once()
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
def test_evidence_vault_includes_gcp_metadata(passport, tmp_path, monkeypatch):
|
|
143
|
+
monkeypatch.setenv("IRIS_ENV", "dev")
|
|
144
|
+
vertexai_module, gm_module, _, _ = _mock_vertex_modules()
|
|
145
|
+
passport.compliance_tags = [ComplianceTag.COLORADO_AI_ACT, _FedrampTag()]
|
|
146
|
+
|
|
147
|
+
with patch.dict(
|
|
148
|
+
"sys.modules",
|
|
149
|
+
{"vertexai": vertexai_module, "vertexai.generative_models": gm_module},
|
|
150
|
+
):
|
|
151
|
+
client = IrisVertexAI(
|
|
152
|
+
passport=passport,
|
|
153
|
+
project="proj-1",
|
|
154
|
+
location="us-central1",
|
|
155
|
+
evidence_vault_dir=tmp_path,
|
|
156
|
+
)
|
|
157
|
+
client._engine = _permit_engine(passport)
|
|
158
|
+
model = client.get_model("gemini-1.5-pro")
|
|
159
|
+
model.generate_content("hello")
|
|
160
|
+
|
|
161
|
+
events_file = Path(tmp_path) / passport.agent_id / "events.jsonl"
|
|
162
|
+
lines = events_file.read_text().strip().splitlines()
|
|
163
|
+
event = json.loads(lines[-1])
|
|
164
|
+
assert event["gcp_project"] == "proj-1"
|
|
165
|
+
assert event["gcp_location"] == "us-central1"
|
|
166
|
+
assert event["additional"]["gcp_project"] == "proj-1"
|
|
167
|
+
assert event["additional"]["gcp_location"] == "us-central1"
|