From b9643e7d5c9d179eb7e44971b9649b686edbeac4 Mon Sep 17 00:00:00 2001
From: wancat <wancat@wancat.cc>
Date: Thu, 27 Jul 2023 17:56:32 +0800
Subject: [PATCH] feat: capture output and exit code of commands

+ ignore interactive commands like vim, top or nano
---
 server/quests/helloworld/Dockerfile           |  3 ++-
 server/quests/helloworld/home/.bashrc         | 20 ++++++++++++++++---
 .../src/api/controllers/commandController.js  |  3 ++-
 server/src/models/command.js                  |  4 ++++
 4 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/server/quests/helloworld/Dockerfile b/server/quests/helloworld/Dockerfile
index 6ef8747..01f3814 100644
--- a/server/quests/helloworld/Dockerfile
+++ b/server/quests/helloworld/Dockerfile
@@ -2,7 +2,8 @@ FROM ubuntu:latest
 
 RUN apt-get update && apt-get install -y \
   curl \
-  vim
+  vim \
+  jq
 
 RUN adduser --disabled-password --gecos '' rudeus
 
diff --git a/server/quests/helloworld/home/.bashrc b/server/quests/helloworld/home/.bashrc
index 7c4d2bf..3a0355f 100644
--- a/server/quests/helloworld/home/.bashrc
+++ b/server/quests/helloworld/home/.bashrc
@@ -4,12 +4,26 @@ preexec() {
     [ "$BASH_COMMAND" = "$PROMPT_COMMAND" ] && return # do nothing if completing
     local command="$BASH_COMMAND"
     local pwd="$PWD"
-    
+
+    if [[ $command =~ ^(vim|top|less|more|nano) ]]; then
+        local payload=$(printf '{"token": "%s", "pwd": "%s", "command": "%s"}' "$TOKEN" "$pwd" "$command")
+        curl -X POST $API_ENDPOINT/api/v1/commands -H 'accept: */*' -H 'Content-Type: application/json' -d "$payload"
+        return
+    fi
+
+    # Execute the command, capture output and exit code
+    local output=$( { eval $command 2>&1; echo $? > /tmp/exit_code; }) 
+    local exit_code=$(cat /tmp/exit_code)
+
+    local output_esc=$(echo $output | jq -Rs .)
+
     # Store command in variable
-    local cmd=$(printf '{"token": "%s", "pwd": "%s", "command": "%s"}' "$TOKEN" "$pwd" "$command")
+    local payload=$(printf '{"token": "%s", "pwd": "%s", "command": "%s", "exit_code": "%s", "output": %s}' "$TOKEN" "$pwd" "$command" "$exit_code" "$output_esc")
     
     # Send command to the API
-    curl -X POST $API_ENDPOINT/api/v1/commands -H 'accept: */*' -H 'Content-Type: application/json' -d "$cmd"
+    curl -X POST $API_ENDPOINT/api/v1/commands -H 'accept: */*' -H 'Content-Type: application/json' -d "$payload" || echo $payload
+
+    BASH_COMMAND=""
 }
 
 # Install the trap
diff --git a/server/src/api/controllers/commandController.js b/server/src/api/controllers/commandController.js
index 92c8374..a886d23 100644
--- a/server/src/api/controllers/commandController.js
+++ b/server/src/api/controllers/commandController.js
@@ -5,7 +5,7 @@ import { verifySessionJWT } from '../../utils/auth.js'
 
 export async function newCommand(req, res) {
   console.log('new command:', req.body)
-  const { token, command, pwd, output } = req.body
+  const { token, command, pwd, exit_code, output } = req.body
 
   let sessionId
   try {
@@ -27,6 +27,7 @@ export async function newCommand(req, res) {
     session,
     command,
     pwd,
+    exit_code,
     output,
   })
 
diff --git a/server/src/models/command.js b/server/src/models/command.js
index 4ea6697..ce64e7c 100644
--- a/server/src/models/command.js
+++ b/server/src/models/command.js
@@ -17,6 +17,10 @@ const Command = model(
         type: String,
         required: true,
       },
+      exit_code: {
+        type: Number,
+        required: false,
+      },
       output: {
         type: String,
         required: false,