[RW-13914][RW-13916][risk=no] Handle vwb egress alert (1st part) (#8990)

* draft * add tests * fix test * fix test * fix test * fix test * fix test * fix test * fix test * fix test
all-of-us · Dec 18, 2024 · 46f30c2 · 46f30c2
1 parent fcac632
commit 46f30c2
Show file tree

Hide file tree

Showing 21 changed files with 661 additions and 74 deletions.
diff --git a/api/db/changelog/db.changelog-246-add-vwb-egress-event-columns.xml b/api/db/changelog/db.changelog-246-add-vwb-egress-event-columns.xml
@@ -0,0 +1,16 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<databaseChangeLog
+        xmlns="http://www.liquibase.org/xml/ns/dbchangelog"
+        xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+        xsi:schemaLocation="http://www.liquibase.org/xml/ns/dbchangelog
+                      http://www.liquibase.org/xml/ns/dbchangelog/dbchangelog-3.1.xsd">
+    <changeSet author="yonghao" id="db.changelog-246-add-is-vwb-egress-event-column">
+        <addColumn tableName="egress_event">
+            <column name="vwb_workspace_id" type="varchar(80)" />
+            <column name="vwb_vm_name" type="varchar(80)" />
+            <column name="gcp_project_id" type="varchar(80)" />
+            <column name="vwb_incident_count" type="bigint" />
+            <column name="vwb_egress_event_id" type="varchar(80)" />
+        </addColumn>
+    </changeSet>
+</databaseChangeLog>
diff --git a/api/db/changelog/db.changelog-master.xml b/api/db/changelog/db.changelog-master.xml
@@ -253,6 +253,7 @@
   <include file="changelog/db.changelog-243-replace-extension-count.xml"/>
   <include file="changelog/db.changelog-244-add-is-vwb-workspace-column.xml"/>
   <include file="changelog/db.changelog-245-add-vwb-template-id-cdr-version.xml"/>
+  <include file="changelog/db.changelog-246-add-vwb-egress-event-columns.xml"/>
   <!--
    Note: to update the DB locally, do the following:
    - Migrate schema changes: `./project.rb run-local-all-migrations`

diff --git a/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditor.java b/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditor.java
@@ -5,6 +5,7 @@
 import org.pmiops.workbench.db.model.DbUser;
 import org.pmiops.workbench.model.SumologicEgressEvent;
 import org.pmiops.workbench.model.SumologicEgressEventRequest;
+import org.pmiops.workbench.model.VwbEgressEventRequest;
 
 /**
  * Auditor service which handles collecting audit logs for high-egress events. These events are
@@ -18,6 +19,12 @@ public interface EgressEventAuditor {
    */
   void fireEgressEvent(SumologicEgressEvent event);
 
+  /**
+   * Decorates a Verily Workbench high-egress event with Workbench metadata and fires an audit event
+   * log in the target workspace.
+   */
+  void fireVwbEgressEvent(VwbEgressEventRequest event, DbUser dbUser);
+
   /**
    * Decorates a Sumologic-reported high-egress event with Workbench metadata and fires an audit
    * event log in the target workspace for the specified user.

diff --git a/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditorImpl.java b/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditorImpl.java
@@ -29,6 +29,7 @@
 import org.pmiops.workbench.model.SumologicEgressEvent;
 import org.pmiops.workbench.model.SumologicEgressEventRequest;
 import org.pmiops.workbench.model.UserRole;
+import org.pmiops.workbench.model.VwbEgressEventRequest;
 import org.pmiops.workbench.workspaces.WorkspaceService;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
@@ -109,6 +110,35 @@ public void fireEgressEvent(SumologicEgressEvent event) {
     fireEvent(agentId, agentEmail, dbWorkspace, event);
   }
 
+  @Override
+  public void fireVwbEgressEvent(VwbEgressEventRequest event, DbUser user) {
+    String actionId = actionIdProvider.get();
+    Builder baseEventBuilder =
+        ActionAuditEvent.builder()
+            .timestamp(clock.millis())
+            .actionId(actionId)
+            .actionType(ActionType.DETECT_HIGH_EGRESS_EVENT)
+            .agentType(AgentType.USER)
+            .agentIdMaybe(user.getUserId())
+            .agentEmailMaybe(user.getUsername())
+            .targetType(TargetType.WORKSPACE);
+    var events = new ArrayList<ActionAuditEvent>();
+    // File 3 events: Workspace, VM, GCP projects. Can add more if needed in the future
+    events.add(
+        baseEventBuilder
+            .targetPropertyMaybe("vwbWorkspaceId")
+            .newValueMaybe(event.getVwbWorkspaceId())
+            .build());
+    events.add(
+        baseEventBuilder
+            .targetPropertyMaybe("gcpProjectId")
+            .newValueMaybe(event.getGcpProjectId())
+            .build());
+    events.add(
+        baseEventBuilder.targetPropertyMaybe("vmName").newValueMaybe(event.getVmName()).build());
+    actionAuditService.send(events);
+  }
+
   @Override
   public void fireEgressEventForUser(SumologicEgressEvent event, DbUser user) {
     DbWorkspace dbWorkspace = getDbWorkspace(event);

diff --git a/api/src/main/java/org/pmiops/workbench/api/CloudTaskEgressController.java b/api/src/main/java/org/pmiops/workbench/api/CloudTaskEgressController.java
@@ -1,24 +1,36 @@
 package org.pmiops.workbench.api;
 
 import static org.pmiops.workbench.exfiltration.ExfiltrationUtils.EGRESS_SUMOLOGIC_SERVICE_QUALIFIER;
+import static org.pmiops.workbench.exfiltration.ExfiltrationUtils.EGRESS_VWB_SERVICE_QUALIFIER;
 
 import org.pmiops.workbench.exfiltration.EgressRemediationService;
 import org.pmiops.workbench.model.ProcessEgressEventRequest;
-import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.RestController;
 
 @RestController
 public class CloudTaskEgressController implements CloudTaskEgressApiDelegate {
 
-  @Autowired
-  @Qualifier(EGRESS_SUMOLOGIC_SERVICE_QUALIFIER)
-  private EgressRemediationService egressRemediationService;
+  private final EgressRemediationService sumologicEgressRemediationService;
+  private final EgressRemediationService vwbEgressRemediationService;
+
+  public CloudTaskEgressController(
+      @Qualifier(EGRESS_SUMOLOGIC_SERVICE_QUALIFIER)
+          EgressRemediationService sumologicEgressRemediationService,
+      @Qualifier(EGRESS_VWB_SERVICE_QUALIFIER)
+          EgressRemediationService vwbEgressRemediationService) {
+    this.sumologicEgressRemediationService = sumologicEgressRemediationService;
+    this.vwbEgressRemediationService = vwbEgressRemediationService;
+  }
 
   @Override
   public ResponseEntity<Void> processEgressEvent(ProcessEgressEventRequest request) {
-    egressRemediationService.remediateEgressEvent(request.getEventId());
+    if (request.isIsVwbEgressEvent()) {
+      vwbEgressRemediationService.remediateEgressEvent(request.getEventId());
+    } else {
+      sumologicEgressRemediationService.remediateEgressEvent(request.getEventId());
+    }
     return ResponseEntity.ok().build();
   }
 }
diff --git a/api/src/main/java/org/pmiops/workbench/api/OfflineEgressController.java b/api/src/main/java/org/pmiops/workbench/api/OfflineEgressController.java
@@ -5,11 +5,11 @@
 import java.time.Duration;
 import java.util.List;
 import java.util.logging.Logger;
-import java.util.stream.Collectors;
 import org.pmiops.workbench.cloudtasks.TaskQueueService;
 import org.pmiops.workbench.db.dao.EgressEventDao;
 import org.pmiops.workbench.db.model.DbEgressEvent;
 import org.pmiops.workbench.db.model.DbEgressEvent.DbEgressEventStatus;
+import org.pmiops.workbench.exfiltration.ExfiltrationUtils;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.RestController;
@@ -27,20 +27,17 @@ public class OfflineEgressController implements OfflineEgressApiDelegate {
   public ResponseEntity<Void> checkPendingEgressEvents() {
     Timestamp latestModifiedTime = Timestamp.from(clock.instant().minus(PENDING_EVENT_LIMIT));
 
-    List<Long> oldPendingEventIds =
-        egressEventDao
-            .findAllByStatusAndLastModifiedTimeLessThan(
-                DbEgressEventStatus.PENDING, latestModifiedTime)
-            .stream()
-            .map(DbEgressEvent::getEgressEventId)
-            .collect(Collectors.toList());
-    for (long eventId : oldPendingEventIds) {
-      taskQueueService.pushEgressEventTask(eventId);
+    List<DbEgressEvent> oldPendingEvents =
+        egressEventDao.findAllByStatusAndLastModifiedTimeLessThan(
+            DbEgressEventStatus.PENDING, latestModifiedTime);
+    for (DbEgressEvent event : oldPendingEvents) {
+      taskQueueService.pushEgressEventTask(
+          event.getEgressEventId(), ExfiltrationUtils.isVwbEgressEvent(event));
     }
 
-    if (oldPendingEventIds.size() > 0) {
+    if (oldPendingEvents.size() > 0) {
       log.warning(
-          String.format("found and re-enqueued %d old PENDING events", oldPendingEventIds.size()));
+          String.format("found and re-enqueued %d old PENDING events", oldPendingEvents.size()));
     }
     return ResponseEntity.noContent().build();
   }

diff --git a/api/src/main/java/org/pmiops/workbench/api/VwbEgressAdminController.java b/api/src/main/java/org/pmiops/workbench/api/VwbEgressAdminController.java
@@ -1,13 +1,52 @@
 package org.pmiops.workbench.api;
 
+import jakarta.inject.Provider;
+import java.util.logging.Logger;
+import org.pmiops.workbench.config.WorkbenchConfig;
+import org.pmiops.workbench.db.model.DbUser;
+import org.pmiops.workbench.exceptions.ForbiddenException;
+import org.pmiops.workbench.exceptions.UnauthorizedException;
+import org.pmiops.workbench.exfiltration.EgressEventService;
 import org.pmiops.workbench.model.VwbEgressEventRequest;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.RestController;
 
 @RestController
 public class VwbEgressAdminController implements VwbEgressAdminApiDelegate {
+  private static final Logger log = Logger.getLogger(NotebooksController.class.getName());
+  private final EgressEventService egressEventService;
+
+  private final Provider<WorkbenchConfig> workbenchConfigProvider;
+
+  private final Provider<DbUser> userProvider;
+
+  @Autowired
+  public VwbEgressAdminController(
+      EgressEventService egressEventService,
+      Provider<WorkbenchConfig> workbenchConfigProvider,
+      Provider<DbUser> userProvider) {
+    this.egressEventService = egressEventService;
+    this.workbenchConfigProvider = workbenchConfigProvider;
+    this.userProvider = userProvider;
+  }
+
   @Override
   public ResponseEntity<Void> createVwbEgressEvent(VwbEgressEventRequest body) {
+    if (!workbenchConfigProvider.get().featureFlags.enableVWBEgressMonitor) {
+      throw new ForbiddenException("VWB Egress Monitor is disabled");
+    }
+    if (!workbenchConfigProvider
+        .get()
+        .vwb
+        .exfilManagerServiceAccount
+        .equals(userProvider.get().getUsername())) {
+      throw new UnauthorizedException(
+          String.format(
+              "User %s is not authorized to create VwbEgressEvent",
+              userProvider.get().getUsername()));
+    }
+    egressEventService.handleVwbEvent(body);
     return ResponseEntity.ok().build();
   }
 }
diff --git a/api/src/main/java/org/pmiops/workbench/cloudtasks/TaskQueueService.java b/api/src/main/java/org/pmiops/workbench/cloudtasks/TaskQueueService.java
@@ -175,11 +175,11 @@ private <T> List<List<T>> groupDeleteWorkspaceTasks(List<T> workspacesToDelete)
     return CloudTasksUtils.partitionList(workspacesToDelete, batchSize);
   }
 
-  public void pushEgressEventTask(Long eventId) {
+  public void pushEgressEventTask(Long eventId, boolean isVwbEgressEvent) {
     createAndPushTask(
         EGRESS_EVENT_QUEUE_NAME,
         EGRESS_EVENT_PATH,
-        new ProcessEgressEventRequest().eventId(eventId));
+        new ProcessEgressEventRequest().eventId(eventId).isVwbEgressEvent(isVwbEgressEvent));
   }
 
   public void pushCreateWorkspaceTask(long operationId, Workspace workspace) {

diff --git a/api/src/main/java/org/pmiops/workbench/db/dao/EgressEventDao.java b/api/src/main/java/org/pmiops/workbench/db/dao/EgressEventDao.java
@@ -1,6 +1,7 @@
 package org.pmiops.workbench.db.dao;
 
 import java.sql.Timestamp;
+import java.util.Collection;
 import java.util.List;
 import org.pmiops.workbench.db.model.DbEgressEvent;
 import org.pmiops.workbench.db.model.DbEgressEvent.DbEgressEventStatus;
@@ -23,7 +24,8 @@ public interface EgressEventDao
   Page<DbEgressEvent> findAllByUserAndWorkspaceOrderByCreationTimeDesc(
       DbUser user, DbWorkspace workspace, Pageable p);
 
-  List<DbEgressEvent> findAllByUserAndStatusNot(DbUser user, DbEgressEventStatus status);
+  List<DbEgressEvent> findAllByUserAndStatusNotIn(
+      DbUser user, Collection<DbEgressEventStatus> status);
 
   List<DbEgressEvent> findAllByUserAndWorkspaceAndCreationTimeBetweenAndCreationTimeNot(
       DbUser user,

diff --git a/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java b/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
@@ -42,6 +42,12 @@ public enum DbEgressEventStatus {
 
   private String bucketAuditEvent;
 
+  private String vwbWorkspaceId;
+  private String vwbVmName;
+  private String gcpProjectId;
+  private String vwbEgressEventId;
+  private int vwbIncidentCount;
+
   @Id
   @GeneratedValue(strategy = GenerationType.IDENTITY)
   @Column(name = "egress_event_id", nullable = false)
@@ -149,6 +155,56 @@ public DbEgressEvent setStatus(DbEgressEventStatus s) {
     return this;
   }
 
+  @Column(name = "vwb_workspace_id")
+  public String getVwbWorkspaceId() {
+    return vwbWorkspaceId;
+  }
+
+  public DbEgressEvent setVwbWorkspaceId(String vwbWorkspaceId) {
+    this.vwbWorkspaceId = vwbWorkspaceId;
+    return this;
+  }
+
+  @Column(name = "vwb_vm_name")
+  public String getVwbVmName() {
+    return vwbVmName;
+  }
+
+  public DbEgressEvent setVwbVmName(String vwbVmName) {
+    this.vwbVmName = vwbVmName;
+    return this;
+  }
+
+  @Column(name = "gcp_project_id")
+  public String getGcpProjectId() {
+    return gcpProjectId;
+  }
+
+  public DbEgressEvent setGcpProjectId(String gcpProjectId) {
+    this.gcpProjectId = gcpProjectId;
+    return this;
+  }
+
+  @Column(name = "vwb_egress_event_id")
+  public String getVwbEgressEventId() {
+    return vwbEgressEventId;
+  }
+
+  public DbEgressEvent setVwbEgressEventId(String vwbEgressEventId) {
+    this.vwbEgressEventId = vwbEgressEventId;
+    return this;
+  }
+
+  @Column(name = "vwb_incident_count")
+  public int getVwbIncidentCount() {
+    return vwbIncidentCount;
+  }
+
+  public DbEgressEvent setVwbIncidentCount(int vwbIncidentCount) {
+    this.vwbIncidentCount = vwbIncidentCount;
+    return this;
+  }
+
   public boolean equals(Object o) {
     if (this == o) {
       return true;
@@ -165,7 +221,11 @@ public boolean equals(Object o) {
         && Objects.equals(creationTime, that.creationTime)
         && Objects.equals(lastModifiedTime, that.lastModifiedTime)
         && status == that.status
-        && Objects.equals(sumologicEvent, that.sumologicEvent);
+        && Objects.equals(sumologicEvent, that.sumologicEvent)
+        && Objects.equals(vwbWorkspaceId, that.vwbWorkspaceId)
+        && Objects.equals(vwbVmName, that.vwbVmName)
+        && Objects.equals(gcpProjectId, that.gcpProjectId)
+        && Objects.equals(vwbIncidentCount, that.vwbIncidentCount);
   }
 
   @Override
@@ -179,6 +239,10 @@ public int hashCode() {
         creationTime,
         lastModifiedTime,
         status,
-        sumologicEvent);
+        sumologicEvent,
+        vwbWorkspaceId,
+        vwbVmName,
+        gcpProjectId,
+        vwbIncidentCount);
   }
 }
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventService.java b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventService.java
@@ -1,7 +1,10 @@
 package org.pmiops.workbench.exfiltration;
 
 import org.pmiops.workbench.model.SumologicEgressEvent;
+import org.pmiops.workbench.model.VwbEgressEventRequest;
 
 public interface EgressEventService {
   void handleEvent(SumologicEgressEvent egressEvent);
+
+  void handleVwbEvent(VwbEgressEventRequest egressEvent);
 }