Merge pull request #39 from fahadkhan-fk/feature-stream

Live stream command support

Author: wh1te909

src/components/agents/CommandStream.vue

@@ -0,0 +1,104 @@
+<template>
+  <q-card-section class="q-px-xs q-pb-md q-pt-xs" v-if="hasText">
+    <script-output-copy-clip label="Live Output" :data="outputText" />
+    <q-separator class="q-my-sm" />
+  </q-card-section>
+  <div class="command-stream" ref="streamContainer" v-if="hasText">
+    <div class="terminal">
+      <pre class="mt-0">{{ outputText }}</pre>
+    </div>
+  </div>
+</template>
+
+<script setup lang="ts">
+import {
+  ref,
+  onMounted,
+  onUnmounted,
+  watchEffect,
+  nextTick,
+  computed,
+} from "vue";
+import { useTemplateRef } from "vue";
+import { useAgentCmdWSConnection } from "@/websocket/agent";
+import ScriptOutputCopyClip from "@/components/scripts/ScriptOutputCopyClip.vue";
+import { uid } from "quasar";
+
+const props = defineProps({
+  agentId: { type: String, required: true },
+  cmd: { type: String, required: true },
+  shell: { type: String, required: true },
+  timeout: { type: Number, default: 10 },
+});
+
+// emits
+const emit = defineEmits(["updateOutput", "streamLoaded", "streamClosed"]);
+
+const cmdId = uid();
+const { send, data, reset, close, status } = useAgentCmdWSConnection(
+  props.agentId,
+  cmdId,
+);
+
+const outputText = ref("");
+const streamContainer = useTemplateRef<HTMLElement>("streamContainer");
+let firstChunk = false;
+
+const hasText = computed(() => outputText.value.trim() !== "");
+
+watchEffect(() => {
+  if (data.value.length) {
+    outputText.value = data.value.map((msg) => msg.output).join("\n");
+    emit("updateOutput", outputText.value);
+
+    if (!firstChunk && data.value.length > 0) {
+      firstChunk = true;
+      emit("streamLoaded");
+    }
+
+    nextTick(() => {
+      if (streamContainer.value) {
+        streamContainer.value.scrollTop = streamContainer.value.scrollHeight;
+      }
+    });
+  }
+});
+
+watchEffect(() => {
+  if (status.value === "CLOSED") emit("streamClosed");
+});
+
+onMounted(() => {
+  outputText.value = "";
+  reset();
+  send(
+    JSON.stringify({
+      shell: props.shell,
+      cmd: props.cmd,
+      timeout: props.timeout,
+      run_as_user: false,
+      custom_shell: "",
+      stream: true,
+      cmd_id: cmdId,
+    }),
+  );
+});
+
+onUnmounted(close);
+</script>
+
+<style scoped>
+.command-stream {
+  font-family: monospace;
+  background: #191818;
+  color: #fff;
+  padding: 0 10px;
+  height: 30vh;
+  overflow-y: auto;
+  border: 1px solid #ccc;
+  border-radius: 4px;
+}
+.terminal {
+  white-space: pre-wrap;
+}
+</style>

src/components/modals/agents/SendCommand.vue

@@ -7,7 +7,7 @@
   >
     <q-card
       class="q-dialog-plugin"
-      :style="{ 'min-width': !ret ? '40vw' : '70vw' }"
+      :style="{ 'min-width': ret || streamOutput ? '70vw' : '40vw' }"
     >
       <q-bar>
         Send command on {{ agent.hostname }}
@@ -82,7 +82,7 @@
             ]"
           />
         </q-card-section>
-        <q-card-section>
+        <q-card-section class="q-pb-xs">
           <q-input
             v-model="state.cmd"
             outlined
@@ -92,17 +92,20 @@
             :rules="[(val) => !!val || '*Required']"
           />
         </q-card-section>
-        <q-card-actions align="right">
-          <q-btn flat dense push label="Cancel" v-close-popup />
-          <q-btn
-            :loading="loading"
-            flat
-            dense
-            push
-            label="Send"
-            color="primary"
-            type="submit"
-          />
+        <q-card-actions align="between">
+          <q-toggle v-model="useStreaming" label="Stream Output" />
+          <div>
+            <q-btn flat dense push label="Cancel" v-close-popup />
+            <q-btn
+              :loading="loading"
+              flat
+              dense
+              push
+              label="Send"
+              color="primary"
+              type="submit"
+            />
+          </div>
         </q-card-actions>
         <q-card-section v-if="ret !== null"
           ><script-output-copy-clip label="Output" :data="ret" /> <q-separator
@@ -114,35 +117,47 @@
         >
           <pre>{{ ret }}</pre>
         </q-card-section>
