updated task response iteration and ui task cli updates

MythicReactUI/src/components/TopAppBarVertical.js

@@ -403,7 +403,7 @@ const SearchTokens = () => {
 const SearchProxies = () => {
     const theme = useTheme();
   return (
-      <StyledListItem className={classes.listSubHeader} component={Link} to='/new/search?tab=socks'>
+      <StyledListItem className={classes.listSubHeader} component={Link} to='/new/search?tab=proxies'>
         <StyledListItemIcon>
             <MythicStyledTooltip title={"Search Proxies"} tooltipStyle={{display: "inline-flex"}}>
                 <FontAwesomeIcon style={{color: theme.navBarTextIconColor}} size={"lg"} icon={faSocks} />

MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingInput.js

@@ -1,4 +1,4 @@
-import { IconButton, Typography } from '@mui/material';
+import { Chip, IconButton, Typography } from '@mui/material';
 import SendIcon from '@mui/icons-material/Send';
 import React from 'react';
 import {TextField} from '@mui/material';
@@ -32,6 +32,7 @@ subscription GetLoadedCommandsSubscription($callback_id: Int!){
             cmd
             id
             attributes
+            description
             payloadtype {
                 name
                 id
@@ -40,6 +41,7 @@ subscription GetLoadedCommandsSubscription($callback_id: Int!){
                 id
                 parameter_type: type 
                 choices
+                description
                 dynamic_query_function
                 required
                 name
@@ -210,6 +212,183 @@ const TaskingContextChip = ({title, label, value, color, callbackColor, emphasiz
         </span>
     )
 }
+const getTaskingParameterLabel = (parameter) => {
+    if(parameter.cli_name){
+        return `-${parameter.cli_name}`;
+    }
+    return parameter.name || parameter.display_name;
+}
+const getTaskingParameterTooltip = (parameter) => {
+    const label = getTaskingParameterLabel(parameter);
+    const displayName = parameter.display_name && parameter.display_name !== label ? `\nDisplay: ${parameter.display_name}` : "";
+    const internalName = parameter.name && parameter.name !== label ? `\nName: ${parameter.name}` : "";
+    const requiredText = parameter.required ? "\nRequired" : "\nOptional";
+    const description = parameter.description ? `\n${parameter.description}` : "";
+    return `${label} · ${parameter.parameter_type}${requiredText}${displayName}${internalName}${description}`;
+}
+const TaskingParameterPreviewChip = ({parameter, required=false, active=false}) => (
+    <MythicStyledTooltip title={getTaskingParameterTooltip(parameter)}>
+        <Chip
+            className={`mythic-tasking-parameter-preview-chip${required ? " mythic-tasking-parameter-preview-chip-required" : ""}${active ? " mythic-tasking-parameter-preview-chip-active" : ""}`}
+            label={
+                <span className="mythic-tasking-parameter-preview-chip-label">
+                    {active &&
+                        <span className="mythic-tasking-parameter-preview-chip-status">Current</span>
+                    }
+                    <span className="mythic-tasking-parameter-preview-chip-name">{getTaskingParameterLabel(parameter)}</span>
+                    <span className="mythic-tasking-parameter-preview-chip-type">{parameter.parameter_type}</span>
+                </span>
+            }
+            size="small"
+        />
+    </MythicStyledTooltip>
+)
+const parseCommandLineForParameterPreview = (commandLine, command) => {
+    if(commandLine.length > 0 && commandLine[0] === "{"){
+        try{
+            const parsedJson = JSON.parse(commandLine);
+            if(['string', 'number', 'boolean', null].includes(typeof parsedJson)){
+                return {"_": []};
+            }
+            return {...parsedJson, "_": []};
+        }catch(error){
+            return {"_": []};
+        }
+    }
+    try{
+        const argv = [];
+        let sQuoted = false;
+        let dQuoted = false;
+        let backslash = false;
+        let buffer = '';
+        commandLine.split('').forEach((value) => {
+            if((sQuoted || dQuoted) && value === "\\"){
+                if(!backslash){
+                    backslash = true;
+                    return;
+                }
+                backslash = false;
+                buffer += "\\";
+                return;
+            }
+            if(!sQuoted && !dQuoted){
+                if(value === `'`){
+                    if(backslash){
+                        backslash = false;
+                        buffer += "'";
+                        return;
+                    }
+                    sQuoted = true;
+                    buffer += value;
+                    return;
+                }
+                if(value === '"'){
+                    if(backslash){
+                        backslash = false;
+                        buffer += '"';
+                        return;
+                    }
+                    dQuoted = true;
+                    buffer += value;
+                    return;
+                }
+                if(value === " "){
+                    if(backslash){
+                        backslash = false;
+                        buffer += "\\";
+                    }
+                    if(buffer.length > 0){
+                        argv.push(buffer[buffer.length-1] === buffer[0] && [`'`, `"`].includes(buffer[0]) ? buffer.slice(1, -1) : buffer);
+                        buffer = '';
+                    }
+                    return;
+                }
+            }else if((sQuoted && value === `'`) || (dQuoted && value === '"')){
+                if(backslash){
+                    backslash = false;
+                    buffer += value;
+                }else{
+                    if(sQuoted){
+                        sQuoted = false;
+                    }
+                    if(dQuoted){
+                        dQuoted = false;
+                    }
+                    buffer += value;
+                }
+                return;
+            }
+            if(backslash){
+                buffer += `\\${value}`;
+                backslash = false;
+            }else{
+                buffer += value;
+            }
+        });
+        if(backslash){
+            buffer += "\\";
+        }
+        if(buffer.length > 0){
+            argv.push(buffer[buffer.length-1] === buffer[0] && [`'`, `"`].includes(buffer[0]) ? buffer.slice(1, -1) : buffer);
+        }
+        if(dQuoted || sQuoted){
+            return {"_": []};
+        }
+        const validCliNames = command.commandparameters.reduce((previous, parameter) => {
+            if(parameter.cli_name){
+                return {...previous, [`-${parameter.cli_name}`]: parameter.cli_name};
+            }
+            return previous;
+        }, {});
+        return argv.reduce((previous, value) => {
+            if(validCliNames[value]){
+                return {...previous, [validCliNames[value]]: true};
+            }
+            return previous;
+        }, {"_": []});
+    }catch(error){
+        return {"_": []};
+    }
+}
+const determineCommandGroupNamesForParameterPreview = (command, parsed) => {
+    if(command.commandparameters.length === 0 || !parsed){
+        return [];
+    }
+    let commandGroupOptions = command.commandparameters.reduce((previous, parameter) => {
+        const groupName = parameter.parameter_group_name || "Default";
+        if(previous.includes(groupName)){
+            return previous;
+        }
+        return [...previous, groupName];
+    }, []);
+    for(const key of Object.keys(parsed)){
+        if(key === "_"){
+            continue;
+        }
+        let parameterGroups = [];
+        let foundParameterGroup = false;
+        for(let i = 0; i < command.commandparameters.length; i++){
+            if(command.commandparameters[i].cli_name === key || command.commandparameters[i].display_name === key || command.commandparameters[i].name === key){
+                foundParameterGroup = true;
+                parameterGroups.push(command.commandparameters[i].parameter_group_name || "Default");
+            }
+        }
+        const intersection = commandGroupOptions.reduce((previous, groupName) => {
+            if(parameterGroups.includes(groupName)){
+                return [...previous, groupName];
+            }
+            return previous;
+        }, []);
+        if(intersection.length === 0){
+            if(foundParameterGroup){
+                return undefined;
+            }
+        }else{
+            commandGroupOptions = [...intersection];
+        }
+    }
+    return commandGroupOptions;
+}
 
 export function CallbacksTabsTaskingInputPreMemo(props){
     const toastId = "tasking-toast-message";
@@ -1648,6 +1827,105 @@ export function CallbacksTabsTaskingInputPreMemo(props){
             inputRef.current.focus();
         }
     }, [props.focus])
+    const getCommandParameterPreview = () => {
+        if(reverseSearching){
+            return {state: "empty", message: "Reverse search active"};
+        }
+        const trimmedMessage = message.trim();
+        if(trimmedMessage === ""){
+            return {state: "empty", message: "Type a loaded command to preview CLI parameters"};
+        }
+        const commandName = trimmedMessage.split(/\s+/)[0];
+        const matchingCommands = loadedOptions.current.filter((command) => {
+            const supportedOS = command.attributes?.supported_os || [];
+            return command.cmd === commandName && (supportedOS.length === 0 || supportedOS.includes(props.callback_os));
+        });
+        if(matchingCommands.length === 0){
+            return {state: "empty", message: "No loaded command selected"};
+        }
+        let command = undefined;
+        if(commandPayloadType !== ""){
+            command = matchingCommands.find((option) => option.payloadtype?.name === commandPayloadType);
+        }
+        if(!command){
+            command = matchingCommands.find((option) => option.payloadtype?.name === props.payloadtype_name);
+        }
+        if(!command){
+            command = matchingCommands[0];
+        }
+        const parameters = command.commandparameters || [];
+        if(parameters.length === 0){
+            return {state: "empty", command, message: "No CLI parameters for this command"};
+        }
+        const commandParametersText = trimmedMessage.split(/\s+/).slice(1).join(" ");
+        const parsed = parseCommandLineForParameterPreview(commandParametersText, command);
+        const commandGroupNames = determineCommandGroupNamesForParameterPreview(command, parsed);
+        if(commandGroupNames === undefined){
+            return {state: "empty", command, message: "Parameters conflict across groups"};
+        }
+        const previewGroupNames = commandGroupNames.length > 0 ? commandGroupNames : parameters.reduce((previous, parameter) => {
+            const groupName = parameter.parameter_group_name || "Default";
+            if(previous.includes(groupName)){
+                return previous;
+            }
+            return [...previous, groupName];
+        }, []);
+        let activeParameter = undefined;
+        try{
+            const [lastSuppliedParameter] = getLastSuppliedArgument(command, trimmedMessage, parsed);
+            if(lastSuppliedParameter?.cli_name && IsCLIPossibleParameterType(lastSuppliedParameter.parameter_type)){
+                const activeParameterOptions = parameters.filter((parameter) =>
+                    parameter.cli_name === lastSuppliedParameter.cli_name &&
+                    previewGroupNames.includes(parameter.parameter_group_name || "Default") &&
+                    IsCLIPossibleParameterType(parameter.parameter_type)
+                );
+                activeParameter = activeParameterOptions.find((parameter) => parameter.required) || activeParameterOptions[0] || lastSuppliedParameter;
+            }
+        }catch(error){
+            activeParameter = undefined;
+        }
+        const availableParameters = parameters.reduce((previous, parameter) => {
+            const groupName = parameter.parameter_group_name || "Default";
+            if(!previewGroupNames.includes(groupName)){
+                return previous;
+            }
+            if(!IsCLIPossibleParameterType(parameter.parameter_type)){
+                return previous;
+            }
+            if(!parameter.cli_name){
+                return previous;
+            }
+            if(activeParameter?.cli_name === parameter.cli_name){
+                return previous;
+            }
+            if(parsed[parameter.cli_name] !== undefined && !IsRepeatableCLIParameterType(parameter.parameter_type)){
+                return previous;
+            }
+            const previousIndex = previous.findIndex((existingParameter) => existingParameter.cli_name === parameter.cli_name);
+            if(previousIndex >= 0){
+                if(parameter.required && !previous[previousIndex].required){
+                    const updatedParameters = [...previous];
+                    updatedParameters[previousIndex] = parameter;
+                    return updatedParameters;
+                }
+                return previous;
+            }
+            return [...previous, parameter];
+        }, []);
+        if(availableParameters.length === 0 && !activeParameter){
+            return {state: "empty", command, message: "All CLI parameters for the current group are supplied"};
+        }
+        const requiredParameters = availableParameters.filter((parameter) => parameter.required);
+        const optionalParameters = availableParameters.filter((parameter) => !parameter.required);
+        return {
+            activeParameter,
+            command,
+            optionalParameters,
+            requiredParameters,
+            state: "parameters",
+        };
+    }
+    const commandParameterPreview = getCommandParameterPreview();
     const showTaskingContext = !hideTaskingContext.current;
     const taskingContextChips = [
         {
@@ -1828,6 +2106,40 @@ export function CallbacksTabsTaskingInputPreMemo(props){
                     }}
                 />
             </div>
+            <div className={`mythic-tasking-parameter-preview${commandParameterPreview.state === "parameters" ? "" : " mythic-tasking-parameter-preview-empty-state"}`}>
+                <div className="mythic-tasking-parameter-preview-heading">
+                    <TerminalIcon fontSize="small" />
+                    <span>CLI parameters</span>
+                </div>
+                {commandParameterPreview.state === "parameters" ? (
+                    <div className="mythic-tasking-parameter-preview-chip-row">
+                    {commandParameterPreview.activeParameter &&
+                        <TaskingParameterPreviewChip key={"active" + commandParameterPreview.activeParameter.id} parameter={commandParameterPreview.activeParameter} active={true} />
+                    }
+                    {commandParameterPreview.requiredParameters.slice(0, commandParameterPreview.activeParameter ? 5 : 6).map((parameter) => (
+                        <TaskingParameterPreviewChip key={"required" + parameter.id} parameter={parameter} required={true} />
+                    ))}
+                    {commandParameterPreview.requiredParameters.length > (commandParameterPreview.activeParameter ? 5 : 6) &&
+                        <span className="mythic-tasking-parameter-preview-more">
+                            +{commandParameterPreview.requiredParameters.length - (commandParameterPreview.activeParameter ? 5 : 6)} required
+                        </span>
+                    }
+                    {commandParameterPreview.optionalParameters.slice(0, commandParameterPreview.requiredParameters.length > 0 ? 4 : (commandParameterPreview.activeParameter ? 5 : 6)).map((parameter) => (
+                        <TaskingParameterPreviewChip key={"optional" + parameter.id} parameter={parameter} />
+                    ))}
+                    {commandParameterPreview.optionalParameters.length > (commandParameterPreview.requiredParameters.length > 0 ? 4 : (commandParameterPreview.activeParameter ? 5 : 6)) &&
+                        <span className="mythic-tasking-parameter-preview-more">
+                            +{commandParameterPreview.optionalParameters.length - (commandParameterPreview.requiredParameters.length > 0 ? 4 : (commandParameterPreview.activeParameter ? 5 : 6))} optional
+                        </span>
+                    }
+                    </div>
+                ) : (
+                    <div className="mythic-tasking-parameter-preview-empty">
+                        <TerminalIcon fontSize="small" />
+                        <span>{commandParameterPreview.message}</span>
+                    </div>
+                )}
+            </div>
             {openFilterOptionsDialog &&
                 <MythicDialog fullWidth={true} maxWidth="md" open={openFilterOptionsDialog}
                               onClose={() => {

MythicReactUI/src/themes/GlobalStyles.js

@@ -502,6 +502,132 @@ tspan {
     justify-content: center;
     width: 28px;
 }
+.mythic-tasking-parameter-preview {
+    align-items: center;
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? alpha(props.theme.palette.common.white, 0.035) : alpha(props.theme.palette.common.black, 0.018)};
+    border: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    display: grid;
+    gap: 0.55rem;
+    grid-template-columns: minmax(7.5rem, auto) minmax(0, 1fr);
+    margin-top: 0.45rem;
+    min-height: 42px;
+    min-width: 0;
+    padding: 0.45rem 0.55rem;
+}
+.mythic-tasking-parameter-preview-empty-state {
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? alpha(props.theme.palette.common.white, 0.025) : alpha(props.theme.palette.common.black, 0.014)};
+}
+.mythic-tasking-parameter-preview-heading {
+    align-items: center;
+    display: flex;
+    gap: 0.35rem;
+    min-width: 0;
+    color: ${(props) => props.theme.palette.text.secondary};
+    font-size: 0.72rem;
+    font-weight: 850;
+    line-height: 1.2;
+}
+.mythic-tasking-parameter-preview-heading svg {
+    color: ${(props) => props.theme.palette.text.disabled};
+    font-size: 1rem;
+}
+.mythic-tasking-parameter-preview-more {
+    align-items: center;
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? "rgba(255,255,255,0.045)" : "rgba(0,0,0,0.025)"};
+    border: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    color: ${(props) => props.theme.palette.text.secondary};
+    display: inline-flex;
+    flex: 0 0 auto;
+    font-size: 0.68rem;
+    font-weight: 750;
+    min-height: 22px;
+    padding: 0.12rem 0.4rem;
+}
+.mythic-tasking-parameter-preview-chip-row {
+    align-items: center;
+    display: flex;
+    flex-wrap: nowrap;
+    gap: 0.28rem;
+    min-width: 0;
+    overflow-x: auto;
+    overflow-y: hidden;
+    padding-bottom: 0.02rem;
+}
+.mythic-tasking-parameter-preview-empty {
+    align-items: center;
+    color: ${(props) => props.theme.palette.text.disabled};
+    display: flex;
+    font-size: 0.74rem;
+    font-weight: 650;
+    gap: 0.35rem;
+    min-width: 0;
+}
+.mythic-tasking-parameter-preview-empty svg {
+    color: ${(props) => props.theme.palette.text.disabled};
+    font-size: 1rem;
+}
+.mythic-tasking-parameter-preview-empty span {
+    min-width: 0;
+    overflow: hidden;
+    text-overflow: ellipsis;
+    white-space: nowrap;
+}
+.mythic-tasking-parameter-preview-chip.MuiChip-root {
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? "rgba(255,255,255,0.045)" : "rgba(0,0,0,0.025)"};
+    border: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    color: ${(props) => props.theme.palette.text.secondary};
+    flex: 0 0 auto;
+    font-size: 0.68rem;
+    font-weight: 750;
+    height: 22px;
+    max-width: min(16rem, 100%);
+}
+.mythic-tasking-parameter-preview-chip-required.MuiChip-root {
+    background-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.14 : 0.08)};
+    border-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.36 : 0.24)};
+    color: ${(props) => props.theme.palette.warning.main};
+}
+.mythic-tasking-parameter-preview-chip-active.MuiChip-root {
+    background-color: ${(props) => alpha(props.theme.palette.info.main, props.theme.palette.mode === "dark" ? 0.14 : 0.08)};
+    border-color: ${(props) => alpha(props.theme.palette.info.main, props.theme.palette.mode === "dark" ? 0.42 : 0.28)};
+    color: ${(props) => props.theme.palette.info.main};
+}
+.mythic-tasking-parameter-preview-chip .MuiChip-label {
+    min-width: 0;
+    overflow: hidden;
+    text-overflow: ellipsis;
+}
+.mythic-tasking-parameter-preview-chip-label {
+    align-items: center;
+    display: inline-flex;
+    gap: 0.25rem;
+    min-width: 0;
+}
+.mythic-tasking-parameter-preview-chip-status {
+    background-color: ${(props) => alpha(props.theme.palette.info.main, props.theme.palette.mode === "dark" ? 0.16 : 0.1)};
+    border-radius: 999px;
+    color: inherit;
+    flex: 0 0 auto;
+    font-size: 0.58rem;
+    font-weight: 850;
+    line-height: 1;
+    padding: 0.14rem 0.28rem;
+}
+.mythic-tasking-parameter-preview-chip-name {
+    min-width: 0;
+    overflow: hidden;
+    text-overflow: ellipsis;
+}
+.mythic-tasking-parameter-preview-chip-type {
+    color: inherit;
+    flex: 0 0 auto;
+    font-size: 0.62rem;
+    font-weight: 650;
+    opacity: 0.66;
+}
 .mythic-tasking-reverse-search {
     align-items: center;
     background-color: ${(props) => props.theme.palette.mode === "dark" ? alpha(props.theme.palette.common.white, 0.045) : alpha(props.theme.palette.common.black, 0.024)};
@@ -721,6 +847,9 @@ tspan {
         max-width: 100%;
         width: 100% !important;
     }
+    .mythic-tasking-parameter-preview {
+        grid-template-columns: 1fr;
+    }
 }
 .mythic-task-parameters-title.MuiDialogTitle-root {
     background-image: ${getSectionHeaderGradient} !important;

mythic-docker/src/database/migrations/003003020_3.3.20.sql

@@ -57,6 +57,25 @@ on "public"."task" using btree (operation_id, display_id);
 create unique index if not exists callback_operation_display_id_unique
 on "public"."callback" using btree (operation_id, display_id);
 
+-- +migrate StatementBegin
+do $$
+begin
+    if exists (
+        select 1
+        from "public"."response"
+        where sequence_number is not null
+        group by task_id, sequence_number
+        having count(*) > 1
+    ) then
+        raise exception 'Cannot add response(task_id, sequence_number) uniqueness because duplicate response sequence numbers already exist';
+    end if;
+end $$;
+-- +migrate StatementEnd
+
+create unique index if not exists response_task_sequence_number_unique
+on "public"."response" using btree (task_id, sequence_number)
+where sequence_number is not null;
+
 -- +migrate StatementBegin
 create or replace function public.new_task_display_id() returns trigger
     language plpgsql
@@ -181,4 +200,5 @@ $$;
 
 drop index if exists "public"."callback_operation_display_id_unique";
 drop index if exists "public"."task_operation_display_id_unique";
+drop index if exists "public"."response_task_sequence_number_unique";
 drop table if exists "public"."operation_display_counters";

mythic-docker/src/rabbitmq/util_agent_message_actions_post_response.go

@@ -34,6 +34,31 @@ import (
 
 const POSTGRES_MAX_INT = 2147483647
 const POSTGRES_MAX_BIGINT = 9223372036854775807
+const defaultAgentMessagePostResponseWorkers = 8
+const defaultAgentMessagePostResponseQueueSize = 4096
+
+const insertAgentMessagePostResponseUserOutputQuery = `INSERT INTO response
+	("timestamp", task_id, response, sequence_number, operation_id)
+	VALUES ($1, $2, $3, $4, $5)
+	ON CONFLICT (task_id, sequence_number) WHERE sequence_number IS NOT NULL DO NOTHING
+	RETURNING id`
+
+const selectAgentMessagePostResponseTasksQuery = `SELECT
+	task.agent_task_id, task.id, task.status, task.completed, task.status_timestamp_processed, task.operator_id, task.operation_id,
+	task.stdout, task.stderr, task.display_id,
+	task.eventstepinstance_id, task.apitokens_id,
+	callback.host "callback.host",
+	callback.user "callback.user",
+	callback.id "callback.id",
+	callback.display_id "callback.display_id",
+	callback.agent_callback_id "callback.agent_callback_id",
+	callback.mythictree_groups "callback.mythictree_groups",
+	payload.payload_type_id "callback.payload.payload_type_id",
+	payload.os "callback.payload.os"
+	FROM task
+	JOIN callback ON task.callback_id = callback.id
+	JOIN payload ON callback.registered_payload_id = payload.id
+	WHERE task.agent_task_id IN (?)`
 
 type agentMessagePostResponseMessage struct {
 	Responses []agentMessagePostResponse `json:"responses" mapstructure:"responses" xml:"responses"`
@@ -275,11 +300,47 @@ type agentAgentMessagePostResponseChannelMessage struct {
 	Task        databaseStructs.Task
 }
 
-var asyncAgentMessagePostResponseChannel = make(chan agentAgentMessagePostResponseChannelMessage, 100)
+var asyncAgentMessagePostResponseWorkerChannels []chan agentAgentMessagePostResponseChannelMessage
+var asyncAgentMessagePostResponseWorkersOnce sync.Once
 
 var responseInterceptMapOperationIDToEventGroupID = make(map[int]int)
 var responseInterceptMapLock = sync.RWMutex{}
 
+func getAgentMessagePostResponseWorkerCount() int {
+	if utils.MythicConfig.AgentMessagePostResponseWorkers > 0 {
+		return utils.MythicConfig.AgentMessagePostResponseWorkers
+	}
+	return defaultAgentMessagePostResponseWorkers
+}
+
+func getAgentMessagePostResponseQueueSize() int {
+	if utils.MythicConfig.AgentMessagePostResponseQueueSize > 0 {
+		return utils.MythicConfig.AgentMessagePostResponseQueueSize
+	}
+	return defaultAgentMessagePostResponseQueueSize
+}
+
+func initializeAsyncAgentMessagePostResponseWorkers() {
+	UpdateCachedResponseIntercept()
+	workerCount := getAgentMessagePostResponseWorkerCount()
+	queueSize := getAgentMessagePostResponseQueueSize()
+	asyncAgentMessagePostResponseWorkerChannels = make([]chan agentAgentMessagePostResponseChannelMessage, workerCount)
+	for i := 0; i < workerCount; i++ {
+		asyncAgentMessagePostResponseWorkerChannels[i] = make(chan agentAgentMessagePostResponseChannelMessage, queueSize)
+		go listenForAsyncAgentMessagePostResponseWorker(i, asyncAgentMessagePostResponseWorkerChannels[i])
+	}
+}
+
+func enqueueAsyncAgentMessagePostResponse(msg agentAgentMessagePostResponseChannelMessage) {
+	asyncAgentMessagePostResponseWorkersOnce.Do(initializeAsyncAgentMessagePostResponseWorkers)
+	workerCount := len(asyncAgentMessagePostResponseWorkerChannels)
+	workerID := msg.Task.ID % workerCount
+	if workerID < 0 {
+		workerID = 0
+	}
+	asyncAgentMessagePostResponseWorkerChannels[workerID] <- msg
+}
+
 func UpdateCachedResponseIntercept() {
 	eventGroups := []databaseStructs.EventGroup{}
 	err := database.DB.Select(&eventGroups, `SELECT id, operation_id 
@@ -298,52 +359,98 @@ func UpdateCachedResponseIntercept() {
 		responseInterceptMapOperationIDToEventGroupID = make(map[int]int)
 		return
 	}
+	newResponseInterceptMap := make(map[int]int, len(eventGroups))
 	for _, eventGroup := range eventGroups {
-		responseInterceptMapOperationIDToEventGroupID[eventGroup.OperationID] = eventGroup.ID
+		newResponseInterceptMap[eventGroup.OperationID] = eventGroup.ID
 	}
+	responseInterceptMapOperationIDToEventGroupID = newResponseInterceptMap
 }
+
+func getCachedResponseInterceptEventGroupID(operationID int) (int, bool) {
+	responseInterceptMapLock.RLock()
+	eventGroupID, ok := responseInterceptMapOperationIDToEventGroupID[operationID]
+	responseInterceptMapLock.RUnlock()
+	return eventGroupID, ok
+}
+
 func listenForAsyncAgentMessagePostResponseContent() {
-	UpdateCachedResponseIntercept()
-	for {
-		msg := <-asyncAgentMessagePostResponseChannel
-		// force chunking user_output can be useful, but might also affect command's browserscripts
-		//chunkSize := 1024 * 1024
-		//chunks := int(len(msg.Response)/chunkSize) + 1 // 1MB chunks
-		//for j := 0; j < chunks; j++ {
-		//	currentChunkStart := j * chunkSize
-		//	currentChunkEnd := currentChunkStart + chunkSize
-		//	if currentChunkEnd > len(msg.Response) {
-		//		currentChunkEnd = len(msg.Response)
-		//	}
-		//	chunkBuf := msg.Response[currentChunkStart:currentChunkEnd]
-		responseInterceptMapLock.RLock()
-		if eventGroupID, ok := responseInterceptMapOperationIDToEventGroupID[msg.Task.OperationID]; ok {
-			output := ""
-			responseID := handleAgentMessagePostResponseUserOutput(msg.Task, agentMessagePostResponse{
-				TaskID:     msg.Task.AgentTaskID,
-				UserOutput: &output,
-			}, false)
-			if responseID > 0 {
-				// we have an interception possibility, so send that off for processing
-				EventingChannel <- EventNotification{
-					Trigger:               eventing.TriggerResponseIntercept,
-					OperationID:           msg.Task.OperationID,
-					EventGroupID:          eventGroupID,
-					ResponseID:            responseID,
-					TaskID:                msg.Task.ID,
-					ResponseInterceptData: msg.Response,
-				}
-			}
-		} else {
-			handleAgentMessagePostResponseUserOutput(msg.Task, agentMessagePostResponse{
-				TaskID:     msg.Task.AgentTaskID,
-				UserOutput: &msg.Response,
-			}, true)
-		}
-		responseInterceptMapLock.RUnlock()
-		//}
+	asyncAgentMessagePostResponseWorkersOnce.Do(initializeAsyncAgentMessagePostResponseWorkers)
+}
+
+func listenForAsyncAgentMessagePostResponseWorker(workerID int, input <-chan agentAgentMessagePostResponseChannelMessage) {
+	insertStatement, err := database.DB.Preparex(insertAgentMessagePostResponseUserOutputQuery)
+	if err != nil {
+		logging.LogError(err, "Failed to prepare user_output insert statement for post_response worker", "worker_id", workerID)
+	}
+	if insertStatement != nil {
+		defer insertStatement.Close()
+	}
+	for msg := range input {
+		processAsyncAgentMessagePostResponseContent(msg, insertStatement)
 	}
 }
+
+func processAsyncAgentMessagePostResponseContent(msg agentAgentMessagePostResponseChannelMessage, insertStatement *sqlx.Stmt) {
+	// force chunking user_output can be useful, but might also affect command's browserscripts
+	//chunkSize := 1024 * 1024
+	//chunks := int(len(msg.Response)/chunkSize) + 1 // 1MB chunks
+	//for j := 0; j < chunks; j++ {
+	//	currentChunkStart := j * chunkSize
+	//	currentChunkEnd := currentChunkStart + chunkSize
+	//	if currentChunkEnd > len(msg.Response) {
+	//		currentChunkEnd = len(msg.Response)
+	//	}
+	//	chunkBuf := msg.Response[currentChunkStart:currentChunkEnd]
+	if eventGroupID, ok := getCachedResponseInterceptEventGroupID(msg.Task.OperationID); ok {
+		output := ""
+		responseID := insertAgentMessagePostResponseUserOutput(msg.Task, output, msg.SequenceNum, false, insertStatement)
+		if responseID > 0 {
+			// we have an interception possibility, so send that off for processing
+			EventingChannel <- EventNotification{
+				Trigger:               eventing.TriggerResponseIntercept,
+				OperationID:           msg.Task.OperationID,
+				EventGroupID:          eventGroupID,
+				ResponseID:            responseID,
+				TaskID:                msg.Task.ID,
+				ResponseInterceptData: msg.Response,
+			}
+		}
+	} else {
+		insertAgentMessagePostResponseUserOutput(msg.Task, msg.Response, msg.SequenceNum, true, insertStatement)
+	}
+	//}
+}
+
+func getAgentMessagePostResponseTasks(responses []agentMessagePostResponse) (map[string]databaseStructs.Task, error) {
+	taskIDs := make([]string, 0, len(responses))
+	seenTaskIDs := make(map[string]bool, len(responses))
+	for _, response := range responses {
+		if response.TaskID == "" || seenTaskIDs[response.TaskID] {
+			continue
+		}
+		seenTaskIDs[response.TaskID] = true
+		taskIDs = append(taskIDs, response.TaskID)
+	}
+	tasksByAgentTaskID := make(map[string]databaseStructs.Task, len(taskIDs))
+	if len(taskIDs) == 0 {
+		return tasksByAgentTaskID, nil
+	}
+
+	query, args, err := sqlx.In(selectAgentMessagePostResponseTasksQuery, taskIDs)
+	if err != nil {
+		return tasksByAgentTaskID, err
+	}
+	query = database.DB.Rebind(query)
+	tasks := []databaseStructs.Task{}
+	if err = database.DB.Select(&tasks, query, args...); err != nil {
+		return tasksByAgentTaskID, err
+	}
+	for _, task := range tasks {
+		tasksByAgentTaskID[task.AgentTaskID] = task
+	}
+	return tasksByAgentTaskID, nil
+}
+
 func handleAgentMessagePostResponse(incoming *map[string]interface{}, uUIDInfo *cachedUUIDInfo) (map[string]interface{}, error) {
 	// got message:
 	/*
@@ -374,6 +481,11 @@ func handleAgentMessagePostResponse(incoming *map[string]interface{}, uUIDInfo *
 		return map[string]interface{}{}, err
 	}
 	responses := []map[string]interface{}{}
+	cachedTaskData, err = getAgentMessagePostResponseTasks(agentMessage.Responses)
+	if err != nil {
+		logging.LogError(err, "Failed to batch load tasks for post_response")
+	}
+	tasksToUpdate := make(map[string]databaseStructs.Task, len(cachedTaskData))
 	// iterate over the agent messages
 	for i, _ := range agentMessage.Responses {
 		mythicResponse := map[string]interface{}{
@@ -382,33 +494,16 @@ func handleAgentMessagePostResponse(incoming *map[string]interface{}, uUIDInfo *
 		}
 		//logging.LogDebug("Got response data from agent", "response data", agentResponse, "extra keys", agentResponse.Other)
 		// every response should be tied to some task
-		currentTask := databaseStructs.Task{AgentTaskID: agentMessage.Responses[i].TaskID}
-		if _, ok := cachedTaskData[currentTask.AgentTaskID]; ok {
-			currentTask = cachedTaskData[currentTask.AgentTaskID]
-		} else {
-			err = database.DB.Get(&currentTask, `SELECT
-			task.id, task.status, task.completed, task.status_timestamp_processed, task.operator_id, task.operation_id,
-			task.stdout, task.stderr, task.display_id,
-			task.eventstepinstance_id, task.apitokens_id,
-			callback.host "callback.host",
-			callback.user "callback.user",
-			callback.id "callback.id",
-			callback.display_id "callback.display_id",
-			callback.agent_callback_id "callback.agent_callback_id",
-			callback.mythictree_groups "callback.mythictree_groups",
-			payload.payload_type_id "callback.payload.payload_type_id",
-			payload.os "callback.payload.os"
-			FROM task
-			JOIN callback ON task.callback_id = callback.id
-			JOIN payload ON callback.registered_payload_id = payload.id
-			WHERE task.agent_task_id=$1`, currentTask.AgentTaskID)
-			if err != nil {
-				logging.LogError(err, "Failed to find task", "task id", currentTask.AgentTaskID)
-				mythicResponse["status"] = "error"
-				mythicResponse["error"] = "Failed to find task"
-				responses = append(responses, mythicResponse)
-				continue
-			}
+		currentTask, ok := tasksToUpdate[agentMessage.Responses[i].TaskID]
+		if !ok {
+			currentTask, ok = cachedTaskData[agentMessage.Responses[i].TaskID]
+		}
+		if !ok {
+			logging.LogError(nil, "Failed to find task", "task id", agentMessage.Responses[i].TaskID)
+			mythicResponse["status"] = "error"
+			mythicResponse["error"] = "Failed to find task"
+			responses = append(responses, mythicResponse)
+			continue
 		}
 
 		// always process here
@@ -481,10 +576,11 @@ func handleAgentMessagePostResponse(incoming *map[string]interface{}, uUIDInfo *
 		if agentMessage.Responses[i].UserOutput != nil && *agentMessage.Responses[i].UserOutput != "" {
 			// do it in the background - the agent doesn't need the result of this directly
 			//handleAgentMessagePostResponseUserOutput(currentTask, agentResponse, true)
-			asyncAgentMessagePostResponseChannel <- agentAgentMessagePostResponseChannelMessage{
-				Task:     currentTask,
-				Response: *agentMessage.Responses[i].UserOutput,
-			}
+			enqueueAsyncAgentMessagePostResponse(agentAgentMessagePostResponseChannelMessage{
+				Task:        currentTask,
+				Response:    *agentMessage.Responses[i].UserOutput,
+				SequenceNum: agentMessage.Responses[i].SequenceNumber,
+			})
 		}
 		if agentMessage.Responses[i].Stdout != nil {
 			currentTask.Stdout += *agentMessage.Responses[i].Stdout
@@ -553,14 +649,14 @@ func handleAgentMessagePostResponse(incoming *map[string]interface{}, uUIDInfo *
 		// this section always happens
 		reflectBackOtherKeys(&mythicResponse, &agentMessage.Responses[i].Other)
 		responses = append(responses, mythicResponse)
-		cachedTaskData[currentTask.AgentTaskID] = currentTask
+		tasksToUpdate[currentTask.AgentTaskID] = currentTask
 	}
 	response := map[string]interface{}{}
 	response["responses"] = responses
 	reflectBackOtherKeys(&response, &agentMessage.Other)
 	// remove responses so that we don't accidentally process it twice
 	delete(*incoming, "responses")
-	for _, currentTask := range cachedTaskData {
+	for _, currentTask := range tasksToUpdate {
 		// always updating at least the timestamp for the last thing that happened
 		_, err = database.DB.NamedExec(`UPDATE task SET
 				status=:status, completed=:completed, status_timestamp_processed=:status_timestamp_processed, "timestamp"=:timestamp,
@@ -647,44 +743,46 @@ func handleAgentMessagePostResponse(incoming *map[string]interface{}, uUIDInfo *
 
 }
 func handleAgentMessagePostResponseUserOutput(task databaseStructs.Task, agentResponse agentMessagePostResponse, emitNotification bool) int {
+	if agentResponse.UserOutput == nil {
+		return 0
+	}
+	return insertAgentMessagePostResponseUserOutput(task, *agentResponse.UserOutput, agentResponse.SequenceNumber, emitNotification, nil)
+}
+
+func insertAgentMessagePostResponseUserOutput(task databaseStructs.Task, userOutput string, sequenceNumber *int64, emitNotification bool, insertStatement *sqlx.Stmt) int {
 	responseOutput := databaseStructs.Response{
 		Timestamp:   time.Now().UTC(),
 		TaskID:      task.ID,
-		Response:    []byte(*agentResponse.UserOutput),
+		Response:    []byte(userOutput),
 		OperationID: task.OperationID,
 	}
-	if len(*agentResponse.UserOutput) == 0 && emitNotification {
+	if len(userOutput) == 0 && emitNotification {
 		//logging.LogError(nil, "Tried to add response of 0 bytes, returning")
 		return 0
 	}
-	if agentResponse.SequenceNumber != nil {
-		// if we're tracking sequence numbers, then there shouldn't be a matching sequence number for this task to prevent replays
+	if sequenceNumber != nil {
 		responseOutput.SequenceNumber.Valid = true
-		responseOutput.SequenceNumber.Int64 = *agentResponse.SequenceNumber
-		if _, err := database.DB.NamedQuery(`SELECT id 
-		FROM response
-		WHERE sequence_number=:sequence_number AND task_id=:task_id`, responseOutput); errors.Is(err, sql.ErrNoRows) {
-			// we don't have this sequence number for this task yet, so we're safe to insert it
-			logging.LogInfo("Sequence number is not NULL!")
-		} else if err != nil {
-			logging.LogError(err, "Failed to fetch responses when looking for an existing sequence number")
-			return 0
-		} else {
-			// this sequence number and task do exist, so don't insert it
-			logging.LogError(nil, "Got a duplicate sequence number for a response", "task_id", responseOutput.TaskID, "sequence number", *agentResponse.SequenceNumber)
+		responseOutput.SequenceNumber.Int64 = *sequenceNumber
+	}
+
+	args := []interface{}{
+		responseOutput.Timestamp,
+		responseOutput.TaskID,
+		responseOutput.Response,
+		responseOutput.SequenceNumber,
+		responseOutput.OperationID,
+	}
+	var err error
+	if insertStatement != nil {
+		err = insertStatement.QueryRowx(args...).Scan(&responseOutput.ID)
+	} else {
+		err = database.DB.QueryRowx(insertAgentMessagePostResponseUserOutputQuery, args...).Scan(&responseOutput.ID)
+	}
+	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) && responseOutput.SequenceNumber.Valid {
+			logging.LogError(nil, "Got a duplicate sequence number for a response", "task_id", responseOutput.TaskID, "sequence number", responseOutput.SequenceNumber.Int64)
 			return 0
 		}
-	}
-	statement, err := database.DB.PrepareNamed(`INSERT INTO response
-		("timestamp", task_id, response, sequence_number, operation_id)
-		VALUES (:timestamp, :task_id, :response, :sequence_number, :operation_id)
-		RETURNING id`)
-	if err != nil {
-		logging.LogError(err, "Failed to prepare new named statement for user_output", "task_id", responseOutput.TaskID)
-		return 0
-	}
-	err = statement.Get(&responseOutput.ID, responseOutput)
-	if err != nil {
 		logging.LogError(err, "Failed to insert new user_output", "task_id", responseOutput.TaskID)
 		return 0
 	}
@@ -1835,11 +1933,11 @@ func HandleAgentMessagePostResponseFileBrowser(task databaseStructs.Task, fileBr
 				logging.LogError(err, "failed to marshal filebrowser data to JSON")
 				return
 			}
-			asyncAgentMessagePostResponseChannel <- agentAgentMessagePostResponseChannelMessage{
+			enqueueAsyncAgentMessagePostResponse(agentAgentMessagePostResponseChannelMessage{
 				Task:        task,
 				Response:    string(outputBytes),
 				SequenceNum: nil,
-			}
+			})
 		}(fileBrowser)
 	}
 	pathData, err := utils.SplitFilePathGetHost(fileBrowser.ParentPath, fileBrowser.Name, []string{})
@@ -2761,11 +2859,11 @@ func handleAgentMessagePostResponseCustomBrowser(task databaseStructs.Task, agen
 				logging.LogError(err, "failed to marshal filebrowser data to JSON")
 				return
 			}
-			asyncAgentMessagePostResponseChannel <- agentAgentMessagePostResponseChannelMessage{
+			enqueueAsyncAgentMessagePostResponse(agentAgentMessagePostResponseChannelMessage{
 				Task:        task,
 				Response:    string(outputBytes),
 				SequenceNum: nil,
-			}
+			})
 		}(agentCustomBrowser)
 	}
 	if agentCustomBrowser.Entries == nil || len(*agentCustomBrowser.Entries) == 0 {

mythic-docker/src/utils/config.go

@@ -34,6 +34,8 @@ type Config struct {
 	MythicServerAllowInviteLinks            bool
 	ServerVersion                           string
 	MythicServerAllowWebhooksOnNewCallbacks bool
+	AgentMessagePostResponseWorkers         int
+	AgentMessagePostResponseQueueSize       int
 
 	// rabbitmq configuration
 	RabbitmqHost     string
@@ -74,6 +76,8 @@ func Initialize() {
 	mythicEnv.SetDefault("debug_level", "warning")
 	mythicEnv.SetDefault("mythic_server_allow_invite_links", false)
 	mythicEnv.SetDefault("mythic_server_allow_webhooks_on_new_callbacks", true)
+	mythicEnv.SetDefault("agent_message_post_response_workers", 8)
+	mythicEnv.SetDefault("agent_message_post_response_queue_size", 4096)
 	// postgres configuration
 	mythicEnv.SetDefault("postgres_host", "mythic_postgres")
 	mythicEnv.SetDefault("postgres_port", 5432)
@@ -167,6 +171,8 @@ func setConfigFromEnv(mythicEnv *viper.Viper) {
 	MythicConfig.DefaultOperationChannel = mythicEnv.GetString("default_operation_webhook_channel")
 	MythicConfig.GlobalServerName = mythicEnv.GetString("global_server_name")
 	MythicConfig.MythicServerAllowInviteLinks = mythicEnv.GetBool("mythic_server_allow_invite_links")
+	MythicConfig.AgentMessagePostResponseWorkers = mythicEnv.GetInt("agent_message_post_response_workers")
+	MythicConfig.AgentMessagePostResponseQueueSize = mythicEnv.GetInt("agent_message_post_response_queue_size")
 	allowedIPBlocks := []*net.IPNet{}
 	for _, ipBlock := range strings.Split(mythicEnv.GetString("allowed_ip_blocks"), ",") {
 		if _, subnet, err := net.ParseCIDR(ipBlock); err != nil {