[feature] Structured Reporting - Process

implants/lib/c2/src/c2.rs

@@ -124,6 +124,17 @@ pub struct TaskOutput {
     #[prost(message, optional, tag = "5")]
     pub exec_finished_at: ::core::option::Option<::prost_types::Timestamp>,
 }
+/// Process running on the system.
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct Process {
+    #[prost(uint64, tag = "1")]
+    pub pid: u64,
+    #[prost(string, tag = "2")]
+    pub name: ::prost::alloc::string::String,
+    #[prost(string, tag = "3")]
+    pub principal: ::prost::alloc::string::String,
+}
 ///
 /// RPC Messages
 #[allow(clippy::derive_partial_eq_without_eq)]
@@ -159,6 +170,17 @@ pub struct DownloadFileResponse {
     #[prost(bytes = "vec", tag = "1")]
     pub chunk: ::prost::alloc::vec::Vec<u8>,
 }
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct ReportProcessListRequest {
+    #[prost(message, repeated, tag = "1")]
+    pub list: ::prost::alloc::vec::Vec<Process>,
+    #[prost(int64, tag = "2")]
+    pub task_id: i64,
+}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct ReportProcessListResponse {}
 /// Generated client implementations.
 pub mod c2_client {
     #![allow(unused_variables, dead_code, missing_docs, clippy::let_unit_value)]
@@ -289,6 +311,31 @@ pub mod c2_client {
             self.inner.unary(req, path, codec).await
         }
         ///
+        /// Report the active list of running processes. This list will replace any previously reported
+        /// lists for the same host.
+        pub async fn report_process_list(
+            &mut self,
+            request: impl tonic::IntoRequest<super::ReportProcessListRequest>,
+        ) -> std::result::Result<
+            tonic::Response<super::ReportProcessListResponse>,
+            tonic::Status,
+        > {
+            self.inner
+                .ready()
+                .await
+                .map_err(|e| {
+                    tonic::Status::new(
+                        tonic::Code::Unknown,
+                        format!("Service was not ready: {}", e.into()),
+                    )
+                })?;
+            let codec = tonic::codec::ProstCodec::default();
+            let path = http::uri::PathAndQuery::from_static("/c2.C2/ReportProcessList");
+            let mut req = request.into_request();
+            req.extensions_mut().insert(GrpcMethod::new("c2.C2", "ReportProcessList"));
+            self.inner.unary(req, path, codec).await
+        }
+        ///
         /// Download a file from the server, returning one or more chunks of data.
         /// The maximum size of these chunks is determined by the server.
         /// The server should reply with two headers:

tavern/internal/c2/api_report_process_list.go

@@ -0,0 +1,84 @@
+package c2
+
+import (
+	"context"
+	"fmt"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"realm.pub/tavern/internal/c2/c2pb"
+	"realm.pub/tavern/internal/ent"
+)
+
+func (srv *Server) ReportProcessList(ctx context.Context, req *c2pb.ReportProcessListRequest) (*c2pb.ReportProcessListResponse, error) {
+	// Validate Arguments
+	if req.TaskId == 0 {
+		return nil, status.Errorf(codes.InvalidArgument, "must provide task id")
+	}
+	if len(req.List) < 1 {
+		return nil, status.Errorf(codes.InvalidArgument, "must provide process list")
+	}
+
+	// Load Task
+	task, err := srv.graph.Task.Get(ctx, int(req.TaskId))
+	if ent.IsNotFound(err) {
+		return nil, status.Errorf(codes.NotFound, "no task found")
+	}
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "failed to load task")
+	}
+
+	// Load Host
+	host, err := task.QueryBeacon().QueryHost().Only(ctx)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "failed to load host")
+	}
+
+	// 	Prepare Transaction
+	tx, err := srv.graph.Tx(ctx)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "failed to initialize transaction: %v", err)
+	}
+	txGraph := tx.Client()
+
+	// Rollback transaction if we panic
+	defer func() {
+		if v := recover(); v != nil {
+			tx.Rollback()
+			panic(v)
+		}
+	}()
+
+	// Create Processes
+	builders := make([]*ent.ProcessCreate, 0, len(req.List))
+	for _, proc := range req.List {
+		builders = append(builders,
+			txGraph.Process.Create().
+				SetHostID(host.ID).
+				SetTaskID(task.ID).
+				SetPid(proc.Pid).
+				SetName(proc.Name).
+				SetPrincipal(proc.Principal),
+		)
+	}
+	processList, err := txGraph.Process.CreateBulk(builders...).Save(ctx)
+	if err != nil {
+		return nil, rollback(tx, fmt.Errorf("failed to create process list: %w", err))
+	}
+
+	// Set new process list for host
+	_, err = txGraph.Host.UpdateOne(host).
+		ClearProcesses().
+		AddProcesses(processList...).
+		Save(ctx)
+	if err != nil {
+		return nil, rollback(tx, fmt.Errorf("failed to set new host process list: %w", err))
+	}
+
+	// Commit the transaction
+	if err := tx.Commit(); err != nil {
+		return nil, rollback(tx, fmt.Errorf("failed to commit transaction: %w", err))
+	}
+
+	return &c2pb.ReportProcessListResponse{}, nil
+}

tavern/internal/c2/api_report_process_list_test.go

