mirror of
https://github.com/prowler-cloud/prowler.git
synced 2026-10-19 18:04:17 +00:00
feat(gcp): add vertexai service with workbench public IP and root access checks (#13062)
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
1 parent
de4db487d5
commit
136e4efdb5
14 files changed
+643
-13
No files matched your search
@@ -0,0 +1 @@
|
||||
Vertex AI service for GCP provider, including the `vertexai_workbench_instance_public_ip_disabled` and `vertexai_workbench_instance_root_access_disabled` checks mapped to PCI DSS 4.0
|
||||
@@ -701,7 +701,9 @@
|
||||
"Id": "1.2.8.37",
|
||||
"Description": "Checks if direct internet access is disabled for a Google AI Platform (formerly AI Platform Notebooks) instance. The rule is NON_COMPLIANT if an AI Platform notebook instance is internet-enabled.",
|
||||
"Name": "Google AI Platform",
|
||||
"Checks": [],
|
||||
"Checks": [
|
||||
"vertexai_workbench_instance_public_ip_disabled"
|
||||
],
|
||||
"Attributes": [
|
||||
{
|
||||
"Section": "1.2.8: Network security controls (NSCs) are configured and maintained.",
|
||||
@@ -1358,7 +1360,9 @@
|
||||
"Id": "1.3.1.41",
|
||||
"Description": "Checks if direct internet access is disabled for a Google AI Platform Notebook instance. The rule is NON_COMPLIANT if an AI Platform Notebook instance is internet-enabled.",
|
||||
"Name": "Google AI Platform",
|
||||
"Checks": [],
|
||||
"Checks": [
|
||||
"vertexai_workbench_instance_public_ip_disabled"
|
||||
],
|
||||
"Attributes": [
|
||||
{
|
||||
"Section": "1.3.1: Network access to and from the cardholder data environment is restricted.",
|
||||
@@ -2029,7 +2033,9 @@
|
||||
"Id": "1.3.2.41",
|
||||
"Description": "Checks if direct internet access is disabled for a Google AI Platform (formerly known as AI Platform Notebooks) instance. The rule is NON_COMPLIANT if an AI Platform Notebook instance is internet-enabled.",
|
||||
"Name": "Google AI Platform",
|
||||
"Checks": [],
|
||||
"Checks": [
|
||||
"vertexai_workbench_instance_public_ip_disabled"
|
||||
],
|
||||
"Attributes": [
|
||||
{
|
||||
"Section": "1.3.2: Network access to and from the cardholder data environment is restricted.",
|
||||
@@ -2805,7 +2811,9 @@
|
||||
"Id": "1.4.2.39",
|
||||
"Description": "Checks if direct internet access is disabled for a Google Cloud AI Platform Notebooks instance. The rule is NON_COMPLIANT if an AI Platform Notebooks instance is internet-enabled.",
|
||||
"Name": "Google Cloud AI Platform Notebooks",
|
||||
"Checks": [],
|
||||
"Checks": [
|
||||
"vertexai_workbench_instance_public_ip_disabled"
|
||||
],
|
||||
"Attributes": [
|
||||
{
|
||||
"Section": "1.4.2: Network connections between trusted and untrusted networks are controlled.",
|
||||
@@ -3616,7 +3624,9 @@
|
||||
"Id": "1.5.1.37",
|
||||
"Description": "Checks if direct internet access is disabled for a Google Cloud Vertex AI Workbench instance. The rule is NON_COMPLIANT if the Vertex AI Workbench instance is internet-enabled.",
|
||||
"Name": "Google Cloud Vertex AI Workbench",
|
||||
"Checks": [],
|
||||
"Checks": [
|
||||
"vertexai_workbench_instance_public_ip_disabled"
|
||||
],
|
||||
"Attributes": [
|
||||
{
|
||||
"Section": "1.5.1: Risks to the CDE from computing devices that are able to connect to both untrusted networks and the CDE are mitigated.",
|
||||
@@ -8491,7 +8501,9 @@
|
||||
"Id": "10.3.2.25",
|
||||
"Description": "Checks if direct internet access is disabled for a Google AI Platform Notebook instance. The rule is NON_COMPLIANT if an AI Platform Notebook instance is internet-enabled.",
|
||||
"Name": "Google AI Platform Notebooks",
|
||||
"Checks": [],
|
||||
"Checks": [
|
||||
"vertexai_workbench_instance_public_ip_disabled"
|
||||
],
|
||||
"Attributes": [
|
||||
{
|
||||
"Section": "10.3.2: Audit logs are protected from destruction and unauthorized modifications.",
|
||||
@@ -12122,7 +12134,9 @@
|
||||
"Id": "3.5.1.3.30",
|
||||
"Description": "Checks if direct internet access is disabled for a Google Cloud AI Platform Notebooks instance. The rule is NON_COMPLIANT if an AI Platform Notebooks instance is internet-enabled.",
|
||||
"Name": "Google Cloud AI Platform Notebooks",
|
||||
"Checks": [],
|
||||
"Checks": [
|
||||
"vertexai_workbench_instance_public_ip_disabled"
|
||||
],
|
||||
"Attributes": [
|
||||
{
|
||||
"Section": "3.5.1.3: Primary account number (PAN) is secured wherever it is stored.",
|
||||
@@ -15753,7 +15767,9 @@
|
||||
"Id": "7.2.1.28",
|
||||
"Description": "Checks if the Google AI Platform Notebooks Root Access setting is enabled for AI Platform Notebook instances. The rule is NON_COMPLIANT if the Root Access setting is set to 'Enabled' for an AI Platform Notebook instance.",
|
||||
"Name": "Google AI Platform Notebooks",
|
||||
"Checks": [],
|
||||
"Checks": [
|
||||
"vertexai_workbench_instance_root_access_disabled"
|
||||
],
|
||||
"Attributes": [
|
||||
{
|
||||
"Section": "7.2.1: Access to system components and data is appropriately defined and assigned.",
|
||||
@@ -16153,7 +16169,9 @@
|
||||
"Id": "7.2.2.28",
|
||||
"Description": "Checks if the Google AI Platform Notebook instance has root access enabled. The rule is NON_COMPLIANT if the root access is set to 'Enabled' for a Google AI Platform Notebook instance.",
|
||||
"Name": "Google AI Platform",
|
||||
"Checks": [],
|
||||
"Checks": [
|
||||
"vertexai_workbench_instance_root_access_disabled"
|
||||
],
|
||||
"Attributes": [
|
||||
{
|
||||
"Section": "7.2.2: Access to system components and data is appropriately defined and assigned.",
|
||||
@@ -16282,7 +16300,9 @@
|
||||
"Id": "7.2.3.9",
|
||||
"Description": "Checks if the Vertex AI Workbench Root Access setting is enabled for Vertex AI Workbench Notebooks. The rule is NON_COMPLIANT if the Root Access setting is set to 'Enabled' for a Vertex AI Workbench Notebook instance.",
|
||||
"Name": "Vertex AI",
|
||||
"Checks": [],
|
||||
"Checks": [
|
||||
"vertexai_workbench_instance_root_access_disabled"
|
||||
],
|
||||
"Attributes": [
|
||||
{
|
||||
"Section": "7.2.3: Access to system components and data is appropriately defined and assigned.",
|
||||
@@ -20547,7 +20567,9 @@
|
||||
"Id": "A1.1.2.21",
|
||||
"Description": "Checks if direct internet access is disabled for a Google AI Platform Notebook instance. The rule is NON_COMPLIANT if an AI Platform Notebook instance is internet-enabled.",
|
||||
"Name": "Google AI Platform Notebooks",
|
||||
"Checks": [],
|
||||
"Checks": [
|
||||
"vertexai_workbench_instance_public_ip_disabled"
|
||||
],
|
||||
"Attributes": [
|
||||
{
|
||||
"Section": "A1.1.2: Multi-tenant service providers protect and separate all customer environments and data.",
|
||||
@@ -21060,7 +21082,9 @@
|
||||
"Id": "A1.1.3.37",
|
||||
"Description": "Checks if direct internet access is disabled for a Google Cloud AI Platform Notebook instance. The rule is NON_COMPLIANT if an AI Platform Notebook instance is internet-enabled.",
|
||||
"Name": "Google Cloud AI Platform",
|
||||
"Checks": [],
|
||||
"Checks": [
|
||||
"vertexai_workbench_instance_public_ip_disabled"
|
||||
],
|
||||
"Attributes": [
|
||||
{
|
||||
"Section": "A1.1.3: Multi-tenant service providers protect and separate all customer environments and data.",
|
||||
@@ -22180,7 +22204,9 @@
|
||||
"Id": "A3.4.1.23",
|
||||
"Description": "Checks if direct internet access is disabled for a Google Cloud AI Platform Notebook instance. The rule is NON_COMPLIANT if an AI Platform Notebook instance is internet-enabled.",
|
||||
"Name": "Google Cloud AI Platform Notebooks",
|
||||
"Checks": [],
|
||||
"Checks": [
|
||||
"vertexai_workbench_instance_public_ip_disabled"
|
||||
],
|
||||
"Attributes": [
|
||||
{
|
||||
"Section": "A3.4.1: Logical access to the cardholder data environment is controlled and managed.",
|
||||
|
||||
Whitespace-only changes.
@@ -0,0 +1,4 @@
|
||||
from prowler.providers.common.provider import Provider
|
||||
from prowler.providers.gcp.services.vertexai.vertexai_service import VertexAI
|
||||
|
||||
vertexai_client = VertexAI(Provider.get_global_provider())
|
||||
@@ -0,0 +1,89 @@
|
||||
from typing import Optional
|
||||
|
||||
from pydantic.v1 import BaseModel
|
||||
|
||||
from prowler.lib.logger import logger
|
||||
from prowler.providers.gcp.config import DEFAULT_RETRY_ATTEMPTS
|
||||
from prowler.providers.gcp.gcp_provider import GcpProvider
|
||||
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"
|
||||
|
||||
|
||||
class VertexAI(GCPService):
|
||||
"""Vertex AI service client.
|
||||
|
||||
Enumerates Vertex AI Workbench instances across every accessible project
|
||||
using the `notebooks.googleapis.com` v2 API and exposes them through the
|
||||
`workbench_instances` attribute.
|
||||
"""
|
||||
|
||||
def __init__(self, provider: GcpProvider) -> None:
|
||||
"""Initialize the service and preload Workbench instances."""
|
||||
super().__init__("notebooks", provider, api_version="v2")
|
||||
self.workbench_instances = []
|
||||
self._get_workbench_instances()
|
||||
|
||||
def _get_workbench_instances(self) -> None:
|
||||
"""Fetch Workbench instances for every project across all locations."""
|
||||
for project_id in self.project_ids:
|
||||
try:
|
||||
instances = self.client.projects().locations().instances()
|
||||
request = instances.list(parent=f"projects/{project_id}/locations/-")
|
||||
while request is not None:
|
||||
response = request.execute(num_retries=DEFAULT_RETRY_ATTEMPTS)
|
||||
for instance in response.get("instances", []):
|
||||
gce_setup = instance.get("gceSetup", {})
|
||||
metadata = gce_setup.get("metadata", {})
|
||||
network_interfaces = gce_setup.get("networkInterfaces", [])
|
||||
# Name format: projects/{project}/locations/{zone}/instances/{instance}
|
||||
name_parts = instance["name"].split("/")
|
||||
self.workbench_instances.append(
|
||||
WorkbenchInstance(
|
||||
id=instance["name"],
|
||||
name=name_parts[-1],
|
||||
project_id=project_id,
|
||||
location=name_parts[3],
|
||||
state=instance.get("state", "STATE_UNSPECIFIED"),
|
||||
public_ip_disabled=gce_setup.get(
|
||||
"disablePublicIp", False
|
||||
),
|
||||
root_access_disabled=str(
|
||||
metadata.get(ROOT_ACCESS_DISABLED_METADATA_KEY, "")
|
||||
).lower()
|
||||
== "true",
|
||||
network=(
|
||||
network_interfaces[0].get("network")
|
||||
if network_interfaces
|
||||
else None
|
||||
),
|
||||
)
|
||||
)
|
||||
# The wildcard list reports locations it could not reach
|
||||
# instead of raising, so the project was not fully scanned.
|
||||
unreachable = response.get("unreachable", [])
|
||||
if unreachable:
|
||||
logger.error(
|
||||
f"{project_id} -- Workbench instances could not be listed in unreachable locations: {', '.join(unreachable)}"
|
||||
)
|
||||
request = instances.list_next(
|
||||
previous_request=request, previous_response=response
|
||||
)
|
||||
except Exception as error:
|
||||
logger.error(
|
||||
f"{project_id} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
|
||||
)
|
||||
|
||||
|
||||
class WorkbenchInstance(BaseModel):
|
||||
"""Vertex AI Workbench instance resource consumed by GCP checks."""
|
||||
|
||||
id: str
|
||||
name: str
|
||||
project_id: str
|
||||
location: str
|
||||
state: str
|
||||
public_ip_disabled: bool = False
|
||||
root_access_disabled: bool = False
|
||||
network: Optional[str] = None
|
||||
Whitespace-only changes.
+40
@@ -0,0 +1,40 @@
|
||||
{
|
||||
"Provider": "gcp",
|
||||
"CheckID": "vertexai_workbench_instance_public_ip_disabled",
|
||||
"CheckTitle": "Vertex AI Workbench instance does not have a public IP address",
|
||||
"CheckType": [],
|
||||
"ServiceName": "vertexai",
|
||||
"SubServiceName": "workbench",
|
||||
"ResourceIdTemplate": "",
|
||||
"Severity": "high",
|
||||
"ResourceType": "notebooks.googleapis.com/Instance",
|
||||
"ResourceGroup": "ai_ml",
|
||||
"Description": "**Vertex AI Workbench instances** are JupyterLab VMs used to explore data, train models, and call generative AI APIs. The finding determines whether an instance was created with **public IP disabled** (`disablePublicIp`), so it has no external address and reaches the internet only through controlled egress such as **Cloud NAT** or **Private Google Access**.",
|
||||
"Risk": "A notebook VM with a public IP has **direct internet access**:\n- **Confidentiality**: training data, models, and service account tokens can be exfiltrated around VPC egress controls\n- **Integrity**: untrusted packages and models are pulled straight into an environment holding project credentials\n- **Exposure**: the Jupyter stack is reachable for scanning and exploitation",
|
||||
"RelatedUrl": "",
|
||||
"AdditionalURLs": [
|
||||
"https://cloud.google.com/vertex-ai/docs/workbench/instances/create",
|
||||
"https://cloud.google.com/vertex-ai/docs/workbench/instances/introduction"
|
||||
],
|
||||
"Remediation": {
|
||||
"Code": {
|
||||
"CLI": "curl -X PATCH -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" -d '{\"gceSetup\": {\"disablePublicIp\": true}}' \"https://notebooks.googleapis.com/v2/projects/<example_project_id>/locations/<example_zone>/instances/<example_resource_name>?updateMask=gceSetup.disable_public_ip\"",
|
||||
"NativeIaC": "",
|
||||
"Other": "Update the existing instance in place; `gcloud workbench instances update` has no public IP flag, so use the Notebooks API.\n1. Confirm the instance subnetwork has Private Google Access or Cloud NAT, so the notebook keeps access to Google APIs and package repositories\n2. Call `instances.patch` with `updateMask=gceSetup.disable_public_ip` and `gceSetup.disablePublicIp` set to `true`, as in the CLI command\n3. Wait for the returned operation to finish, then confirm in Compute Engine > VM instances that the instance VM has no external IP\n4. Only if the update is rejected, create a new instance with `gcloud workbench instances create --disable-public-ip`, migrate notebooks and data, and delete the old instance",
|
||||
"Terraform": "```hcl\nresource \"google_workbench_instance\" \"<example_resource_name>\" {\n name = \"<example_resource_name>\"\n location = \"<example_zone>\"\n\n gce_setup {\n disable_public_ip = true # FIX: No external IP is assigned to the notebook VM\n\n network_interfaces {\n network = \"<example_network_id>\"\n subnet = \"<example_subnet_id>\"\n }\n }\n}\n```"
|
||||
},
|
||||
"Recommendation": {
|
||||
"Text": "Create Workbench instances **without external IPs** and keep them on private networks.\n- Use **Cloud NAT** for controlled egress and **Private Google Access** for Google APIs\n- Reach JupyterLab through the IAM-authenticated **Workbench proxy** or **IAP**\n- Place AI projects inside a **VPC Service Controls** perimeter\n- Enforce the `constraints/compute.vmExternalIpAccess` organization policy",
|
||||
"Url": "https://hub.prowler.com/check/vertexai_workbench_instance_public_ip_disabled"
|
||||
}
|
||||
},
|
||||
"Categories": [
|
||||
"internet-exposed",
|
||||
"gen-ai"
|
||||
],
|
||||
"DependsOn": [],
|
||||
"RelatedTo": [
|
||||
"compute_instance_public_ip"
|
||||
],
|
||||
"Notes": "Workbench instances are backed by Compute Engine VMs, which compute_instance_public_ip also evaluates; this check reports the finding against the Workbench resource that owners manage."
|
||||
}
|
||||
+33
@@ -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_public_ip_disabled(Check):
|
||||
"""Ensure Vertex AI Workbench instances do not have a public IP address.
|
||||
|
||||
A Workbench instance with an external IP is directly reachable from, and can
|
||||
reach, the internet without passing through VPC egress controls.
|
||||
- PASS: The instance was created with `disablePublicIp` set.
|
||||
- FAIL: The instance has a public IP address assigned.
|
||||
"""
|
||||
|
||||
def execute(self) -> list[Check_Report_GCP]:
|
||||
"""Execute the public IP 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.public_ip_disabled:
|
||||
report.status = "PASS"
|
||||
report.status_extended = f"Vertex AI Workbench instance {instance.name} does not have a public IP address."
|
||||
else:
|
||||
report.status = "FAIL"
|
||||
report.status_extended = f"Vertex AI Workbench instance {instance.name} has a public IP address assigned, allowing direct internet access."
|
||||
findings.append(report)
|
||||
return findings
|
||||
Whitespace-only changes.
+38
@@ -0,0 +1,38 @@
|
||||
{
|
||||
"Provider": "gcp",
|
||||
"CheckID": "vertexai_workbench_instance_root_access_disabled",
|
||||
"CheckTitle": "Vertex AI Workbench instance has root access disabled",
|
||||
"CheckType": [],
|
||||
"ServiceName": "vertexai",
|
||||
"SubServiceName": "workbench",
|
||||
"ResourceIdTemplate": "",
|
||||
"Severity": "medium",
|
||||
"ResourceType": "notebooks.googleapis.com/Instance",
|
||||
"ResourceGroup": "ai_ml",
|
||||
"Description": "**Vertex AI Workbench instances** give the JupyterLab user **passwordless sudo** by default. The finding determines whether root access is disabled by setting the `notebook-disable-root` instance metadata key to `true`.",
|
||||
"Risk": "With root access, a notebook user, or any malicious package, model, or notebook they run, can:\n- **Integrity**: tamper with the OS, Jupyter server, and libraries, or install persistent backdoors\n- **Confidentiality**: read every user's files and credentials on the VM\n- **Defense evasion**: disable monitoring agents and audit tooling",
|
||||
"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-root=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-root with the value true\n4. Click Submit and restart the instance if prompted",
|
||||
"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-root = \"true\" # FIX: Removes passwordless sudo for the notebook user\n }\n }\n}\n```"
|
||||
},
|
||||
"Recommendation": {
|
||||
"Text": "Disable root access on Workbench instances and apply **least privilege** to notebook users.\n- Install required packages in user space (`pip install --user`, virtual environments, or custom containers)\n- Build approved images for system-level dependencies instead of granting sudo\n- Attach a dedicated, minimally scoped **service account** to each instance",
|
||||
"Url": "https://hub.prowler.com/check/vertexai_workbench_instance_root_access_disabled"
|
||||
}
|
||||
},
|
||||
"Categories": [
|
||||
"identity-access",
|
||||
"gen-ai"
|
||||
],
|
||||
"DependsOn": [],
|
||||
"RelatedTo": [],
|
||||
"Notes": ""
|
||||
}
|
||||
+33
@@ -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_root_access_disabled(Check):
|
||||
"""Ensure Vertex AI Workbench instances have root access disabled.
|
||||
|
||||
Workbench grants the notebook user passwordless sudo unless the
|
||||
`notebook-disable-root` metadata key is set to `true`.
|
||||
- PASS: The instance has `notebook-disable-root` set to `true`.
|
||||
- FAIL: The notebook user can obtain root on the instance.
|
||||
"""
|
||||
|
||||
def execute(self) -> list[Check_Report_GCP]:
|
||||
"""Execute the root 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.root_access_disabled:
|
||||
report.status = "PASS"
|
||||
report.status_extended = f"Vertex AI Workbench instance {instance.name} has root access disabled."
|
||||
else:
|
||||
report.status = "FAIL"
|
||||
report.status_extended = f"Vertex AI Workbench instance {instance.name} has root access enabled."
|
||||
findings.append(report)
|
||||
return findings
|
||||
@@ -0,0 +1,158 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from prowler.providers.gcp.services.vertexai.vertexai_service import VertexAI
|
||||
from tests.providers.gcp.gcp_fixtures import (
|
||||
GCP_PROJECT_ID,
|
||||
mock_is_api_active,
|
||||
set_mocked_gcp_provider,
|
||||
)
|
||||
|
||||
_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"
|
||||
|
||||
|
||||
def _make_notebooks_client(instances_list=None, list_error=None):
|
||||
"""Return a mock Notebooks v2 API client for Workbench instances."""
|
||||
client = MagicMock()
|
||||
instances = client.projects().locations().instances()
|
||||
if list_error:
|
||||
instances.list().execute.side_effect = list_error
|
||||
else:
|
||||
instances.list().execute.return_value = {"instances": instances_list or []}
|
||||
instances.list_next.return_value = None
|
||||
return client
|
||||
|
||||
|
||||
def _build_service(client):
|
||||
with (
|
||||
patch(
|
||||
"prowler.providers.gcp.lib.service.service.GCPService.__is_api_active__",
|
||||
new=mock_is_api_active,
|
||||
),
|
||||
patch(
|
||||
"prowler.providers.gcp.lib.service.service.GCPService.__generate_client__",
|
||||
new=lambda *_: client,
|
||||
),
|
||||
):
|
||||
return VertexAI(set_mocked_gcp_provider(project_ids=[GCP_PROJECT_ID]))
|
||||
|
||||
|
||||
class TestVertexAIService:
|
||||
def test_service_uses_notebooks_v2_api(self):
|
||||
vertexai_client = _build_service(_make_notebooks_client())
|
||||
|
||||
assert vertexai_client.service == "notebooks"
|
||||
assert vertexai_client.api_version == "v2"
|
||||
assert vertexai_client.project_ids == [GCP_PROJECT_ID]
|
||||
assert vertexai_client.workbench_instances == []
|
||||
|
||||
def test_lists_instances_across_all_locations(self):
|
||||
client = _make_notebooks_client()
|
||||
_build_service(client)
|
||||
|
||||
client.projects().locations().instances().list.assert_called_with(
|
||||
parent=f"projects/{GCP_PROJECT_ID}/locations/-"
|
||||
)
|
||||
|
||||
def test_get_workbench_instance_hardened(self):
|
||||
vertexai_client = _build_service(
|
||||
_make_notebooks_client(
|
||||
instances_list=[
|
||||
{
|
||||
"name": _INSTANCE_ID,
|
||||
"state": "ACTIVE",
|
||||
"gceSetup": {
|
||||
"disablePublicIp": True,
|
||||
"metadata": {"notebook-disable-root": "true"},
|
||||
"networkInterfaces": [{"network": _NETWORK}],
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
assert len(vertexai_client.workbench_instances) == 1
|
||||
instance = vertexai_client.workbench_instances[0]
|
||||
assert instance.id == _INSTANCE_ID
|
||||
assert instance.name == _INSTANCE_NAME
|
||||
assert instance.project_id == GCP_PROJECT_ID
|
||||
assert instance.location == _ZONE
|
||||
assert instance.state == "ACTIVE"
|
||||
assert instance.public_ip_disabled is True
|
||||
assert instance.root_access_disabled is True
|
||||
assert instance.network == _NETWORK
|
||||
|
||||
def test_get_workbench_instance_defaults(self):
|
||||
vertexai_client = _build_service(
|
||||
_make_notebooks_client(
|
||||
instances_list=[{"name": _INSTANCE_ID, "gceSetup": {}}]
|
||||
)
|
||||
)
|
||||
|
||||
instance = vertexai_client.workbench_instances[0]
|
||||
assert instance.state == "STATE_UNSPECIFIED"
|
||||
assert instance.public_ip_disabled is False
|
||||
assert instance.root_access_disabled is False
|
||||
assert instance.network is None
|
||||
|
||||
def test_root_access_metadata_is_case_insensitive(self):
|
||||
vertexai_client = _build_service(
|
||||
_make_notebooks_client(
|
||||
instances_list=[
|
||||
{
|
||||
"name": _INSTANCE_ID,
|
||||
"gceSetup": {"metadata": {"notebook-disable-root": "TRUE"}},
|
||||
}
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
assert vertexai_client.workbench_instances[0].root_access_disabled is True
|
||||
|
||||
def test_root_access_metadata_false(self):
|
||||
vertexai_client = _build_service(
|
||||
_make_notebooks_client(
|
||||
instances_list=[
|
||||
{
|
||||
"name": _INSTANCE_ID,
|
||||
"gceSetup": {"metadata": {"notebook-disable-root": "false"}},
|
||||
}
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
assert vertexai_client.workbench_instances[0].root_access_disabled is False
|
||||
|
||||
def test_unreachable_locations_return_no_instances(self):
|
||||
client = _make_notebooks_client()
|
||||
client.projects().locations().instances().list().execute.return_value = {
|
||||
"unreachable": ["us-west1-a", "us-central1-b"]
|
||||
}
|
||||
|
||||
vertexai_client = _build_service(client)
|
||||
|
||||
assert vertexai_client.workbench_instances == []
|
||||
|
||||
def test_unreachable_on_later_page_keeps_listed_instances(self):
|
||||
client = _make_notebooks_client(
|
||||
instances_list=[{"name": _INSTANCE_ID, "gceSetup": {}}]
|
||||
)
|
||||
second_page = MagicMock()
|
||||
second_page.execute.return_value = {"unreachable": ["us-west1-a"]}
|
||||
client.projects().locations().instances().list_next.side_effect = [
|
||||
second_page,
|
||||
None,
|
||||
]
|
||||
|
||||
vertexai_client = _build_service(client)
|
||||
|
||||
assert [i.name for i in vertexai_client.workbench_instances] == [_INSTANCE_NAME]
|
||||
|
||||
def test_list_error_returns_no_instances(self):
|
||||
vertexai_client = _build_service(
|
||||
_make_notebooks_client(list_error=Exception("permission denied"))
|
||||
)
|
||||
|
||||
assert vertexai_client.workbench_instances == []
|
||||
+104
@@ -0,0 +1,104 @@
|
||||
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_public_ip_disabled."
|
||||
"vertexai_workbench_instance_public_ip_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_public_ip_disabled.vertexai_workbench_instance_public_ip_disabled import (
|
||||
vertexai_workbench_instance_public_ip_disabled,
|
||||
)
|
||||
|
||||
return vertexai_workbench_instance_public_ip_disabled().execute()
|
||||
|
||||
|
||||
def _workbench_instance(name, public_ip_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,
|
||||
public_ip_disabled=public_ip_disabled,
|
||||
)
|
||||
|
||||
|
||||
class Test_vertexai_workbench_instance_public_ip_disabled:
|
||||
def test_no_instances(self):
|
||||
result = _run_check(_mock_vertexai_client())
|
||||
|
||||
assert len(result) == 0
|
||||
|
||||
def test_instance_public_ip_disabled(self):
|
||||
result = _run_check(
|
||||
_mock_vertexai_client(
|
||||
instances=[_workbench_instance("private-nb", public_ip_disabled=True)]
|
||||
)
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].status == "PASS"
|
||||
assert (
|
||||
result[0].status_extended
|
||||
== "Vertex AI Workbench instance private-nb does not have a public IP address."
|
||||
)
|
||||
assert result[0].resource_id == _instance_id("private-nb")
|
||||
assert result[0].resource_name == "private-nb"
|
||||
assert result[0].project_id == GCP_PROJECT_ID
|
||||
assert result[0].location == _ZONE
|
||||
|
||||
def test_instance_public_ip_enabled(self):
|
||||
result = _run_check(
|
||||
_mock_vertexai_client(
|
||||
instances=[_workbench_instance("public-nb", public_ip_disabled=False)]
|
||||
)
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].status == "FAIL"
|
||||
assert (
|
||||
result[0].status_extended
|
||||
== "Vertex AI Workbench instance public-nb has a public IP address assigned, allowing direct internet access."
|
||||
)
|
||||
assert result[0].resource_id == _instance_id("public-nb")
|
||||
|
||||
def test_deleted_instance_is_skipped(self):
|
||||
result = _run_check(
|
||||
_mock_vertexai_client(
|
||||
instances=[
|
||||
_workbench_instance(
|
||||
"gone-nb", public_ip_disabled=False, state="DELETED"
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
assert len(result) == 0
|
||||
+104
@@ -0,0 +1,104 @@
|
||||
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_root_access_disabled."
|
||||
"vertexai_workbench_instance_root_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_root_access_disabled.vertexai_workbench_instance_root_access_disabled import (
|
||||
vertexai_workbench_instance_root_access_disabled,
|
||||
)
|
||||
|
||||
return vertexai_workbench_instance_root_access_disabled().execute()
|
||||
|
||||
|
||||
def _workbench_instance(name, root_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,
|
||||
root_access_disabled=root_access_disabled,
|
||||
)
|
||||
|
||||
|
||||
class Test_vertexai_workbench_instance_root_access_disabled:
|
||||
def test_no_instances(self):
|
||||
result = _run_check(_mock_vertexai_client())
|
||||
|
||||
assert len(result) == 0
|
||||
|
||||
def test_instance_root_access_disabled(self):
|
||||
result = _run_check(
|
||||
_mock_vertexai_client(
|
||||
instances=[_workbench_instance("locked-nb", root_access_disabled=True)]
|
||||
)
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].status == "PASS"
|
||||
assert (
|
||||
result[0].status_extended
|
||||
== "Vertex AI Workbench instance locked-nb has root 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_root_access_enabled(self):
|
||||
result = _run_check(
|
||||
_mock_vertexai_client(
|
||||
instances=[_workbench_instance("sudo-nb", root_access_disabled=False)]
|
||||
)
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].status == "FAIL"
|
||||
assert (
|
||||
result[0].status_extended
|
||||
== "Vertex AI Workbench instance sudo-nb has root access enabled."
|
||||
)
|
||||
assert result[0].resource_id == _instance_id("sudo-nb")
|
||||
|
||||
def test_deleted_instance_is_skipped(self):
|
||||
result = _run_check(
|
||||
_mock_vertexai_client(
|
||||
instances=[
|
||||
_workbench_instance(
|
||||
"gone-nb", root_access_disabled=False, state="DELETED"
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
assert len(result) == 0
|
||||
Reference in new issue
Block a user