+        <q-card-section v-if="showStream" class="q-py-xs">
+          <command-stream
+            :key="`${runId}`"
+            :agent-id="agent.agent_id"
+            :cmd="streamCmd"
+            :shell="state.shell"
+            :timeout="state.timeout"
+            @updateOutput="(val) => (streamOutput = val)"
+            @streamLoaded="loading = false"
+            @streamClosed="loading = false"
+          />
+        </q-card-section>
       </q-form>
     </q-card>
   </q-dialog>
 </template>
 
 <script>
 // composition imports
-import { ref } from "vue";
+import { ref, nextTick } from "vue";
 import { useDialogPluginComponent } from "quasar";
 import { sendAgentCommand } from "@/api/agents";
 import { cmdPlaceholder } from "@/composables/agents";
 import { runAsUserToolTip } from "@/constants/constants";
 
 import ScriptOutputCopyClip from "@/components/scripts/ScriptOutputCopyClip.vue";
+import CommandStream from "@/components/agents/CommandStream.vue";
 
 export default {
   name: "SendCommand",
   components: {
     ScriptOutputCopyClip,
+    CommandStream,
   },
   emits: [...useDialogPluginComponent.emits],
   props: {
     agent: !Object,
   },
   setup(props) {
-    // setup quasar dialog plugin
     const { dialogRef, onDialogHide } = useDialogPluginComponent();
 
-    // run command logic
     const state = ref({
       shell: props.agent.plat === "windows" ? "cmd" : "/bin/bash",
       cmd: null,
@@ -153,10 +168,30 @@ export default {
 
     const loading = ref(false);
     const ret = ref(null);
+    const useStreaming = ref(false);
+    const showStream = ref(false);
+    const streamOutput = ref("");
+    const streamCmd = ref("");
+    const streamShell = ref("");
+    const streamTimeout = ref(30);
+    const runId = ref(0);
 
     async function submit() {
       loading.value = true;
       ret.value = null;
+      streamOutput.value = "";
+      showStream.value = false;
+
+      if (useStreaming.value) {
+        streamCmd.value = state.value.cmd;
+        streamShell.value = state.value.shell;
+        streamTimeout.value = state.value.timeout;
+        await nextTick();
+        runId.value++;
+        showStream.value = true;
+        return;
+      }
+
       try {
         ret.value = await sendAgentCommand(props.agent.agent_id, state.value);
       } catch (e) {
@@ -167,20 +202,21 @@ export default {
 
     return {
       // reactive data
+      dialogRef,
+      onDialogHide,
       state,
       loading,
       ret,
-
-      // non reactivete data
-      runAsUserToolTip,
-
-      // methods
+      useStreaming,
+      streamOutput,
+      showStream,
+      streamCmd,
+      streamShell,
+      streamTimeout,
+      runId,
       submit,
+      runAsUserToolTip,
       cmdPlaceholder,
-
-      // quasar dialog
-      dialogRef,
-      onDialogHide,
     };
   },
 };

src/websocket/agent.ts

@@ -0,0 +1,74 @@
+import { ref } from "vue";
+import { useWebSocket } from "@vueuse/core";
+import { useAuthStore } from "@/stores/auth";
+import { getWSUrl } from "./websocket";
+import { Notify } from "quasar";
+interface CmdMessage {
+  cmd_id: string;
+  output?: string;
+  done?: boolean;
+  exit_code?: number;
+  [key: string]: unknown;
+}
+
+export function useAgentCmdWSConnection(agentId: string, cmdId: string) {
+  const auth = useAuthStore();
+  const lines = ref<CmdMessage[]>([]);
+  const url = getWSUrl(`agent/${agentId}/cmd`, auth.token);
+  const ws = useWebSocket(url, {
+    autoReconnect: false,
+    onMessage(_, ev) {
+      try {
+        const parsed = JSON.parse(ev.data);
+        if (parsed?.error) {
+          const msg = `${parsed.error ? parsed.error : "Unknown WebSocket error"}`;
+          const caption = parsed.status
+            ? `${parsed.status}: Forbidden`
+            : "Error";
+          Notify.create({
+            message: msg,
+            color: "negative",
+            position: "top",
+            caption,
+            timeout: 4000,
+          });
+          ws.close();
+          return;
+        }
+        if (parsed?.cmd_id !== cmdId) return;
+        if (
+          parsed?.output != null &&
+          !(
+            (lines.value.length === 0 || lines.value.length === 1) &&
+            parsed.output.trim() === ""
+          )
+        ) {
+          lines.value.push(parsed);
+        }
+      } catch {
+        lines.value.push({
+          cmd_id,
+          output: "[Error] Unable to parse server output",
+        });
+      }
+    },
+  });
+
+  function reset() {
+    lines.value = [];
+  }
+
+  function closeConnection() {
+    ws.close();
+    lines.value = [];
+  }
+
+  return {
+    status: ws.status,
+    data: lines,
+    send: ws.send,
+    open: ws.open,
+    reset,
+    close: closeConnection,
+  };
+}