ben.carrasco/prowler
1
0 Fork 0
mirror of https://github.com/prowler-cloud/prowler.git synced 2026-10-25 04:44:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(gcp): add vertexai workbench terminal, download and auto-upgrade… (#13077)

Browse Source 
Co-authored-by: Rubén De la Torre Vico <98956809+puchy22@users.noreply.github.com>
This commit is contained in:
Rajesh NadimintiandRubén De la Torre Vico authored and GitHub committed 2026-10-20 12:43:06 +02:00
1 parent 4d922da2a6
commit 4980ded0ae
18 files changed
+660 -15
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

prowler/changelog.d/gcp-vertexai-workbench-mlde-checks.added.md
+1
View File
@@ -0,0 +1 @@
`vertexai_workbench_instance_terminal_access_disabled`, `vertexai_workbench_instance_file_download_disabled` and `vertexai_workbench_instance_auto_upgrade_enabled` checks for GCP provider, with Vertex AI Workbench checks mapped to CCC MLDE requirements
prowler/compliance/gcp/ccc_gcp.json
+24 -8
View File
@@ -6815,7 +6815,9 @@
]
}
],
"Checks": []
"Checks": [
"vertexai_workbench_instance_root_access_disabled"
]
},
{
"Id": "CCC.MLDE.CN03.AR02",
@@ -6871,7 +6873,9 @@
]
}
],
"Checks": []
"Checks": [
"vertexai_workbench_instance_root_access_disabled"
]
},
{
"Id": "CCC.MLDE.CN04.AR01",
@@ -6923,7 +6927,9 @@
]
}
],
"Checks": []
"Checks": [
"vertexai_workbench_instance_terminal_access_disabled"
]
},
{
"Id": "CCC.MLDE.CN04.AR02",
@@ -6978,7 +6984,9 @@
]
}
],
"Checks": []
"Checks": [
"vertexai_workbench_instance_terminal_access_disabled"
]
},
{
"Id": "CCC.MLDE.CN02.AR01",
@@ -7033,7 +7041,9 @@
]
}
],
"Checks": []
"Checks": [
"vertexai_workbench_instance_file_download_disabled"
]
},
{
"Id": "CCC.MLDE.CN02.AR02",
@@ -7250,7 +7260,9 @@
]
}
],
"Checks": []
"Checks": [
"vertexai_workbench_instance_auto_upgrade_enabled"
]
},
{
"Id": "CCC.MLDE.CN06.AR02",
@@ -7360,7 +7372,9 @@
]
}
],
"Checks": []
"Checks": [
"vertexai_workbench_instance_public_ip_disabled"
]
},
{
"Id": "CCC.MLDE.CN07.AR02",
@@ -7416,7 +7430,9 @@
]
}
],
"Checks": []
"Checks": [
"vertexai_workbench_instance_public_ip_disabled"
]
},
{
"Id": "CCC.MLDE.CN08.AR01",
prowler/providers/gcp/services/vertexai/vertexai_service.py
+29 -4
View File
@@ -9,6 +9,17 @@ from prowler.providers.gcp.lib.service.service import GCPService
# Workbench metadata key that controls sudo access for the notebook user.
ROOT_ACCESS_DISABLED_METADATA_KEY = "notebook-disable-root"
# Workbench metadata key that removes the JupyterLab terminal.
TERMINAL_ACCESS_DISABLED_METADATA_KEY = "notebook-disable-terminal"
# Workbench metadata key that blocks file downloads from the JupyterLab UI.
FILE_DOWNLOAD_DISABLED_METADATA_KEY = "notebook-disable-downloads"
# Workbench metadata key holding the Unix-cron schedule for automatic upgrades.
UPGRADE_SCHEDULE_METADATA_KEY = "notebook-upgrade-schedule"
def _metadata_flag_enabled(metadata: dict, key: str) -> bool:
"""Return True when a boolean Workbench metadata key is set to `true`."""
return str(metadata.get(key, "")).lower() == "true"
class VertexAI(GCPService):
@@ -49,10 +60,21 @@ class VertexAI(GCPService):
public_ip_disabled=gce_setup.get(
"disablePublicIp", False
),
root_access_disabled=str(
metadata.get(ROOT_ACCESS_DISABLED_METADATA_KEY, "")
).lower()
== "true",
root_access_disabled=_metadata_flag_enabled(
metadata, ROOT_ACCESS_DISABLED_METADATA_KEY
),
terminal_access_disabled=_metadata_flag_enabled(
metadata, TERMINAL_ACCESS_DISABLED_METADATA_KEY
),
file_download_disabled=_metadata_flag_enabled(
metadata, FILE_DOWNLOAD_DISABLED_METADATA_KEY
),
upgrade_schedule=(
str(
metadata.get(UPGRADE_SCHEDULE_METADATA_KEY, "")
).strip()
or None
),
network=(
network_interfaces[0].get("network")
if network_interfaces
@@ -86,4 +108,7 @@ class WorkbenchInstance(BaseModel):
state: str
public_ip_disabled: bool = False
root_access_disabled: bool = False
terminal_access_disabled: bool = False
file_download_disabled: bool = False
upgrade_schedule: Optional[str] = None
network: Optional[str] = None
prowler/providers/gcp/services/vertexai/vertexai_workbench_instance_auto_upgrade_enabled/__init__.py
View File
Whitespace-only changes.
prowler/providers/gcp/services/vertexai/vertexai_workbench_instance_auto_upgrade_enabled/vertexai_workbench_instance_auto_upgrade_enabled.metadata.json
+38
View File
@@ -0,0 +1,38 @@
{
"Provider": "gcp",
"CheckID": "vertexai_workbench_instance_auto_upgrade_enabled",
"CheckTitle": "Vertex AI Workbench instance has automatic upgrades scheduled",
"CheckType": [],
"ServiceName": "vertexai",
"SubServiceName": "",
"ResourceIdTemplate": "",
"Severity": "low",
"ResourceType": "notebooks.googleapis.com/Instance",
"ResourceGroup": "ai_ml",
"Description": "**Vertex AI Workbench instances** are not upgraded automatically by default. The finding determines whether automatic upgrades are scheduled by setting the `notebook-upgrade-schedule` instance metadata key to a Unix-cron expression (GMT).",
"Risk": "Without scheduled upgrades, long-lived instances keep running outdated images:\n- **Known vulnerabilities**: JupyterLab, Python ML libraries, and OS packages stay unpatched against published CVEs\n- **Drift**: instances fall behind the image versions the rest of the organization has validated",
"RelatedUrl": "",
"AdditionalURLs": [
"https://cloud.google.com/vertex-ai/docs/workbench/instances/manage-metadata",
"https://cloud.google.com/vertex-ai/docs/workbench/instances/upgrade"
],
"Remediation": {
"Code": {
"CLI": "gcloud workbench instances update <example_resource_name> --location=<example_zone> --metadata=notebook-upgrade-schedule=\"00 19 * * MON\"",
"NativeIaC": "",
"Other": "1. In Google Cloud Console, go to Vertex AI > Workbench > Instances\n2. Click the instance name to open the Instance details page\n3. In the Environment auto-upgrade section, set the upgrade schedule\n4. Click Submit",
"Terraform": "```hcl\nresource \"google_workbench_instance\" \"<example_resource_name>\" {\n name = \"<example_resource_name>\"\n location = \"<example_zone>\"\n\n gce_setup {\n metadata = {\n notebook-upgrade-schedule = \"00 19 * * MON\" # FIX: Upgrades the instance every Monday at 19:00 GMT\n }\n }\n}\n```"
},
"Recommendation": {
"Text": "Schedule automatic upgrades for every Workbench instance and pick a window that does not interrupt long-running jobs.\n- Store notebooks and data on the data disk or in Cloud Storage so upgrades do not lose work\n- Review release notes for breaking changes in preinstalled frameworks",
"Url": "https://hub.prowler.com/check/vertexai_workbench_instance_auto_upgrade_enabled"
}
},
"Categories": [
"vulnerabilities",
"gen-ai"
],
"DependsOn": [],
"RelatedTo": [],
"Notes": "The check only verifies that a schedule is set. It does not validate the cron expression or how often it runs."
}
prowler/providers/gcp/services/vertexai/vertexai_workbench_instance_auto_upgrade_enabled/vertexai_workbench_instance_auto_upgrade_enabled.py
+34
View File
@@ -0,0 +1,34 @@
from prowler.lib.check.models import Check, Check_Report_GCP
from prowler.providers.gcp.services.vertexai.vertexai_client import vertexai_client
class vertexai_workbench_instance_auto_upgrade_enabled(Check):
"""Ensure Vertex AI Workbench instances have automatic upgrades scheduled.
Workbench only upgrades an instance automatically when the
`notebook-upgrade-schedule` metadata key holds a Unix-cron schedule.
The check does not evaluate how often the schedule runs.
- PASS: The instance has a non-empty `notebook-upgrade-schedule`.
- FAIL: The instance is never upgraded automatically.
"""
def execute(self) -> list[Check_Report_GCP]:
"""Execute the automatic upgrade check across all Workbench instances.
Returns:
A list of `Check_Report_GCP` findings, one per non-deleted
Workbench instance.
"""
findings = []
for instance in vertexai_client.workbench_instances:
if instance.state == "DELETED":
continue
report = Check_Report_GCP(metadata=self.metadata(), resource=instance)
if instance.upgrade_schedule:
report.status = "PASS"
report.status_extended = f"Vertex AI Workbench instance {instance.name} has automatic upgrades scheduled ({instance.upgrade_schedule})."
else:
report.status = "FAIL"
report.status_extended = f"Vertex AI Workbench instance {instance.name} does not have automatic upgrades scheduled."
findings.append(report)
return findings
prowler/providers/gcp/services/vertexai/vertexai_workbench_instance_file_download_disabled/__init__.py
View File
Whitespace-only changes.
prowler/providers/gcp/services/vertexai/vertexai_workbench_instance_file_download_disabled/vertexai_workbench_instance_file_download_disabled.metadata.json
+39
View File
@@ -0,0 +1,39 @@
{
"Provider": "gcp",
"CheckID": "vertexai_workbench_instance_file_download_disabled",
"CheckTitle": "Vertex AI Workbench instance has file downloads disabled",
"CheckType": [],
"ServiceName": "vertexai",
"SubServiceName": "",
"ResourceIdTemplate": "",
"Severity": "medium",
"ResourceType": "notebooks.googleapis.com/Instance",
"ResourceGroup": "ai_ml",
"Description": "**Vertex AI Workbench instances** let JupyterLab users **download any file** from the instance through the browser by default. The finding determines whether file downloads are disabled by setting the `notebook-disable-downloads` instance metadata key to `true`.",
"Risk": "With downloads enabled, a notebook user, or anyone who hijacks their browser session, can:\n- **Exfiltration**: save training data, query results from BigQuery or Cloud Storage, and model artifacts to an unmanaged device in a few clicks\n- **Bypass data controls**: move data outside VPC Service Controls perimeters and data loss prevention controls that apply to cloud storage",
"RelatedUrl": "",
"AdditionalURLs": [
"https://cloud.google.com/vertex-ai/docs/workbench/instances/manage-metadata",
"https://cloud.google.com/vertex-ai/docs/workbench/instances/introduction"
],
"Remediation": {
"Code": {
"CLI": "gcloud workbench instances update <example_resource_name> --location=<example_zone> --metadata=notebook-disable-downloads=true",
"NativeIaC": "",
"Other": "1. In Google Cloud Console, go to Vertex AI > Workbench > Instances\n2. Click the instance name and open the Software and security tab\n3. Under Metadata, add the key notebook-disable-downloads with the value true\n4. Click Submit",
"Terraform": "```hcl\nresource \"google_workbench_instance\" \"<example_resource_name>\" {\n name = \"<example_resource_name>\"\n location = \"<example_zone>\"\n\n gce_setup {\n metadata = {\n notebook-disable-downloads = \"true\" # FIX: Blocks file downloads from the JupyterLab UI\n }\n }\n}\n```"
},
"Recommendation": {
"Text": "Disable file downloads on Workbench instances that handle sensitive data, and keep data inside governed storage.\n- Share results through Cloud Storage or BigQuery, where IAM and audit logging apply\n- Combine with **VPC Service Controls** to stop data leaving the perimeter by other paths\n- Disable terminal access as well, since a shell can still copy data off the VM",
"Url": "https://hub.prowler.com/check/vertexai_workbench_instance_file_download_disabled"
}
},
"Categories": [
"gen-ai"
],
"DependsOn": [],
"RelatedTo": [
"vertexai_workbench_instance_terminal_access_disabled"
],
"Notes": "This setting only blocks downloads through the JupyterLab UI. Users who can run code or open a terminal can still copy data off the instance, so pair it with vertexai_workbench_instance_terminal_access_disabled."
}
prowler/providers/gcp/services/vertexai/vertexai_workbench_instance_file_download_disabled/vertexai_workbench_instance_file_download_disabled.py
+33
View File
@@ -0,0 +1,33 @@
from prowler.lib.check.models import Check, Check_Report_GCP
from prowler.providers.gcp.services.vertexai.vertexai_client import vertexai_client
class vertexai_workbench_instance_file_download_disabled(Check):
"""Ensure Vertex AI Workbench instances have file downloads disabled.
JupyterLab lets users download any file from the instance through the
browser unless the `notebook-disable-downloads` metadata key is `true`.
- PASS: The instance has `notebook-disable-downloads` set to `true`.
- FAIL: Notebook users can download files from the JupyterLab UI.
"""
def execute(self) -> list[Check_Report_GCP]:
"""Execute the file download check across all Workbench instances.
Returns:
A list of `Check_Report_GCP` findings, one per non-deleted
Workbench instance.
"""
findings = []
for instance in vertexai_client.workbench_instances:
if instance.state == "DELETED":
continue
report = Check_Report_GCP(metadata=self.metadata(), resource=instance)
if instance.file_download_disabled:
report.status = "PASS"
report.status_extended = f"Vertex AI Workbench instance {instance.name} has file downloads disabled."
else:
report.status = "FAIL"
report.status_extended = f"Vertex AI Workbench instance {instance.name} has file downloads enabled."
findings.append(report)
return findings
prowler/providers/gcp/services/vertexai/vertexai_workbench_instance_public_ip_disabled/vertexai_workbench_instance_public_ip_disabled.metadata.json
+1 -1
View File
@@ -4,7 +4,7 @@
"CheckTitle": "Vertex AI Workbench instance does not have a public IP address",
"CheckType": [],
"ServiceName": "vertexai",
"SubServiceName": "workbench",
"SubServiceName": "",
"ResourceIdTemplate": "",
"Severity": "high",
"ResourceType": "notebooks.googleapis.com/Instance",
prowler/providers/gcp/services/vertexai/vertexai_workbench_instance_root_access_disabled/vertexai_workbench_instance_root_access_disabled.metadata.json
+1 -1
View File
@@ -4,7 +4,7 @@
"CheckTitle": "Vertex AI Workbench instance has root access disabled",
"CheckType": [],
"ServiceName": "vertexai",
"SubServiceName": "workbench",
"SubServiceName": "",
"ResourceIdTemplate": "",
"Severity": "medium",
"ResourceType": "notebooks.googleapis.com/Instance",
prowler/providers/gcp/services/vertexai/vertexai_workbench_instance_terminal_access_disabled/__init__.py
View File
Whitespace-only changes.
prowler/providers/gcp/services/vertexai/vertexai_workbench_instance_terminal_access_disabled/vertexai_workbench_instance_terminal_access_disabled.metadata.json
+41
View File
@@ -0,0 +1,41 @@
{
"Provider": "gcp",
"CheckID": "vertexai_workbench_instance_terminal_access_disabled",
"CheckTitle": "Vertex AI Workbench instance has terminal access disabled",
"CheckType": [],
"ServiceName": "vertexai",
"SubServiceName": "",
"ResourceIdTemplate": "",
"Severity": "medium",
"ResourceType": "notebooks.googleapis.com/Instance",
"ResourceGroup": "ai_ml",
"Description": "**Vertex AI Workbench instances** give every JupyterLab user a **terminal** by default, which is an interactive shell on the VM. The finding determines whether terminal access is disabled by setting the `notebook-disable-terminal` instance metadata key to `true`.",
"Risk": "From the terminal, a notebook user, or anyone who hijacks their session, can:\n- **Credential access**: read the attached service account's access token from the metadata server and use it outside the notebook\n- **Exfiltration**: copy datasets and model artifacts off the VM with `gsutil`, `curl`, or `scp`\n- **Persistence**: start long-running processes that outlive the notebook session",
"RelatedUrl": "",
"AdditionalURLs": [
"https://cloud.google.com/vertex-ai/docs/workbench/instances/manage-metadata",
"https://cloud.google.com/vertex-ai/docs/workbench/instances/introduction"
],
"Remediation": {
"Code": {
"CLI": "gcloud workbench instances update <example_resource_name> --location=<example_zone> --metadata=notebook-disable-terminal=true",
"NativeIaC": "",
"Other": "1. In Google Cloud Console, go to Vertex AI > Workbench > Instances\n2. Click the instance name and open the Software and security tab\n3. Under Metadata, add the key notebook-disable-terminal with the value true\n4. Click Submit",
"Terraform": "```hcl\nresource \"google_workbench_instance\" \"<example_resource_name>\" {\n name = \"<example_resource_name>\"\n location = \"<example_zone>\"\n\n gce_setup {\n metadata = {\n notebook-disable-terminal = \"true\" # FIX: Removes the JupyterLab terminal\n }\n }\n}\n```"
},
"Recommendation": {
"Text": "Disable terminal access on Workbench instances that handle sensitive data.\n- Let users run code only through notebooks, where it can be reviewed and audited\n- Use a separate, tightly controlled instance for the rare tasks that need a shell\n- Attach a dedicated, minimally scoped **service account** to each instance to limit what a stolen token can reach",
"Url": "https://hub.prowler.com/check/vertexai_workbench_instance_terminal_access_disabled"
}
},
"Categories": [
"identity-access",
"gen-ai"
],
"DependsOn": [],
"RelatedTo": [
"vertexai_workbench_instance_root_access_disabled",
"vertexai_workbench_instance_file_download_disabled"
],
"Notes": "Disabling the terminal does not stop code execution from notebook cells. Pair it with vertexai_workbench_instance_root_access_disabled and a minimally scoped service account."
}
prowler/providers/gcp/services/vertexai/vertexai_workbench_instance_terminal_access_disabled/vertexai_workbench_instance_terminal_access_disabled.py
+33
View File
@@ -0,0 +1,33 @@
from prowler.lib.check.models import Check, Check_Report_GCP
from prowler.providers.gcp.services.vertexai.vertexai_client import vertexai_client
class vertexai_workbench_instance_terminal_access_disabled(Check):
"""Ensure Vertex AI Workbench instances have terminal access disabled.
Workbench exposes a JupyterLab terminal (a shell on the VM) unless the
`notebook-disable-terminal` metadata key is set to `true`.
- PASS: The instance has `notebook-disable-terminal` set to `true`.
- FAIL: Notebook users can open a terminal on the instance.
"""
def execute(self) -> list[Check_Report_GCP]:
"""Execute the terminal access check across all Workbench instances.
Returns:
A list of `Check_Report_GCP` findings, one per non-deleted
Workbench instance.
"""
findings = []
for instance in vertexai_client.workbench_instances:
if instance.state == "DELETED":
continue
report = Check_Report_GCP(metadata=self.metadata(), resource=instance)
if instance.terminal_access_disabled:
report.status = "PASS"
report.status_extended = f"Vertex AI Workbench instance {instance.name} has terminal access disabled."
else:
report.status = "FAIL"
report.status_extended = f"Vertex AI Workbench instance {instance.name} has terminal access enabled."
findings.append(report)
return findings
tests/providers/gcp/services/vertexai/vertexai_service_test.py
+69 -1
View File
@@ -11,6 +11,7 @@ _ZONE = "us-central1-a"
_INSTANCE_NAME = "my-workbench"
_INSTANCE_ID = f"projects/{GCP_PROJECT_ID}/locations/{_ZONE}/instances/{_INSTANCE_NAME}"
_NETWORK = f"projects/{GCP_PROJECT_ID}/global/networks/default"
_UPGRADE_SCHEDULE = "00 19 * * MON"
def _make_notebooks_client(instances_list=None, list_error=None):
@@ -65,7 +66,12 @@ class TestVertexAIService:
"state": "ACTIVE",
"gceSetup": {
"disablePublicIp": True,
"metadata": {"notebook-disable-root": "true"},
"metadata": {
"notebook-disable-root": "true",
"notebook-disable-terminal": "true",
"notebook-disable-downloads": "true",
"notebook-upgrade-schedule": _UPGRADE_SCHEDULE,
},
"networkInterfaces": [{"network": _NETWORK}],
},
}
@@ -82,6 +88,9 @@ class TestVertexAIService:
assert instance.state == "ACTIVE"
assert instance.public_ip_disabled is True
assert instance.root_access_disabled is True
assert instance.terminal_access_disabled is True
assert instance.file_download_disabled is True
assert instance.upgrade_schedule == _UPGRADE_SCHEDULE
assert instance.network == _NETWORK
def test_get_workbench_instance_defaults(self):
@@ -95,8 +104,67 @@ class TestVertexAIService:
assert instance.state == "STATE_UNSPECIFIED"
assert instance.public_ip_disabled is False
assert instance.root_access_disabled is False
assert instance.terminal_access_disabled is False
assert instance.file_download_disabled is False
assert instance.upgrade_schedule is None
assert instance.network is None
def test_disable_metadata_flags_are_case_insensitive(self):
vertexai_client = _build_service(
_make_notebooks_client(
instances_list=[
{
"name": _INSTANCE_ID,
"gceSetup": {
"metadata": {
"notebook-disable-terminal": "TRUE",
"notebook-disable-downloads": "True",
}
},
}
]
)
)
instance = vertexai_client.workbench_instances[0]
assert instance.terminal_access_disabled is True
assert instance.file_download_disabled is True
def test_disable_metadata_flags_false(self):
vertexai_client = _build_service(
_make_notebooks_client(
instances_list=[
{
"name": _INSTANCE_ID,
"gceSetup": {
"metadata": {
"notebook-disable-terminal": "false",
"notebook-disable-downloads": "false",
}
},
}
]
)
)
instance = vertexai_client.workbench_instances[0]
assert instance.terminal_access_disabled is False
assert instance.file_download_disabled is False
def test_blank_upgrade_schedule_is_treated_as_unset(self):
vertexai_client = _build_service(
_make_notebooks_client(
instances_list=[
{
"name": _INSTANCE_ID,
"gceSetup": {"metadata": {"notebook-upgrade-schedule": " "}},
}
]
)
)
assert vertexai_client.workbench_instances[0].upgrade_schedule is None
def test_root_access_metadata_is_case_insensitive(self):
vertexai_client = _build_service(
_make_notebooks_client(
tests/providers/gcp/services/vertexai/vertexai_workbench_instance_auto_upgrade_enabled/vertexai_workbench_instance_auto_upgrade_enabled_test.py
+103
View File
@@ -0,0 +1,103 @@
from unittest import mock
from tests.providers.gcp.gcp_fixtures import GCP_PROJECT_ID, set_mocked_gcp_provider
_ZONE = "us-central1-a"
_UPGRADE_SCHEDULE = "00 19 * * MON"
_CHECK_PATH = (
"prowler.providers.gcp.services.vertexai."
"vertexai_workbench_instance_auto_upgrade_enabled."
"vertexai_workbench_instance_auto_upgrade_enabled"
)
_CLIENT_PATH = f"{_CHECK_PATH}.vertexai_client"
def _instance_id(name: str) -> str:
return f"projects/{GCP_PROJECT_ID}/locations/{_ZONE}/instances/{name}"
def _mock_vertexai_client(instances=None):
vertexai_client = mock.MagicMock()
vertexai_client.workbench_instances = instances or []
return vertexai_client
def _run_check(vertexai_client):
with (
mock.patch(
"prowler.providers.common.provider.Provider.get_global_provider",
return_value=set_mocked_gcp_provider(),
),
mock.patch(_CLIENT_PATH, new=vertexai_client),
):
from prowler.providers.gcp.services.vertexai.vertexai_workbench_instance_auto_upgrade_enabled.vertexai_workbench_instance_auto_upgrade_enabled import (
vertexai_workbench_instance_auto_upgrade_enabled,
)
return vertexai_workbench_instance_auto_upgrade_enabled().execute()
def _workbench_instance(name, upgrade_schedule=None, state="ACTIVE"):
from prowler.providers.gcp.services.vertexai.vertexai_service import (
WorkbenchInstance,
)
return WorkbenchInstance(
id=_instance_id(name),
name=name,
project_id=GCP_PROJECT_ID,
location=_ZONE,
state=state,
upgrade_schedule=upgrade_schedule,
)
class Test_vertexai_workbench_instance_auto_upgrade_enabled:
def test_no_resources(self):
result = _run_check(_mock_vertexai_client())
assert len(result) == 0
def test_instance_auto_upgrade_enabled(self):
result = _run_check(
_mock_vertexai_client(
instances=[
_workbench_instance(
"patched-nb", upgrade_schedule=_UPGRADE_SCHEDULE
)
]
)
)
assert len(result) == 1
assert result[0].status == "PASS"
assert (
result[0].status_extended
== f"Vertex AI Workbench instance patched-nb has automatic upgrades scheduled ({_UPGRADE_SCHEDULE})."
)
assert result[0].resource_id == _instance_id("patched-nb")
assert result[0].resource_name == "patched-nb"
assert result[0].project_id == GCP_PROJECT_ID
assert result[0].location == _ZONE
def test_instance_auto_upgrade_disabled(self):
result = _run_check(
_mock_vertexai_client(instances=[_workbench_instance("stale-nb")])
)
assert len(result) == 1
assert result[0].status == "FAIL"
assert (
result[0].status_extended
== "Vertex AI Workbench instance stale-nb does not have automatic upgrades scheduled."
)
assert result[0].resource_id == _instance_id("stale-nb")
def test_deleted_instance_is_skipped(self):
result = _run_check(
_mock_vertexai_client(
instances=[_workbench_instance("gone-nb", state="DELETED")]
)
)
assert len(result) == 0
tests/providers/gcp/services/vertexai/vertexai_workbench_instance_file_download_disabled/vertexai_workbench_instance_file_download_disabled_test.py
+106
View File
@@ -0,0 +1,106 @@
from unittest import mock
from tests.providers.gcp.gcp_fixtures import GCP_PROJECT_ID, set_mocked_gcp_provider
_ZONE = "us-central1-a"
_CHECK_PATH = (
"prowler.providers.gcp.services.vertexai."
"vertexai_workbench_instance_file_download_disabled."
"vertexai_workbench_instance_file_download_disabled"
)
_CLIENT_PATH = f"{_CHECK_PATH}.vertexai_client"
def _instance_id(name: str) -> str:
return f"projects/{GCP_PROJECT_ID}/locations/{_ZONE}/instances/{name}"
def _mock_vertexai_client(instances=None):
vertexai_client = mock.MagicMock()
vertexai_client.workbench_instances = instances or []
return vertexai_client
def _run_check(vertexai_client):
with (
mock.patch(
"prowler.providers.common.provider.Provider.get_global_provider",
return_value=set_mocked_gcp_provider(),
),
mock.patch(_CLIENT_PATH, new=vertexai_client),
):
from prowler.providers.gcp.services.vertexai.vertexai_workbench_instance_file_download_disabled.vertexai_workbench_instance_file_download_disabled import (
vertexai_workbench_instance_file_download_disabled,
)
return vertexai_workbench_instance_file_download_disabled().execute()
def _workbench_instance(name, file_download_disabled, state="ACTIVE"):
from prowler.providers.gcp.services.vertexai.vertexai_service import (
WorkbenchInstance,
)
return WorkbenchInstance(
id=_instance_id(name),
name=name,
project_id=GCP_PROJECT_ID,
location=_ZONE,
state=state,
file_download_disabled=file_download_disabled,
)
class Test_vertexai_workbench_instance_file_download_disabled:
def test_no_resources(self):
result = _run_check(_mock_vertexai_client())
assert len(result) == 0
def test_instance_file_download_disabled(self):
result = _run_check(
_mock_vertexai_client(
instances=[
_workbench_instance("locked-nb", file_download_disabled=True)
]
)
)
assert len(result) == 1
assert result[0].status == "PASS"
assert (
result[0].status_extended
== "Vertex AI Workbench instance locked-nb has file downloads disabled."
)
assert result[0].resource_id == _instance_id("locked-nb")
assert result[0].resource_name == "locked-nb"
assert result[0].project_id == GCP_PROJECT_ID
assert result[0].location == _ZONE
def test_instance_file_download_enabled(self):
result = _run_check(
_mock_vertexai_client(
instances=[_workbench_instance("open-nb", file_download_disabled=False)]
)
)
assert len(result) == 1
assert result[0].status == "FAIL"
assert (
result[0].status_extended
== "Vertex AI Workbench instance open-nb has file downloads enabled."
)
assert result[0].resource_id == _instance_id("open-nb")
def test_deleted_instance_is_skipped(self):
result = _run_check(
_mock_vertexai_client(
instances=[
_workbench_instance(
"gone-nb", file_download_disabled=False, state="DELETED"
)
]
)
)
assert len(result) == 0
tests/providers/gcp/services/vertexai/vertexai_workbench_instance_terminal_access_disabled/vertexai_workbench_instance_terminal_access_disabled_test.py
+108
View File
@@ -0,0 +1,108 @@
from unittest import mock
from tests.providers.gcp.gcp_fixtures import GCP_PROJECT_ID, set_mocked_gcp_provider
_ZONE = "us-central1-a"
_CHECK_PATH = (
"prowler.providers.gcp.services.vertexai."
"vertexai_workbench_instance_terminal_access_disabled."
"vertexai_workbench_instance_terminal_access_disabled"
)
_CLIENT_PATH = f"{_CHECK_PATH}.vertexai_client"
def _instance_id(name: str) -> str:
return f"projects/{GCP_PROJECT_ID}/locations/{_ZONE}/instances/{name}"
def _mock_vertexai_client(instances=None):
vertexai_client = mock.MagicMock()
vertexai_client.workbench_instances = instances or []
return vertexai_client
def _run_check(vertexai_client):
with (
mock.patch(
"prowler.providers.common.provider.Provider.get_global_provider",
return_value=set_mocked_gcp_provider(),
),
mock.patch(_CLIENT_PATH, new=vertexai_client),
):
from prowler.providers.gcp.services.vertexai.vertexai_workbench_instance_terminal_access_disabled.vertexai_workbench_instance_terminal_access_disabled import (
vertexai_workbench_instance_terminal_access_disabled,
)
return vertexai_workbench_instance_terminal_access_disabled().execute()
def _workbench_instance(name, terminal_access_disabled, state="ACTIVE"):
from prowler.providers.gcp.services.vertexai.vertexai_service import (
WorkbenchInstance,
)
return WorkbenchInstance(
id=_instance_id(name),
name=name,
project_id=GCP_PROJECT_ID,
location=_ZONE,
state=state,
terminal_access_disabled=terminal_access_disabled,
)
class Test_vertexai_workbench_instance_terminal_access_disabled:
def test_no_resources(self):
result = _run_check(_mock_vertexai_client())
assert len(result) == 0
def test_instance_terminal_access_disabled(self):
result = _run_check(
_mock_vertexai_client(
instances=[
_workbench_instance("locked-nb", terminal_access_disabled=True)
]
)
)
assert len(result) == 1
assert result[0].status == "PASS"
assert (
result[0].status_extended
== "Vertex AI Workbench instance locked-nb has terminal access disabled."
)
assert result[0].resource_id == _instance_id("locked-nb")
assert result[0].resource_name == "locked-nb"
assert result[0].project_id == GCP_PROJECT_ID
assert result[0].location == _ZONE
def test_instance_terminal_access_enabled(self):
result = _run_check(
_mock_vertexai_client(
instances=[
_workbench_instance("shell-nb", terminal_access_disabled=False)
]
)
)
assert len(result) == 1
assert result[0].status == "FAIL"
assert (
result[0].status_extended
== "Vertex AI Workbench instance shell-nb has terminal access enabled."
)
assert result[0].resource_id == _instance_id("shell-nb")
def test_deleted_instance_is_skipped(self):
result = _run_check(
_mock_vertexai_client(
instances=[
_workbench_instance(
"gone-nb", terminal_access_disabled=False, state="DELETED"
)
]
)
)
assert len(result) == 0