Add process collection for command

Author: Tzvonimir

collector/collector.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"strconv"
 	"strings"
 	"sync"
 
@@ -287,7 +288,7 @@ func (c *Collector) handleSocketCollection(con net.Conn) error {
 
 	c.logger.Debug().Msgf("Received: %s", string(buf[:n]))
 
-	if len(parts) != 7 {
+	if len(parts) != 8 {
 		c.logger.Error().Msg("Invalid command format")
 		return fmt.Errorf("invalid command format")
 	}
@@ -321,13 +322,16 @@ func (c *Collector) handleStartCommand(parts []string) error {
 		c.logger.Error().Err(err).Msg("Failed to get repository name")
 	}
 
+	pid, _ := strconv.ParseInt(parts[5], 10, 64)
+
 	command := Command{
 		Category:   ParseCommand(parts[1]),
 		Command:    parts[1],
 		Directory:  parts[2],
 		User:       parts[3],
 		StartTime:  time.Now().UnixMilli(), // TODO: there are some issues with sending time through shell because of ms support on MAC, explore more
 		Repository: repo,
+		PID:        pid,
 	}
 
 	c.collectionConfig.ongoingCommands[parts[4]] = command
@@ -349,8 +353,8 @@ func (c *Collector) handleEndCommand(parts []string) error {
 	if command, exists := c.collectionConfig.ongoingCommands[parts[4]]; exists {
 		command.EndTime = time.Now().UnixMilli()
 		command.ExecutionTime = command.EndTime - command.StartTime
-		command.Result = parts[5]
-		command.Status = parts[6]
+		command.Result = parts[6]
+		command.Status = parts[7]
 
 		c.logger.Debug().Msgf("Command: %+v", command)
 		if err := InsertCommand(command); err != nil {

collector/command.go

@@ -23,6 +23,7 @@ type Command struct {
 	Status        string `json:"status" db:"status"`
 	Result        string `json:"result" db:"result"`
 	Repository    string `json:"repository" db:"repository"`
+	PID           int64  `json:"pid" db:"pid"`
 }
 
 // GetCommandById fetches a command by its ID
@@ -91,8 +92,8 @@ func DeleteCommandsByDays(days int) error {
 
 // InsertCommand inserts a command into the database
 func InsertCommand(command Command) error {
-	query := `INSERT INTO commands (category, command, user, directory, execution_time, start_time, end_time, status, result, repository)
-	VALUES (:category, :command, :user, :directory, :execution_time, :start_time, :end_time, :status, :result, :repository)`
+	query := `INSERT INTO commands (category, command, user, directory, execution_time, start_time, end_time, status, result, repository, pid)
+	VALUES (:category, :command, :user, :directory, :execution_time, :start_time, :end_time, :status, :result, :repository, :pid)`
 
 	_, err := database.DB.NamedExec(query, command)
 
@@ -139,5 +140,6 @@ func MapCommandToProto(command Command) *gen.Command {
 		Status:        command.Status,
 		Result:        command.Result,
 		Repository:    command.Repository,
+		Pid:           command.PID,
 	}
 }

database/migrations.go

@@ -74,7 +74,8 @@ func createCommandsTable() {
 			end_time INTEGER,
 			status TEXT,
 			result TEXT,
-			repository TEXT
+			repository TEXT,
+			pid INTEGER
 		);`
 
 		_, err := DB.Exec(createCommandsTableSQL)

gen/api/v1/collector.pb.go

@@ -29,6 +29,7 @@ message Command {
   string result = 9; // Result of executed command => success/failure
   string status = 10; // Status of executed command 
   string repository = 11; // Repository is repository where commands are executed
+  int64 pid = 12; // PID of the command
 }
 
 // Define a message representing a process, including its metadata and resource usage.

proto/api/v1/collector.proto

@@ -2,13 +2,18 @@ generate_uuid() {
     echo "$(date +%s)-$$-$RANDOM"
 }
 
+generate_ppid() {
+  echo "$$"
+}
+
 preexec_invoke_exec() {
     # Avoid running preexec_invoke_exec for PROMPT_COMMAND
     if [[ "$BASH_COMMAND" != "$PROMPT_COMMAND" ]]; then
         export UUID=$(generate_uuid)
+        export PID=$(generate_ppid)
         export LAST_COMMAND="$BASH_COMMAND"
         # Send a start execution message
-        {{.CommandScriptPath}} "start" "$LAST_COMMAND" "$PWD" "$USER" "$UUID"
+        {{.CommandScriptPath}} "start" "$LAST_COMMAND" "$PWD" "$USER" "$UUID" "$PID"
     fi
 }
 trap 'preexec_invoke_exec' DEBUG
@@ -22,7 +27,7 @@ precmd_invoke_cmd() {
     fi
 
     # Send an end execution message with the result and exit status
-    {{.CommandScriptPath}} "end" "$LAST_COMMAND" "$PWD" "$USER" "$UUID" "$result" "$exit_status"
+    {{.CommandScriptPath}} "end" "$LAST_COMMAND" "$PWD" "$USER" "$UUID" "$PID" "$result" "$exit_status"
 }
 
 # Update PROMPT_COMMAND to invoke precmd_invoke_cmd

shell/scripts/bash.sh

@@ -7,6 +7,7 @@
 # $3 - Working directory
 # $4 - User who executed the command
 # $5 - Unique identifier
+# #6 - PID of the command
 # $6 - Command result (success/failure)
 # $7 - Exit status
 
@@ -46,7 +47,7 @@ send_via_perl() {
 }
 
 # Construct the log message including result and exit status
-LOG_MESSAGE="$1|$2|$3|$4|$5|$6|$7"
+LOG_MESSAGE="$1|$2|$3|$4|$5|$6|$7|$8"
 
 # Send the log message to the Go application via UNIX socket
 if command_exists nc && nc_supports_U; then

shell/scripts/collector.sh

@@ -2,11 +2,16 @@ function generate_uuid
     echo (date +%s)"-"(echo %self)"-"(random)
 end
 
+function generate_ppid
+    echo (echo %self)
+end
+
 function fish_preexec --on-event fish_preexec
     set -gx LAST_COMMAND $argv[1]
     set -gx UUID (generate_uuid)
+    set -gx PID (generate_ppid)
     # Send a start execution message
-    {{.CommandScriptPath}} "start" "$LAST_COMMAND" "$PWD" "$USER" "$UUID"
+    {{.CommandScriptPath}} "start" "$LAST_COMMAND" "$PWD" "$USER" "$UUID" "$PID"
 end
 
 function fish_postexec --on-event fish_postexec
@@ -18,5 +23,5 @@ function fish_postexec --on-event fish_postexec
     end
 
     # Send an end execution message with result and exit status
-    {{.CommandScriptPath}} "end" "$LAST_COMMAND" "$PWD" "$USER" "$UUID" $result $exit_status
+    {{.CommandScriptPath}} "end" "$LAST_COMMAND" "$PWD" "$USER" "$UUID" "$PID" "$result" "$exit_status"
 end

shell/scripts/fish.sh

@@ -2,11 +2,16 @@ generate_uuid() {
   echo "$(date +%s)-$$-$RANDOM"
 }
 
+generate_ppid() {
+  echo "$$"
+}
+
 preexec() {
   export LAST_COMMAND=$1
   UUID=$(generate_uuid)
+  PID=$(generate_ppid)
   # Send a start execution message
-  {{.CommandScriptPath}} "start" "$LAST_COMMAND" "$PWD" "$USER" "$UUID"
+  {{.CommandScriptPath}} "start" "$LAST_COMMAND" "$PWD" "$USER" "$UUID" "$PID"
 }
 
 precmd() {
@@ -18,5 +23,5 @@ precmd() {
   fi
 
   # Send an end execution message with result and exit status
-  {{.CommandScriptPath}} "end" "$LAST_COMMAND" "$PWD" "$USER" "$UUID" "$result" "$exit_status"
+  {{.CommandScriptPath}} "end" "$LAST_COMMAND" "$PWD" "$USER" "$UUID" "$PID" "$result" "$exit_status"
 }