@@ -0,0 +1,135 @@
+package c2_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/testing/protocmp"
+	"realm.pub/tavern/internal/c2/c2pb"
+	"realm.pub/tavern/internal/c2/c2test"
+	"realm.pub/tavern/internal/ent"
+)
+
+func TestReportProcessList(t *testing.T) {
+	// Setup Dependencies
+	ctx := context.Background()
+	client, graph, close := c2test.New(t)
+	defer close()
+
+	// Test Data
+	existingBeacon := c2test.NewRandomBeacon(ctx, graph)
+	existingTask := c2test.NewRandomAssignedTask(ctx, graph, existingBeacon.Identifier)
+	existingHost := existingBeacon.QueryHost().OnlyX(ctx)
+
+	// Test Cases
+	tests := []struct {
+		name string
+		host *ent.Host
+		task *ent.Task
+		req  *c2pb.ReportProcessListRequest
+
+		wantResp     *c2pb.ReportProcessListResponse
+		wantCode     codes.Code
+		wantHostPIDs []uint64
+		wantTaskPIDs []uint64
+	}{
+		{
+			name: "New_List",
+			host: existingHost,
+			task: existingTask,
+			req: &c2pb.ReportProcessListRequest{
+				TaskId: int64(existingTask.ID),
+				List: []*c2pb.Process{
+					{Pid: 1, Name: "systemd", Principal: "root"},
+					{Pid: 2321, Name: "/bin/sh", Principal: "root"},
+					{Pid: 4505, Name: "/usr/bin/sshd", Principal: "root"},
+				},
+			},
+			wantResp:     &c2pb.ReportProcessListResponse{},
+			wantCode:     codes.OK,
+			wantHostPIDs: []uint64{1, 2321, 4505},
+			wantTaskPIDs: []uint64{1, 2321, 4505},
+		},
+		{
+			name: "Updated_List",
+			host: existingHost,
+			task: existingTask,
+			req: &c2pb.ReportProcessListRequest{
+				TaskId: int64(existingTask.ID),
+				List: []*c2pb.Process{
+					{Pid: 1, Name: "systemd", Principal: "root"},
+					{Pid: 4505, Name: "/usr/bin/sshd", Principal: "root"},
+					{Pid: 4809, Name: "/usr/bin/nginx", Principal: "root"},
+				},
+			},
+			wantResp:     &c2pb.ReportProcessListResponse{},
+			wantCode:     codes.OK,
+			wantHostPIDs: []uint64{1, 4505, 4809},
+			wantTaskPIDs: []uint64{1, 2321, 4505, 1, 4505, 4809},
+		},
+		{
+			name: "No_TaskID",
+			host: existingHost,
+			task: existingTask,
+			req: &c2pb.ReportProcessListRequest{
+				List: []*c2pb.Process{
+					{Pid: 1, Name: "systemd", Principal: "root"},
+				},
+			},
+			wantResp: nil,
+			wantCode: codes.InvalidArgument,
+		},
+		{
+			name: "No_Processes",
+			host: existingHost,
+			task: existingTask,
+			req: &c2pb.ReportProcessListRequest{
+				TaskId: int64(existingTask.ID),
+				List:   []*c2pb.Process{},
+			},
+			wantResp: nil,
+			wantCode: codes.InvalidArgument,
+		},
+	}
+
+	// Run Tests
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			// gRPC
+			resp, err := client.ReportProcessList(ctx, tc.req)
+
+			// Assert Response Code
+			require.Equal(t, tc.wantCode.String(), status.Code(err).String(), err)
+			if status.Code(err) != codes.OK {
+				// Do not continue if we expected error code
+				return
+			}
+
+			// Assert Response
+			if diff := cmp.Diff(tc.wantResp, resp, protocmp.Transform()); diff != "" {
+				t.Errorf("invalid response (-want +got): %v", diff)
+			}
+
+			// Assert Task Processes
+			var taskPIDs []uint64
+			taskProcessList := tc.task.QueryReportedProcesses().AllX(ctx)
+			for _, proc := range taskProcessList {
+				taskPIDs = append(taskPIDs, proc.Pid)
+			}
+			assert.ElementsMatch(t, tc.wantTaskPIDs, taskPIDs)
+
+			// Assert Host Processes
+			var hostPIDs []uint64
+			hostProcessList := tc.host.QueryProcesses().AllX(ctx)
+			for _, proc := range hostProcessList {
+				hostPIDs = append(hostPIDs, proc.Pid)
+			}
+			assert.ElementsMatch(t, tc.wantHostPIDs, hostPIDs)
+		})
+	}
+}

tavern/internal/c2/c2.proto

@@ -69,6 +69,13 @@ message TaskOutput {
     google.protobuf.Timestamp exec_finished_at = 5;
 }
 
+// Process running on the system.
+message Process {
+    uint64 pid = 1;
+    string name = 2;
+    string principal = 3;
+}
+
 /*
  * RPC Messages
  */
@@ -94,6 +101,13 @@ message DownloadFileResponse {
     bytes chunk = 1;
 }
 
+message ReportProcessListRequest {
+    repeated Process list = 1;
+    int64 task_id = 2;
+}
+
+message ReportProcessListResponse {}
+
 /*
  * Service
  */
@@ -102,6 +116,12 @@ service C2 {
     rpc ClaimTasks(ClaimTasksRequest) returns (ClaimTasksResponse) {}
     rpc ReportTaskOutput(ReportTaskOutputRequest) returns (ReportTaskOutputResponse) {}
 
+    /*
+     * Report the active list of running processes. This list will replace any previously reported
+     * lists for the same host.
+     */
+    rpc ReportProcessList(ReportProcessListRequest) returns (ReportProcessListResponse);
+
     /*
      * Download a file from the server, returning one or more chunks of data.
      * The maximum size of these chunks is determined by the server.