Add support for Vertex AI notebooks to rpe lib.

cleardataeng · Nov 27, 2024 · 8025468 · 8025468
1 parent 313e88b
commit 8025468
Show file tree

Hide file tree

Showing 6 changed files with 185 additions and 0 deletions.
diff --git a/rpe/extractors/gcp_auditlogs.py b/rpe/extractors/gcp_auditlogs.py
@@ -443,4 +443,15 @@ def add_resource():
                 resource_data["repository"] = name_bits[5]
                 add_resource()
 
+        # Vertex AI Notebook
+        elif res_type == "audited_resource" and prop('resource.labels.service') == 'notebooks.googleapis.com':
+            name_bits = prop('protoPayload.resourceName').split('/')
+            resource_data = {
+                'name': name_bits[5],
+                'project_id': name_bits[1],
+                'location': name_bits[3],
+                'resource_type': 'notebooks.googleapis.com/Instance',
+            }
+            add_resource()
+
         return resources
diff --git a/rpe/resources/gcp.py b/rpe/resources/gcp.py
@@ -1146,3 +1146,29 @@ def _get_request_args(self):
 
     def _get_iam_request_args(self):
         return {"resource": self._get_resource_string()}
+
+
+class GcpNotebookInstance(GoogleAPIResource):
+    """
+    Vertex AI Workbench Instance.
+    """
+    service_name = "notebooks"
+    resource_path = "projects.locations.instances"
+    version = "v2"
+
+    required_resource_data = ["name", "location", "project_id"]
+
+    resource_type = "notebooks.googleapis.com/Instance"
+
+    inferred_data_map = {
+        "uniquifier": "id",
+    }
+
+    def _get_request_args(self):
+        return {
+            "name": "projects/{}/locations/{}/instances/{}".format(
+                self._resource_data["project_id"],
+                self._resource_data["location"],
+                self._resource_data["name"],
+            ),
+        }
diff --git a/tests/data/notebook_instance_create_01.json b/tests/data/notebook_instance_create_01.json
@@ -0,0 +1,66 @@
+{
+  "protoPayload": {
+    "@type": "type.googleapis.com/google.cloud.audit.AuditLog",
+    "status": {},
+    "authenticationInfo": {
+      "principalEmail": "[email protected]",
+      "principalSubject": "user:[email protected]"
+    },
+    "requestMetadata": {
+      "callerIp": "70.113.39.123",
+      "callerSuppliedUserAgent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36,gzip(gfe),gzip(gfe)",
+      "requestAttributes": {
+        "time": "2024-11-22T19:39:10.510275284Z",
+        "auth": {}
+      },
+      "destinationAttributes": {}
+    },
+    "serviceName": "notebooks.googleapis.com",
+    "methodName": "google.cloud.notebooks.v2.NotebookService.CreateInstance",
+    "authorizationInfo": [
+      {
+        "resource": "projects/cd-onboarding-graph-test-4/locations/us-central1-a/instances/test-create-event",
+        "permission": "notebooks.instances.create",
+        "granted": true,
+        "resourceAttributes": {
+          "service": "notebooks.googleapis.com",
+          "name": "projects/cd-onboarding-graph-test-4/locations/us-central1-a/instances/test-create-event"
+        },
+        "permissionType": "ADMIN_WRITE"
+      }
+    ],
+    "resourceName": "projects/cd-onboarding-graph-test-4/locations/us-central1-a/instances/test-create-event",
+    "request": {
+      "instance": {},
+      "parent": "projects/cd-onboarding-graph-test-4/locations/us-central1-a",
+      "@type": "type.googleapis.com/google.cloud.notebooks.v2.CreateInstanceRequest",
+      "instance_id": "test-create-event"
+    },
+    "response": {
+      "@type": "type.googleapis.com/google.longrunning.Operation"
+    },
+    "resourceLocation": {
+      "currentLocations": [
+        "us-central1-a"
+      ]
+    }
+  },
+  "insertId": "dlmvxecay0",
+  "resource": {
+    "type": "audited_resource",
+    "labels": {
+      "method": "google.cloud.notebooks.v2.NotebookService.CreateInstance",
+      "project_id": "cd-onboarding-graph-test-4",
+      "service": "notebooks.googleapis.com"
+    }
+  },
+  "timestamp": "2024-11-22T19:39:11.999291772Z",
+  "severity": "NOTICE",
+  "logName": "projects/cd-onboarding-graph-test-4/logs/cloudaudit.googleapis.com%2Factivity",
+  "operation": {
+    "id": "projects/cd-onboarding-graph-test-4/locations/us-central1-a/operations/operation-1732304350517-6278589749919-3cddbad0-d4e5f89e",
+    "producer": "notebooks.googleapis.com",
+    "first": true
+  },
+  "receiveTimestamp": "2024-11-22T19:39:12.712491438Z"
+}
diff --git a/tests/data/notebook_instance_update_01.json b/tests/data/notebook_instance_update_01.json
@@ -0,0 +1,62 @@
+{
+  "protoPayload": {
+    "@type": "type.googleapis.com/google.cloud.audit.AuditLog",
+    "status": {
+      "code": 9,
+      "message": "instance in state \"ACTIVE\" must be stopped before updating one of the following: [accelerator_configs machine_type shielded_instance_config reservation_affinity], failed precondition"
+    },
+    "authenticationInfo": {
+      "principalEmail": "[email protected]",
+      "principalSubject": "user:[email protected]"
+    },
+    "requestMetadata": {
+      "callerIp": "70.113.39.123",
+      "callerSuppliedUserAgent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36,gzip(gfe),gzip(gfe)",
+      "requestAttributes": {
+        "time": "2024-11-22T21:34:16.882098660Z",
+        "auth": {}
+      },
+      "destinationAttributes": {}
+    },
+    "serviceName": "notebooks.googleapis.com",
+    "methodName": "google.cloud.notebooks.v2.NotebookService.UpdateInstance",
+    "authorizationInfo": [
+      {
+        "resource": "projects/cd-onboarding-graph-test-4/locations/us-central1-a/instances/test-notebook",
+        "permission": "notebooks.instances.update",
+        "granted": true,
+        "resourceAttributes": {
+          "service": "notebooks.googleapis.com",
+          "name": "projects/cd-onboarding-graph-test-4/locations/us-central1-a/instances/test-notebook"
+        },
+        "permissionType": "ADMIN_WRITE"
+      }
+    ],
+    "resourceName": "projects/cd-onboarding-graph-test-4/locations/us-central1-a/instances/test-notebook",
+    "request": {
+      "@type": "type.googleapis.com/google.cloud.notebooks.v2.UpdateInstanceRequest",
+      "instance": {}
+    },
+    "response": {
+      "@type": "type.googleapis.com/google.longrunning.Operation"
+    },
+    "resourceLocation": {
+      "currentLocations": [
+        "us-central1-a"
+      ]
+    }
+  },
+  "insertId": "dlmvxeceen",
+  "resource": {
+    "type": "audited_resource",
+    "labels": {
+      "service": "notebooks.googleapis.com",
+      "project_id": "cd-onboarding-graph-test-4",
+      "method": "google.cloud.notebooks.v2.NotebookService.UpdateInstance"
+    }
+  },
+  "timestamp": "2024-11-22T21:34:16.891531432Z",
+  "severity": "ERROR",
+  "logName": "projects/cd-onboarding-graph-test-4/logs/cloudaudit.googleapis.com%2Factivity",
+  "receiveTimestamp": "2024-11-22T21:34:17.868613326Z"
+}
diff --git a/tests/test_extractors.py b/tests/test_extractors.py
@@ -226,6 +226,18 @@ def get_test_data(filename):
         "update",
         "test-repository",
     ),
+    (
+        "notebook_instance_create_01.json",
+        "notebooks.googleapis.com/Instance",
+        "create",
+        "test-create-event",
+    ),
+    (
+        "notebook_instance_update_01.json",
+        "notebooks.googleapis.com/Instance",
+        "update",
+        "test-notebook",
+    ),
 ]
 
 test_micromanager_log = [

diff --git a/tests/test_resources_cai.py b/tests/test_resources_cai.py
@@ -39,6 +39,7 @@
     GcpIamServiceAccount,
     GcpIamServiceAccountKey,
     GcpMemcacheInstance,
+    GcpNotebookInstance,
     GcpOrganization,
     GcpProject,
     GcpProjectService,
@@ -244,6 +245,13 @@
         },
         resource_cls=GcpDataformWorkspace,
     ),
+    CaiTestCase(
+        data={
+            "name": "//notebooks.googleapis.com/projects/test-001/locations/us-central1-a/instances/test-instance-01",
+            "asset_type": "notebooks.googleapis.com/Instance",
+        },
+        resource_cls=GcpNotebookInstance,
+    ),
 ]