better rabbitmq handling and task filter ui updates

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

@@ -1,63 +1,40 @@
 import React, {useEffect} from 'react';
-import { styled } from '@mui/material/styles';
+import Box from '@mui/material/Box';
 import Button from '@mui/material/Button';
-import DialogActions from '@mui/material/DialogActions';
+import Checkbox from '@mui/material/Checkbox';
+import Chip from '@mui/material/Chip';
 import DialogContent from '@mui/material/DialogContent';
 import DialogTitle from '@mui/material/DialogTitle';
-import MythicTextField from '../../MythicComponents/MythicTextField';
-import Switch from '@mui/material/Switch';
-import Select from '@mui/material/Select';
-import Chip from '@mui/material/Chip';
-import Input from '@mui/material/Input';
-import InputLabel from '@mui/material/InputLabel';
-import MenuItem from '@mui/material/MenuItem';
 import FormControl from '@mui/material/FormControl';
 import ListItemText from '@mui/material/ListItemText';
-import Checkbox from '@mui/material/Checkbox';
+import MenuItem from '@mui/material/MenuItem';
+import OutlinedInput from '@mui/material/OutlinedInput';
+import Select from '@mui/material/Select';
+import Switch from '@mui/material/Switch';
+import Typography from '@mui/material/Typography';
+import ClearIcon from '@mui/icons-material/Clear';
+import FilterAltIcon from '@mui/icons-material/FilterAlt';
+import MythicTextField from '../../MythicComponents/MythicTextField';
 import {useQuery, gql } from '@apollo/client';
 import { meState } from '../../../cache';
 import {useReactiveVar} from '@apollo/client';
+import {
+  MythicDialogBody,
+  MythicDialogButton,
+  MythicDialogFooter,
+  MythicDialogGrid,
+  MythicDialogSection,
+  MythicFormField,
+  MythicFormSwitchRow
+} from '../../MythicComponents/MythicDialogLayout';
 
-const PREFIX = 'CallbacksTabsTaskingFilterDialog';
-
-const classes = {
-  formControl: `${PREFIX}-formControl`,
-  chips: `${PREFIX}-chips`,
-  chip: `${PREFIX}-chip`,
-  noLabel: `${PREFIX}-noLabel`
-};
-
-const Root = styled('div')((
-  {
-    theme
-  }
-) => ({
-  [`& .${classes.formControl}`]: {
-    margin: theme.spacing(1),
-    width: "100%",
-  },
-
-  [`& .${classes.chips}`]: {
-    display: 'flex',
-    flexWrap: 'wrap',
-  },
-
-  [`& .${classes.chip}`]: {
-    margin: 2,
-  },
-
-  [`& .${classes.noLabel}`]: {
-    marginTop: theme.spacing(2),
-  }
-}));
-
-const ITEM_HEIGHT = 48;
+const ITEM_HEIGHT = 42;
 const ITEM_PADDING_TOP = 8;
 const MenuProps = {
   PaperProps: {
     style: {
-      maxHeight: ITEM_HEIGHT * 4.5 + ITEM_PADDING_TOP,
-      width: 250,
+      maxHeight: ITEM_HEIGHT * 6 + ITEM_PADDING_TOP,
+      width: 320,
     },
   },
   variant: "menu",
@@ -72,6 +49,60 @@ query operatorQuery($operation_id: Int!) {
     }
   }
 }`;
+
+const selectedLabel = (values, singular, plural = singular) => {
+  if(values.length === 0){
+    return "";
+  }
+  return `${values.length} ${values.length === 1 ? singular : plural}`;
+}
+
+const FilterSummaryChip = ({icon, label, muted=false}) => (
+  <Chip
+    className={`mythic-tasking-filter-summary-chip${muted ? " mythic-tasking-filter-summary-chip-muted" : ""}`}
+    icon={icon}
+    label={label}
+    size="small"
+  />
+)
+
+const MultiSelectField = ({label, value, options, onChange, emptyLabel}) => {
+  const renderValue = (selected) => {
+    if(selected.length === 0){
+      return <span className="mythic-tasking-filter-select-empty">{emptyLabel}</span>
+    }
+    return (
+      <Box className="mythic-tasking-filter-select-chips">
+        {selected.map((selectedValue) => (
+          <Chip key={selectedValue} label={selectedValue} size="small" className="mythic-tasking-filter-selected-chip" />
+        ))}
+      </Box>
+    );
+  }
+
+  return (
+    <FormControl fullWidth={true} size="small" className="mythic-tasking-filter-select">
+      <Select
+        multiple
+        displayEmpty
+        value={value}
+        onChange={onChange}
+        input={<OutlinedInput />}
+        renderValue={renderValue}
+        MenuProps={MenuProps}
+        aria-label={label}
+      >
+        {options.map((name) => (
+          <MenuItem key={name} value={name}>
+            <Checkbox color="primary" checked={value.indexOf(name) > -1} />
+            <ListItemText primary={name} />
+          </MenuItem>
+        ))}
+      </Select>
+    </FormControl>
+  )
+}
+
 export function CallbacksTabsTaskingFilterDialog(props) {
   const me = useReactiveVar(meState);
   const [onlyOperators, setOnlyOperators] = React.useState([]);
@@ -111,7 +142,7 @@ export function CallbacksTabsTaskingFilterDialog(props) {
       const commandOptionNames = props.filterCommandOptions.map(c => c.cmd);
       setCommandOptions(commandOptionNames);
     }
-  }, [props.filterOptions]);
+  }, [props.filterOptions, props.filterCommandOptions]);
   const onSubmit = () => {
     props.onSubmit({
       "operatorsList": onlyOperators,
@@ -133,17 +164,20 @@ export function CallbacksTabsTaskingFilterDialog(props) {
     setHideErrors(event.target.checked);
   }
   const handleOperatorChange = (event) => {
-    setOnlyOperators(event.target.value);
+    const value = typeof event.target.value === "string" ? event.target.value.split(",") : event.target.value;
+    setOnlyOperators(value);
   }
   const handleOnlyCommandsChange = (event) => {
-    setOnlyCommands(event.target.value);
-    if(event.target.value.length > 0){
+    const value = typeof event.target.value === "string" ? event.target.value.split(",") : event.target.value;
+    setOnlyCommands(value);
+    if(value.length > 0){
       setEverythingBut([]);
     }
   }
   const handleEverythingButChange = (event) => {
-    setEverythingBut(event.target.value);
-    if(event.target.value.length > 0){
+    const value = typeof event.target.value === "string" ? event.target.value.split(",") : event.target.value;
+    setEverythingBut(value);
+    if(value.length > 0){
       setOnlyCommands([]);
     }
   }
@@ -153,110 +187,178 @@ export function CallbacksTabsTaskingFilterDialog(props) {
   const clearAllEverythingBut = () => {
     setEverythingBut([]);
   }
+  const clearAllFilters = () => {
+    setOnlyOperators([]);
+    setOnlyHasComments(false);
+    setOnlyCommands([]);
+    setEverythingBut([]);
+    setOnlyParameters("");
+    setHideErrors(false);
+  }
+  const activeFilters = [
+    selectedLabel(onlyOperators, "operator"),
+    onlyHasComments ? "Has comments" : "",
+    hideErrors ? "Errors hidden" : "",
+    selectedLabel(onlyCommands, "included command", "included commands"),
+    selectedLabel(everythingBut, "excluded command", "excluded commands"),
+    onlyParameters !== "" ? "Parameter regex" : "",
+  ].filter(Boolean);
+
   return (
-    <Root>
-        <DialogTitle id="form-dialog-title">Filter Which Tasks Are Visible in this Callback</DialogTitle>
-        <DialogContent dividers={true} style={{overflow: "hidden"}}>
-            <React.Fragment>
-                <FormControl className={classes.formControl}>
-                <InputLabel id="operator-chip-label">Only Show Tasks by the Following Operators</InputLabel>
-                <Select
-                  labelId="operator-chip-label"
-                  multiple
-                  id="operator-chip"
+    <Box className="mythic-tasking-filter-dialog">
+      <DialogTitle id="mythic-draggable-title" className="mythic-tasking-filter-dialog-title">
+        <Box className="mythic-tasking-filter-title-row">
+          <Box className="mythic-tasking-filter-title-icon">
+            <FilterAltIcon fontSize="small" />
+          </Box>
+          <Box sx={{minWidth: 0}}>
+            <Typography component="div" className="mythic-tasking-filter-title-main">
+              Task visibility filters
+            </Typography>
+            <Typography component="div" className="mythic-tasking-filter-title-subtitle">
+              Control which tasks are shown for this callback.
+            </Typography>
+          </Box>
+        </Box>
+      </DialogTitle>
+      <DialogContent dividers={true} className="mythic-tasking-filter-dialog-content">
+        <MythicDialogBody compact={true}>
+          <Box className="mythic-tasking-filter-summary">
+            {activeFilters.length > 0 ? (
+              activeFilters.map((filterLabel) => (
+                <FilterSummaryChip key={filterLabel} icon={<FilterAltIcon />} label={filterLabel} />
+              ))
+            ) : (
+              <FilterSummaryChip muted={true} icon={<FilterAltIcon />} label="No active filters" />
+            )}
+          </Box>
+          <MythicDialogSection
+            title="Operator and task state"
+            description="Narrow the task list by operator, comments, or error state."
+          >
+            <MythicDialogGrid minWidth="18rem">
+              <MythicFormField
+                label="Operators"
+                description="Only show tasks created by selected operators."
+              >
+                <MultiSelectField
+                  label="Operators"
                   value={onlyOperators}
+                  options={operatorUsernames}
                   onChange={handleOperatorChange}
-                  input={<Input />}
-                  renderValue={(selected) => (
-                    <div className={classes.chips}>
-                      {selected.map((value) => (
-                        <Chip key={value} label={value} className={classes.chip} />
-                      ))}
-                    </div>
-                  )}
-                  MenuProps={MenuProps}
-                >
-                  {operatorUsernames.map((name) => (
-                    <MenuItem key={name} value={name}>
-                      <Checkbox color="primary" checked={onlyOperators.indexOf(name) > -1} />
-                      <ListItemText primary={name} />
-                    </MenuItem>
-                  ))}
-                </Select>
-              </FormControl>
-                Only Show Tasks with Comments: <Switch checked={onlyHasComments} onChange={handleCommentsChange} color="primary" name="Only Comments" inputProps={{'aria-label': 'primary checkbox'}}/>
-                <br/>
-              Hide Error Tasks:<Switch checked={hideErrors} onChange={handleHideErrorsChange} color="primary" name="Hide Errors" inputProps={{'aria-label': 'primary checkbox'}}/>
-              <FormControl className={classes.formControl}>
-                  <InputLabel id="include-chip-label">Only Show These Commands</InputLabel>
-                  <Select
-                    labelId="include-chip-label"
-                    multiple
-                    id="include-chip"
-                    value={onlyCommands}
-                    onChange={handleOnlyCommandsChange}
-                    input={<Input />}
-                    renderValue={(selected) => (
-                      <div className={classes.chips}>
-                        {selected.map((value) => (
-                          <Chip key={value} label={value} className={classes.chip} />
-                        ))}
-                      </div>
-                    )}
-                    MenuProps={MenuProps}
+                  emptyLabel="Any operator"
+                />
+              </MythicFormField>
+              <Box className="mythic-tasking-filter-switch-stack">
+                <MythicFormSwitchRow
+                  label="Only tasks with comments"
+                  description="Require at least one comment."
+                  control={
+                    <Switch
+                      checked={onlyHasComments}
+                      onChange={handleCommentsChange}
+                      color="primary"
+                      name="Only Comments"
+                      inputProps={{'aria-label': 'Only tasks with comments'}}
+                    />
+                  }
+                />
+                <MythicFormSwitchRow
+                  label="Hide error tasks"
+                  description="Remove tasks that are currently marked as errors."
+                  control={
+                    <Switch
+                      checked={hideErrors}
+                      onChange={handleHideErrorsChange}
+                      color="primary"
+                      name="Hide Errors"
+                      inputProps={{'aria-label': 'Hide error tasks'}}
+                    />
+                  }
+                />
+              </Box>
+            </MythicDialogGrid>
+          </MythicDialogSection>
+          <MythicDialogSection
+            title="Command scope"
+            description="Choose commands to include, or choose commands to hide."
+          >
+            <Box className="mythic-tasking-filter-command-grid">
+              <MythicFormField
+                label="Only show commands"
+                description="When set, only matching command names remain visible."
+              >
+                <MultiSelectField
+                  label="Only show commands"
+                  value={onlyCommands}
+                  options={commandOptions}
+                  onChange={handleOnlyCommandsChange}
+                  emptyLabel="No include filter"
+                />
+                {onlyCommands.length > 0 &&
+                  <Button
+                    className="mythic-tasking-filter-clear-button"
+                    onClick={clearAllOnlyCommands}
+                    size="small"
+                    startIcon={<ClearIcon fontSize="small" />}
+                    variant="outlined"
                   >
-                    {commandOptions.map((name) => (
-                      <MenuItem key={name} value={name}>
-                        <Checkbox color="primary" checked={onlyCommands.indexOf(name) > -1} />
-                        <ListItemText primary={name} />
-                      </MenuItem>
-                    ))}
-                  </Select>
-                </FormControl>
-              {onlyCommands.length > 0 &&
-                  <Button onClick={clearAllOnlyCommands} variant={"contained"}>Clear</Button>
-              }
-                <FormControl className={classes.formControl}>
-                  <InputLabel id="exclude-chip-label">Do Not Show These Commands</InputLabel>
-                  <Select
-                    labelId="exclude-chip-label"
-                    multiple
-                    id="exclude-chip"
-                    value={everythingBut}
-                    onChange={handleEverythingButChange}
-                    input={<Input />}
-                    renderValue={(selected) => (
-                      <div className={classes.chips}>
-                        {selected.map((value) => (
-                          <Chip key={value} label={value} className={classes.chip} />
-                        ))}
-                      </div>
-                    )}
-                    MenuProps={MenuProps}
+                    Clear include
+                  </Button>
+                }
+              </MythicFormField>
+              <Box className="mythic-tasking-filter-choice-divider">or</Box>
+              <MythicFormField
+                label="Hide commands"
+                description="When set, selected command names are removed from view."
+              >
+                <MultiSelectField
+                  label="Hide commands"
+                  value={everythingBut}
+                  options={commandOptions}
+                  onChange={handleEverythingButChange}
+                  emptyLabel="No exclude filter"
+                />
+                {everythingBut.length > 0 &&
+                  <Button
+                    className="mythic-tasking-filter-clear-button"
+                    onClick={clearAllEverythingBut}
+                    size="small"
+                    startIcon={<ClearIcon fontSize="small" />}
+                    variant="outlined"
                   >
-                    {commandOptions.map((name) => (
-                      <MenuItem key={name} value={name}>
-                        <Checkbox color="primary" checked={everythingBut.indexOf(name) > -1} />
-                        <ListItemText primary={name} />
-                      </MenuItem>
-                    ))}
-                  </Select>
-                </FormControl>
-              {everythingBut.length > 0 &&
-                  <Button onClick={clearAllEverythingBut} variant={"contained"}>Clear</Button>
-              }
-                <MythicTextField value={onlyParameters} onChange={onChange} name="Only Show Tasks with the Following Parameter Regex"/>
-            </React.Fragment>
-        </DialogContent>
-        <DialogActions>
-          <Button onClick={props.onClose} variant="contained" >
-            Close
-          </Button>
-          <Button onClick={onSubmit} color="success" variant="contained" >
-            Filter
-          </Button>
-        </DialogActions>
-  </Root>
+                    Clear exclude
+                  </Button>
+                }
+              </MythicFormField>
+            </Box>
+          </MythicDialogSection>
+          <MythicDialogSection
+            title="Parameter matching"
+            description="Apply a regular expression to task parameters."
+          >
+            <MythicTextField
+              value={onlyParameters}
+              onChange={onChange}
+              name="Parameter regex"
+              placeholder="Regex to match task parameters"
+              marginBottom="0px"
+              marginTop="0px"
+            />
+          </MythicDialogSection>
+        </MythicDialogBody>
+      </DialogContent>
+      <MythicDialogFooter className="mythic-tasking-filter-dialog-actions">
+        <MythicDialogButton onClick={clearAllFilters}>
+          Clear all
+        </MythicDialogButton>
+        <MythicDialogButton onClick={props.onClose}>
+          Close
+        </MythicDialogButton>
+        <MythicDialogButton onClick={onSubmit} intent="primary">
+          Apply filters
+        </MythicDialogButton>
+      </MythicDialogFooter>
+    </Box>
   );
 }
-

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

@@ -3,6 +3,7 @@ import SendIcon from '@mui/icons-material/Send';
 import React from 'react';
 import {TextField} from '@mui/material';
 import TuneIcon from '@mui/icons-material/Tune';
+import TerminalIcon from '@mui/icons-material/Terminal';
 import { MythicDialog } from '../../MythicComponents/MythicDialog';
 import {CallbacksTabsTaskingFilterDialog} from './CallbacksTabsTaskingFilterDialog';
 import {CallbacksTabsTaskingInputTokenSelect} from './CallbacksTabsTaskingInputTokenSelect';
@@ -20,6 +21,7 @@ import {GetMythicSetting} from "../../MythicComponents/MythicSavedUserSetting";
 import {getSkewedNow} from "../../utilities/Time";
 import { useTheme } from '@mui/material/styles';
 import {MythicStyledTooltip} from "../../MythicComponents/MythicStyledTooltip";
+import {getReadableTextColor, isValidHexColor} from "../../MythicComponents/MythicColorInput";
 
 const GetLoadedCommandsSubscription = gql`
 subscription GetLoadedCommandsSubscription($callback_id: Int!){
@@ -169,6 +171,27 @@ const IsRepeatableCLIParameterType = (parameter_type) => {
             return false;
     }
 }
+const TaskingContextChip = ({title, label, value, color, callbackColor, emphasized=false}) => {
+    const safeColor = isValidHexColor(color) ? color : "#000000";
+    const borderColor = isValidHexColor(callbackColor) ? callbackColor : undefined;
+    const tooltipTitle = `${title}: ${value}`;
+    return (
+        <MythicStyledTooltip title={tooltipTitle}>
+            <span className={`mythic-tasking-context-chip${emphasized ? " mythic-tasking-context-chip-emphasized" : ""}`}
+                  style={{
+                      backgroundColor: safeColor,
+                      borderColor: borderColor,
+                      color: getReadableTextColor(safeColor),
+                  }}
+                  title={tooltipTitle}>
+                {label !== "" &&
+                    <span className="mythic-tasking-context-chip-label">{label}</span>
+                }
+                <span className="mythic-tasking-context-chip-value">{value}</span>
+            </span>
+        </MythicStyledTooltip>
+    )
+}
 
 export function CallbacksTabsTaskingInputPreMemo(props){
     const toastId = "tasking-toast-message";
@@ -1607,177 +1630,186 @@ export function CallbacksTabsTaskingInputPreMemo(props){
             inputRef.current.focus();
         }
     }, [props.focus])
+    const showTaskingContext = !hideTaskingContext.current;
+    const taskingContextChips = [
+        {
+            key: "impersonation_context",
+            title: "Impersonation Context",
+            label: "User",
+            value: callbackContext?.impersonation_context,
+            color: theme.taskContextImpersonationColor,
+            emphasized: true,
+        },
+        {
+            key: "user",
+            title: "User Context" + (callbackContext.integrity_level > 2 ? " (high integrity)" : ""),
+            label: "User",
+            value: callbackContext?.user ? `${callbackContext.user}${callbackContext.integrity_level > 2 ? "*" : ""}` : "",
+            color: theme.taskContextColor,
+            emphasized: callbackContext.integrity_level > 2,
+        },
+        {
+            key: "cwd",
+            title: "Current Working Directory",
+            label: "Dir",
+            value: callbackContext?.cwd,
+            color: theme.taskContextColor,
+        },
+        {
+            key: "host",
+            title: "Hostname",
+            label: "Host",
+            value: callbackContext?.host,
+            color: theme.taskContextColor,
+        },
+        {
+            key: "ip",
+            title: "First IP Address",
+            label: "IP",
+            value: callbackContext?.ip,
+            color: theme.taskContextColor,
+        },
+        {
+            key: "pid",
+            title: "Process ID",
+            label: "PID",
+            value: callbackContext?.pid,
+            color: theme.taskContextColor,
+        },
+        {
+            key: "architecture",
+            title: "Process Architecture",
+            label: "Arch",
+            value: callbackContext?.architecture,
+            color: theme.taskContextColor,
+        },
+        {
+            key: "process_short_name",
+            title: "Process Name",
+            label: "Process",
+            value: callbackContext?.process_short_name,
+            color: theme.taskContextColor,
+        },
+        {
+            key: "extra_info",
+            title: "Extra Callback Context",
+            label: "",
+            value: callbackContext?.extra_info,
+            color: theme.taskContextExtraColor,
+            emphasized: true,
+        },
+    ].filter((chip) => showTaskingContext && taskingContextFields.current.includes(chip.key) && chip.value !== undefined && chip.value !== "");
     return (
-        <div style={{position: "relative"}}>
+        <div className="mythic-tasking-composer">
             {backdropOpen && <Backdrop open={backdropOpen} style={{zIndex: 2, position: "absolute"}} invisible={false}>
                 <CircularProgress color="inherit" size={30}/>
             </Backdrop>
             }
             {reverseSearching &&
+                <div className="mythic-tasking-reverse-search">
+                    <Typography component="span" className="mythic-tasking-reverse-search-label">
+                        reverse-i-search
+                    </Typography>
+                    <TextField
+                        placeholder={"Search previous commands"}
+                        onKeyDown={onReverseSearchKeyDown}
+                        onChange={handleReverseSearchInputChange}
+                        size="small"
+                        color={"secondary"}
+                        autoFocus={true}
+                        variant="outlined"
+                        value={reverseSearchString}
+                        fullWidth={true}
+                        className="mythic-tasking-reverse-search-input"
+                        InputProps={{
+                            type: 'search',
+                        }}
+                    />
+                </div>
+            }
+            {taskingContextChips.length > 0 &&
+                <div className="mythic-tasking-context-row">
+                    {taskingContextChips.map((chip) => (
+                        <TaskingContextChip
+                            key={chip.key}
+                            title={chip.title}
+                            label={chip.label}
+                            value={chip.value}
+                            color={chip.color}
+                            callbackColor={callbackContext.color}
+                            emphasized={chip.emphasized}
+                        />
+                    ))}
+                </div>
+            }
+            <div className="mythic-tasking-command-row">
+                {tokenOptions.current.length > 0 ? (
+                    <CallbacksTabsTaskingInputTokenSelect
+                        options={tokenOptions.current}
+                        changeSelectedToken={props.changeSelectedToken}
+                        width={"13rem"}
+                        modern={true}
+                        className="mythic-tasking-token-select"
+                    />
+                ) : null}
                 <TextField
-                    placeholder={"Search previous commands"}
-                    onKeyDown={onReverseSearchKeyDown}
-                    onChange={handleReverseSearchInputChange}
+                    placeholder={"Task an agent..."}
+                    onKeyDown={onKeyDown}
+                    onChange={handleInputChange}
                     size="small"
                     color={"secondary"}
-                    autoFocus={true}
                     variant="outlined"
-                    value={reverseSearchString}
+                    multiline={true}
+                    maxRows={15}
+                    disabled={reverseSearching}
+                    value={message}
+                    autoFocus={true}
                     fullWidth={true}
+                    inputRef={inputRef}
+                    className="mythic-tasking-command-input"
                     InputProps={{
                         type: 'search',
-                        startAdornment: <React.Fragment><Typography
-                            style={{width: "10%"}}>reverse-i-search:</Typography></React.Fragment>
+                        spellCheck: false,
+                        autoFocus: true,
+                        startAdornment:
+                            <span className="mythic-tasking-command-prefix">
+                                <TerminalIcon fontSize="small" />
+                            </span>,
+                        endAdornment:
+                            <div className="mythic-tasking-action-row">
+                                {commandPayloadType !== "" &&
+                                    <MythicStyledTooltip title={commandPayloadType}>
+                                        <span className="mythic-tasking-payload-chip">
+                                            <MythicAgentSVGIcon payload_type={commandPayloadType}
+                                                                style={{width: "20px", height: "20px"}}/>
+                                        </span>
+                                    </MythicStyledTooltip>
+                                }
+                                {props.filterTasks &&
+                                    <MythicStyledTooltip title={activeFiltering ? "Adjust active task filters" : "Filter task history"}>
+                                        <IconButton
+                                            className={`mythic-tasking-action-button ${activeFiltering ? "mythic-tasking-action-button-warning" : "mythic-tasking-action-button-neutral"}`}
+                                            onClick={onClickFilter}
+                                            disableRipple={true}
+                                            disableFocusRipple={true}
+                                            size="small"
+                                            aria-label="Filter task history"><TuneIcon fontSize="small"/></IconButton>
+                                    </MythicStyledTooltip>
+                                }
+                                <MythicStyledTooltip title={"Submit task"}>
+                                    <IconButton
+                                        className="mythic-tasking-action-button mythic-tasking-action-button-success"
+                                        disableRipple={true}
+                                        disableFocusRipple={true}
+                                        onClick={onSubmitCommandLine}
+                                        size="small"
+                                        aria-label="Submit task"><SendIcon fontSize="small"/>
+                                    </IconButton>
+                                </MythicStyledTooltip>
+                            </div>
                     }}
                 />
-            }
-            {callbackContext?.impersonation_context !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("impersonation_context") &&
-                <MythicStyledTooltip title={"Impersonation Context"}>
-                    <span className={"rounded-tab"} style={{
-                        backgroundColor: theme.taskContextImpersonationColor,
-                        borderColor: callbackContext.color === "" ? "" : callbackContext.color
-                    }}>
-                        <b>{"User: "}</b>{callbackContext.impersonation_context}
-                    </span>
-                </MythicStyledTooltip>
-
-            }
-            {callbackContext?.user !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("user") &&
-                <MythicStyledTooltip title={"User Context" + (callbackContext.integrity_level > 2 ? " (high integrity)" : "")}>
-                    <span className={"rounded-tab"} style={{
-                        backgroundColor: theme.taskContextColor,
-                        borderColor: callbackContext.color === "" ? "" : callbackContext.color
-                    }}>
-                        <b>{"User: "}</b>{callbackContext.user}{callbackContext.integrity_level > 2 ? "*" : ""}
-                    </span>
-                </MythicStyledTooltip>
-            }
-            {callbackContext?.cwd !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("cwd") &&
-                <MythicStyledTooltip title={"Current Working Directory"}>
-                    <span className={"rounded-tab"} style={{
-                        backgroundColor: theme.taskContextColor,
-                        borderColor: callbackContext.color === "" ? "" : callbackContext.color
-                    }}>
-                        <b>{"Dir: "}</b>{callbackContext.cwd}
-                    </span>
-                </MythicStyledTooltip>
-            }
-            {callbackContext?.host !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("host") &&
-                <MythicStyledTooltip title={"Hostname"}>
-                    <span className={"rounded-tab"} style={{
-                        backgroundColor: theme.taskContextColor,
-                        borderColor: callbackContext.color === "" ? "" : callbackContext.color
-                    }}>
-                        <b>{"Host: "}</b>{callbackContext.host}
-                    </span>
-                </MythicStyledTooltip>
-            }
-            {callbackContext?.ip !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("ip") &&
-                <MythicStyledTooltip title={"First IP Address"}>
-                    <span className={"rounded-tab"} style={{
-                        backgroundColor: theme.taskContextColor,
-                        borderColor: callbackContext.color === "" ? "" : callbackContext.color
-                    }}>
-                        <b>{"IP: "}</b>{callbackContext.ip}
-                    </span>
-                </MythicStyledTooltip>
-            }
-            {callbackContext?.pid !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("pid") &&
-                <MythicStyledTooltip title={"Process ID"}>
-                    <span className={"rounded-tab"} style={{
-                        backgroundColor: theme.taskContextColor,
-                        borderColor: callbackContext.color === "" ? "" : callbackContext.color
-                    }}>
-                        <b>{"PID: "}</b>{callbackContext.pid}
-                    </span>
-                </MythicStyledTooltip>
-            }
-            {callbackContext?.architecture !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("architecture") &&
-                <MythicStyledTooltip title={"Process Architecture"}>
-                    <span className={"rounded-tab"} style={{
-                        backgroundColor: theme.taskContextColor,
-                        borderColor: callbackContext.color === "" ? "" : callbackContext.color
-                    }}>
-                        <b>{"Arch: "}</b>{callbackContext.architecture}
-                    </span>
-                </MythicStyledTooltip>
-            }
-            {callbackContext?.process_short_name !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("process_short_name") &&
-                <MythicStyledTooltip title={"Process Name"}>
-                    <span className={"rounded-tab"} style={{
-                        backgroundColor: theme.taskContextColor,
-                        borderColor: callbackContext.color === "" ? "" : callbackContext.color
-                    }}>
-                        <b>{"Process: "}</b>{callbackContext.process_short_name}
-                    </span>
-                </MythicStyledTooltip>
-            }
-            {callbackContext?.extra_info !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("extra_info") &&
-                <MythicStyledTooltip title={"Extra Callback Context"}>
-                    <span className={"rounded-tab"} style={{
-                        backgroundColor: theme.taskContextExtraColor,
-                        borderColor: callbackContext.color === "" ? "" : callbackContext.color
-                    }}>
-                        {callbackContext.extra_info}
-                    </span>
-                </MythicStyledTooltip>
-            }
-            <TextField
-                placeholder={"Task an agent..."}
-                onKeyDown={onKeyDown}
-                onChange={handleInputChange}
-                size="small"
-                color={"secondary"}
-                variant="outlined"
-                multiline={true}
-                maxRows={15}
-                disabled={reverseSearching}
-                value={message}
-                autoFocus={true}
-                fullWidth={true}
-                inputRef={inputRef}
-                style={{marginBottom: "0px", marginTop: "0px", paddingTop: "0px"}}
-                InputProps={{
-                    type: 'search',
-                    spellCheck: false,
-                    autoFocus: true,
-                    style: {paddingTop: "0px", paddingBottom: "0px", paddingRight: "5px"},
-                    endAdornment:
-                        <React.Fragment>
-                            <IconButton
-                                color="info"
-                                variant="contained"
-                                disableRipple={true}
-                                disableFocusRipple={true}
-                                onClick={onSubmitCommandLine}
-                                size="large"><SendIcon/>
-                            </IconButton>
-                            {props.filterTasks &&
-                                <IconButton
-                                    color={activeFiltering ? "warning" : "secondary"}
-                                    variant="contained"
-                                    onClick={onClickFilter}
-                                    style={{paddingLeft: 0}}
-                                    disableRipple={true}
-                                    disableFocusRipple={true}
-                                    size="large"><TuneIcon/></IconButton>
-                            }
-                            {commandPayloadType !== "" &&
-                                <MythicAgentSVGIcon payload_type={commandPayloadType}
-                                                    style={{width: "35px", height: "35px"}}/>
-                            }
-                        </React.Fragment>
-                    ,
-                    startAdornment: <React.Fragment>
-                        {tokenOptions.current.length > 0 ? (
-                            <CallbacksTabsTaskingInputTokenSelect options={tokenOptions.current}
-                                                                  changeSelectedToken={props.changeSelectedToken}/>
-                        ) : null}
-
-                    </React.Fragment>
-
-                }}
-            />
+            </div>
             {openFilterOptionsDialog &&
                 <MythicDialog fullWidth={true} maxWidth="md" open={openFilterOptionsDialog}
                               onClose={() => {

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

@@ -42,15 +42,16 @@ export function CallbacksTabsTaskingInputTokenSelect(props) {
       }
     }
   return (
-      <FormControl style={{width: props.width ? props.width : "20%"}}>
+      <FormControl className={props.className || ""} size={props.modern ? "small" : undefined} style={{width: props.width ? props.width : "20%"}}>
         <Select
           labelId="demo-dialog-select-label"
           id="demo-dialog-select"
           value={selected}
           onChange={handleChange}
-          variant="filled"
+          variant={props.modern ? "outlined" : "filled"}
+          size={props.modern ? "small" : undefined}
           renderValue={renderValue}
-          input={<Input style={{width: "100%"}}/>}
+          input={props.modern ? undefined : <Input style={{width: "100%"}}/>}
         >
           <MenuItem value={"Default Token"} key={0}>Default Token</MenuItem>
           {options.map( (opt) => (
@@ -60,4 +61,3 @@ export function CallbacksTabsTaskingInputTokenSelect(props) {
       </FormControl>
   );
 }
-

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

@@ -180,9 +180,16 @@ const NonInteractiveResponseDisplay = (props) => {
     fetchPolicy: "no-cache",
     onCompleted: (data) => {
       const responseArray = data.response.map(decodeResponse);
-      rawResponsesRef.current = responseArray;
-      setRawResponses(responseArray);
-      setTrackedTotalCount(getAggregateCount(data));
+      const shouldPreserveStreamedResponses = search === "" && currentPageRef.current === 1;
+      const nextResponses = shouldPreserveStreamedResponses ? mergeResponsesById({
+        existingResponses: rawResponsesRef.current,
+        incomingResponses: responseArray,
+        limit: pageSize,
+      }) : responseArray;
+      rawResponsesRef.current = nextResponses;
+      setRawResponses(nextResponses);
+      const aggregateCount = Math.max(getAggregateCount(data), shouldPreserveStreamedResponses ? nextResponses.length : 0);
+      setTrackedTotalCount(aggregateCount);
       highestKnownResponseID.current = Math.max(highestKnownResponseID.current, getLatestResponseID(data));
       pageDataLoaded.current = true;
       setOpenBackdrop(false);
@@ -195,15 +202,20 @@ const NonInteractiveResponseDisplay = (props) => {
     fetchPolicy: "no-cache",
     onCompleted: (data) => {
       const responseArray = data.response.map(decodeResponse);
-      rawResponsesRef.current = responseArray;
-      setRawResponses(responseArray);
-      const aggregateCount = getAggregateCount(data);
+      const nextResponses = search === "" ? mergeResponsesById({
+        existingResponses: rawResponsesRef.current,
+        incomingResponses: responseArray,
+        limit: MAX_SELECT_ALL_RESPONSES,
+      }) : responseArray;
+      rawResponsesRef.current = nextResponses;
+      setRawResponses(nextResponses);
+      const aggregateCount = Math.max(getAggregateCount(data), search === "" ? nextResponses.length : 0);
       setTrackedTotalCount(aggregateCount);
       highestKnownResponseID.current = Math.max(highestKnownResponseID.current, getLatestResponseID(data));
       pageDataLoaded.current = true;
       currentPageRef.current = 1;
       setCurrentPage(1);
-      warnSelectAllCapped(responseArray.length, aggregateCount);
+      warnSelectAllCapped(nextResponses.length, aggregateCount);
       setOpenBackdrop(false);
     },
     onError: (data) => {
@@ -262,17 +274,32 @@ const NonInteractiveResponseDisplay = (props) => {
     return () => clearTimeout(timeoutID);
   }, [openBackdrop]);
   React.useEffect( () => {
-    taskResponseCountRef.current = props.task.response_count || 0;
+    const nextResponseCount = props.task.response_count || 0;
+    const previousTaskResponseCount = taskResponseCountRef.current;
+    const previousKnownResponseCount = Math.max(totalCountRef.current, previousTaskResponseCount);
+    taskResponseCountRef.current = nextResponseCount;
     if(search !== ""){
       return;
     }
-    if(props.task.response_count > totalCountRef.current){
-      setTrackedTotalCount(props.task.response_count);
+    if(nextResponseCount > totalCountRef.current){
+      setTrackedTotalCount(nextResponseCount);
     }
-  }, [props.task.response_count, search, setTrackedTotalCount]);
+    if(nextResponseCount <= previousKnownResponseCount){
+      return;
+    }
+    const pageStartIndex = (currentPageRef.current - 1) * pageSize;
+    const pageEndIndex = currentPageRef.current * pageSize;
+    const currentPageCouldContainNewResponses = previousKnownResponseCount < pageEndIndex &&
+        nextResponseCount > pageStartIndex;
+    if(props.selectAllOutput){
+      fetchSelectAllResponses();
+    }else if(rawResponsesRef.current.length === 0 || currentPageCouldContainNewResponses){
+      fetchResponsePage(currentPageRef.current, false);
+    }
+  }, [fetchResponsePage, fetchSelectAllResponses, pageSize, props.selectAllOutput, props.task.response_count, search, setTrackedTotalCount]);
   const subscriptionDataCallback = React.useCallback( ({data}) => {
     const streamedResponses = data?.data?.response_stream || [];
-    if(streamedResponses.length === 0 || !pageDataLoaded.current){
+    if(streamedResponses.length === 0){
       return;
     }
     const previousHighestID = highestKnownResponseID.current;

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

@@ -869,11 +869,23 @@ export const ResponseDisplayInteractive = (props) =>{
         setSearch(newSearch || "");
     }, []);
     React.useEffect( () => {
-        taskResponseCountRef.current = props.task.response_count || 0;
+        const nextResponseCount = props.task.response_count || 0;
+        const previousKnownResponseCount = Math.max(totalCountRef.current, taskResponseCountRef.current);
+        taskResponseCountRef.current = Math.max(taskResponseCountRef.current, nextResponseCount);
         if(search === ""){
-            setTrackedTotalCount(props.task.response_count || 0);
+            setTrackedTotalCount(nextResponseCount);
         }
-    }, [props.task.response_count, search, setTrackedTotalCount]);
+        if(search !== "" || nextResponseCount <= previousKnownResponseCount){
+            return;
+        }
+        const pageStartIndex = (currentPageRef.current - 1) * pageSize.current;
+        const pageEndIndex = currentPageRef.current * pageSize.current;
+        const currentPageCouldContainNewResponses = previousKnownResponseCount < pageEndIndex &&
+            nextResponseCount > pageStartIndex;
+        if(props.selectAllOutput || rawResponsesRef.current.length === 0 || currentPageCouldContainNewResponses){
+            fetchResponsePage(currentPageRef.current, false);
+        }
+    }, [fetchResponsePage, props.selectAllOutput, props.task.response_count, search, setTrackedTotalCount]);
     const toggleANSIColor = () => {
         setUseANSIColor(!useASNIColor);
     }

MythicReactUI/src/themes/GlobalStyles.js

@@ -330,6 +330,383 @@ tspan {
     background-color: ${(props) => alpha(props.theme.palette.error.main, props.theme.palette.mode === "dark" ? 0.2 : 0.12)};
     color: ${(props) => props.theme.palette.error.main};
 }
+.mythic-tasking-composer {
+    background-color: ${(props) => alpha(props.theme.palette.background.paper, props.theme.palette.mode === "dark" ? 0.96 : 0.98)};
+    border-top: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    box-shadow: ${(props) => props.theme.palette.mode === "dark" ? "0 -10px 28px rgba(0, 0, 0, 0.24)" : "0 -10px 28px rgba(15, 23, 42, 0.06)"};
+    flex: 0 0 auto;
+    min-width: 0;
+    padding: 0.55rem 0.65rem 0.6rem;
+    position: relative;
+    width: 100%;
+}
+.mythic-tasking-context-row {
+    align-items: center;
+    display: flex;
+    gap: 0.28rem;
+    margin-bottom: 0.45rem;
+    min-height: 24px;
+    min-width: 0;
+    overflow-x: auto;
+    overflow-y: hidden;
+    padding-bottom: 0.05rem;
+}
+.mythic-tasking-context-chip {
+    align-items: center;
+    border: 1px solid ${(props) => alpha(props.theme.palette.common.black, props.theme.palette.mode === "dark" ? 0.22 : 0.1)};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    display: inline-flex;
+    flex: 0 0 auto;
+    font-size: 0.72rem;
+    font-weight: 750;
+    gap: 0.28rem;
+    line-height: 1;
+    max-width: min(28rem, 72vw);
+    min-height: 22px;
+    overflow: hidden;
+    padding: 0.22rem 0.45rem;
+    text-overflow: ellipsis;
+    white-space: nowrap;
+}
+.mythic-tasking-context-chip-emphasized {
+    box-shadow: inset 0 0 0 1px ${(props) => alpha(props.theme.palette.common.white, props.theme.palette.mode === "dark" ? 0.12 : 0.18)};
+}
+.mythic-tasking-context-chip-label {
+    flex: 0 0 auto;
+    font-size: 0.64rem;
+    font-weight: 850;
+    letter-spacing: 0;
+    opacity: 0.78;
+}
+.mythic-tasking-context-chip-value {
+    min-width: 0;
+    overflow: hidden;
+    text-overflow: ellipsis;
+}
+.mythic-tasking-command-row {
+    align-items: flex-end;
+    display: flex;
+    gap: 0.45rem;
+    min-width: 0;
+    width: 100%;
+}
+.mythic-tasking-token-select {
+    flex: 0 0 13rem;
+    max-width: 40%;
+    min-width: 9.5rem;
+}
+.mythic-tasking-token-select .MuiOutlinedInput-root {
+    background-color: ${(props) => props.theme.palette.background.paper};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    min-height: 37px;
+}
+.mythic-tasking-token-select .MuiSelect-select {
+    font-size: 0.78rem;
+    font-weight: 650;
+    min-height: 0 !important;
+    overflow: hidden;
+    padding-bottom: 0.52rem;
+    padding-top: 0.52rem;
+    text-overflow: ellipsis;
+    white-space: nowrap;
+}
+.mythic-tasking-command-input {
+    flex: 1 1 auto;
+    min-width: 10rem;
+}
+.mythic-tasking-command-input .MuiOutlinedInput-root {
+    align-items: flex-end;
+    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-radius: ${(props) => props.theme.shape.borderRadius}px;
+    padding: 0.1rem 0.24rem 0.1rem 0;
+}
+.mythic-tasking-command-input .MuiOutlinedInput-root:hover {
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? alpha(props.theme.palette.common.white, 0.05) : alpha(props.theme.palette.common.black, 0.026)};
+}
+.mythic-tasking-command-input .MuiInputBase-input {
+    font-family: ${(props) => props.theme.typography.fontFamily};
+    font-size: 0.86rem;
+    line-height: 1.45;
+    padding-bottom: 0.48rem;
+    padding-top: 0.48rem;
+}
+.mythic-tasking-command-prefix {
+    align-items: center;
+    color: ${(props) => props.theme.palette.text.secondary};
+    display: inline-flex;
+    flex: 0 0 auto;
+    height: 30px;
+    justify-content: center;
+    margin-left: 0.35rem;
+    margin-right: 0.35rem;
+    margin-top: 0.12rem;
+    opacity: 0.86;
+}
+.mythic-tasking-action-row {
+    align-items: center;
+    align-self: flex-end;
+    display: inline-flex;
+    flex: 0 0 auto;
+    gap: 0.18rem;
+    margin-bottom: 0.12rem;
+    padding-left: 0.28rem;
+}
+.mythic-tasking-action-button.MuiIconButton-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};
+    height: 28px;
+    padding: 0;
+    transition: background-color 120ms ease, border-color 120ms ease, color 120ms ease;
+    width: 28px;
+}
+.mythic-tasking-action-button-neutral.MuiIconButton-root:hover {
+    background-color: ${(props) => alpha(props.theme.palette.info.main, props.theme.palette.mode === "dark" ? 0.18 : 0.1)};
+    border-color: ${(props) => alpha(props.theme.palette.info.main, props.theme.palette.mode === "dark" ? 0.48 : 0.32)};
+    color: ${(props) => props.theme.palette.info.main};
+}
+.mythic-tasking-action-button-warning.MuiIconButton-root:hover {
+    background-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.2 : 0.12)};
+    border-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.52 : 0.36)};
+    color: ${(props) => props.theme.palette.warning.main};
+}
+.mythic-tasking-action-button-success.MuiIconButton-root:hover {
+    background-color: ${(props) => alpha(props.theme.palette.success.main, props.theme.palette.mode === "dark" ? 0.18 : 0.1)};
+    border-color: ${(props) => alpha(props.theme.palette.success.main, props.theme.palette.mode === "dark" ? 0.48 : 0.32)};
+    color: ${(props) => props.theme.palette.success.main};
+}
+.mythic-tasking-payload-chip {
+    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;
+    display: inline-flex;
+    flex: 0 0 auto;
+    height: 28px;
+    justify-content: center;
+    width: 28px;
+}
+.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)};
+    border: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    display: flex;
+    gap: 0.45rem;
+    margin-bottom: 0.45rem;
+    min-width: 0;
+    padding: 0.28rem;
+}
+.mythic-tasking-reverse-search-label {
+    color: ${(props) => props.theme.palette.text.secondary};
+    flex: 0 0 auto;
+    font-size: 0.72rem;
+    font-weight: 800;
+    padding-left: 0.3rem;
+}
+.mythic-tasking-reverse-search-input .MuiOutlinedInput-root {
+    background-color: ${(props) => props.theme.palette.background.paper};
+}
+.mythic-tasking-filter-dialog {
+    background-color: ${(props) => props.theme.surfaces?.raised || props.theme.palette.background.paper};
+    color: ${(props) => props.theme.palette.text.primary};
+}
+.mythic-tasking-filter-dialog-title.MuiDialogTitle-root {
+    background-image: ${getSectionHeaderGradient} !important;
+    border-bottom: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor} !important;
+    color: ${(props) => props.theme.pageHeaderText?.main || props.theme.palette.text.primary};
+    cursor: move;
+    padding: 1rem 1.15rem !important;
+    position: relative;
+}
+.mythic-tasking-filter-dialog-title.MuiDialogTitle-root::before {
+    background-color: ${getSectionHeaderAccent};
+    bottom: 0;
+    box-shadow: 0 0 0 1px ${(props) => alpha(props.theme.pageHeaderText?.main || props.theme.palette.text.primary, 0.2)};
+    content: "";
+    left: 0;
+    position: absolute;
+    top: 0;
+    width: 6px;
+}
+.mythic-tasking-filter-title-row {
+    align-items: center;
+    display: flex;
+    gap: 0.8rem;
+    min-width: 0;
+    min-height: 42px;
+    padding-left: 0.55rem;
+}
+.mythic-tasking-filter-title-icon {
+    align-items: center;
+    background-color: ${(props) => alpha(props.theme.pageHeaderText?.main || props.theme.palette.text.primary, 0.12)};
+    border: 1px solid ${(props) => alpha(props.theme.pageHeaderText?.main || props.theme.palette.text.primary, 0.24)};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    display: inline-flex;
+    flex: 0 0 auto;
+    height: 32px;
+    justify-content: center;
+    width: 32px;
+}
+.mythic-tasking-filter-title-main {
+    color: ${(props) => props.theme.pageHeaderText?.main || props.theme.palette.text.primary};
+    font-size: 1rem;
+    font-weight: 850;
+    line-height: 1.15;
+}
+.mythic-tasking-filter-title-subtitle {
+    color: ${(props) => alpha(props.theme.pageHeaderText?.main || props.theme.palette.text.primary, 0.78)};
+    font-size: 0.76rem;
+    font-weight: 600;
+    line-height: 1.25;
+    margin-top: 0.15rem;
+}
+.mythic-tasking-filter-dialog-actions.MuiDialogActions-root {
+    background-color: ${(props) => props.theme.surfaces?.muted || props.theme.palette.background.paper} !important;
+    border-top: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor} !important;
+    gap: 0.65rem !important;
+    padding: 0.85rem 1rem !important;
+}
+.mythic-tasking-filter-dialog-actions.MuiDialogActions-root .MuiButton-root {
+    min-height: 36px;
+    min-width: 8.25rem;
+    padding-left: 1rem;
+    padding-right: 1rem;
+}
+.mythic-tasking-filter-dialog-actions.MuiDialogActions-root .MuiButton-root:first-of-type {
+    margin-right: auto;
+}
+.mythic-tasking-filter-dialog-content.MuiDialogContent-root {
+    background-color: ${(props) => props.theme.surfaces?.raised || props.theme.palette.background.paper};
+    max-height: min(72vh, 48rem);
+    min-width: min(46rem, calc(100vw - 3rem));
+    overflow: auto;
+    padding: 1rem !important;
+}
+.mythic-tasking-filter-summary {
+    align-items: center;
+    display: flex;
+    flex-wrap: wrap;
+    gap: 0.35rem;
+    min-height: 28px;
+}
+.mythic-tasking-filter-summary-chip.MuiChip-root {
+    background-color: ${(props) => alpha(props.theme.palette.primary.main, props.theme.palette.mode === "dark" ? 0.18 : 0.1)};
+    border: 1px solid ${(props) => alpha(props.theme.palette.primary.main, props.theme.palette.mode === "dark" ? 0.42 : 0.28)};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    color: ${(props) => props.theme.palette.primary.main};
+    font-size: 0.72rem;
+    font-weight: 750;
+}
+.mythic-tasking-filter-summary-chip .MuiChip-icon {
+    color: inherit;
+    font-size: 1rem;
+}
+.mythic-tasking-filter-summary-chip-muted.MuiChip-root {
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? "rgba(255,255,255,0.045)" : "rgba(0,0,0,0.025)"};
+    border-color: ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    color: ${(props) => props.theme.palette.text.secondary};
+}
+.mythic-tasking-filter-switch-stack {
+    display: flex;
+    flex-direction: column;
+    gap: 0.55rem;
+    min-width: 0;
+}
+.mythic-tasking-filter-select .MuiOutlinedInput-root {
+    background-color: ${(props) => props.theme.palette.background.paper};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    min-height: 38px;
+}
+.mythic-tasking-filter-select .MuiSelect-select {
+    align-items: center;
+    display: flex;
+    min-height: 0 !important;
+    padding-bottom: 0.42rem;
+    padding-top: 0.42rem;
+}
+.mythic-tasking-filter-select-empty {
+    color: ${(props) => props.theme.palette.text.secondary};
+    font-size: 0.78rem;
+    font-weight: 650;
+}
+.mythic-tasking-filter-select-chips {
+    align-items: center;
+    display: flex;
+    flex-wrap: wrap;
+    gap: 0.25rem;
+    min-width: 0;
+}
+.mythic-tasking-filter-selected-chip.MuiChip-root {
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? "rgba(255,255,255,0.07)" : "rgba(0,0,0,0.045)"};
+    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.primary};
+    font-size: 0.72rem;
+    font-weight: 700;
+}
+.mythic-tasking-filter-command-grid {
+    align-items: start;
+    display: grid;
+    gap: 0.6rem;
+    grid-template-columns: minmax(0, 1fr) auto minmax(0, 1fr);
+    min-width: 0;
+    width: 100%;
+}
+.mythic-tasking-filter-choice-divider {
+    align-self: center;
+    color: ${(props) => props.theme.palette.text.secondary};
+    font-size: 0.72rem;
+    font-weight: 750;
+    padding-top: 1.55rem;
+    text-align: center;
+}
+.mythic-tasking-filter-clear-button.MuiButton-root {
+    align-self: flex-start;
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? "rgba(255,255,255,0.05)" : "rgba(0,0,0,0.025)"} !important;
+    border: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor} !important;
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    box-shadow: none;
+    color: ${(props) => props.theme.palette.text.secondary} !important;
+    font-size: 0.72rem;
+    font-weight: 750;
+    margin-top: 0.4rem;
+    min-height: 28px;
+    text-transform: none;
+}
+.mythic-tasking-filter-clear-button.MuiButton-root:hover {
+    background-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.18 : 0.1)} !important;
+    border-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.48 : 0.32)} !important;
+    color: ${(props) => props.theme.palette.warning.main} !important;
+}
+@media screen and (max-width: 760px) {
+    .mythic-tasking-filter-dialog-content.MuiDialogContent-root {
+        min-width: 0;
+    }
+    .mythic-tasking-filter-command-grid {
+        grid-template-columns: 1fr;
+    }
+    .mythic-tasking-filter-choice-divider {
+        padding-top: 0;
+    }
+    .mythic-tasking-filter-dialog-actions.MuiDialogActions-root .MuiButton-root,
+    .mythic-tasking-filter-dialog-actions.MuiDialogActions-root .MuiButton-root:first-of-type {
+        flex: 1 1 100%;
+        margin-right: 0;
+    }
+}
+@media screen and (max-width: 760px) {
+    .mythic-tasking-command-row {
+        align-items: stretch;
+        flex-direction: column;
+    }
+    .mythic-tasking-token-select {
+        flex: 0 0 auto;
+        max-width: 100%;
+        width: 100% !important;
+    }
+}
 .mythic-file-browser-iconButtonCompound.MuiIconButton-root {
     gap: 0;
     padding: 0 0.15rem;

Mythic_CLI/src/cmd/config/env.go

@@ -82,7 +82,7 @@ func GetIntendedMythicServiceNames() ([]string, error) {
 					if mythicEnv.GetBool("postgres_debug") {
 						containerList = append(containerList, service)
 					}
-			
+
 			*/
 		}
 	}
@@ -321,6 +321,9 @@ Setting this to "true" means that the local Mythic/rabbitmq-docker/Dockerfile is
 	mythicEnv.SetDefault("rabbitmq_vhost", "mythic_vhost")
 	mythicEnvInfo["rabbitmq_vhost"] = `The VHost attribute on RabbitMQ allows you to logically separate queues into separate "hosts" while keeping the same names. This helps with collisions if you have multiple instances of something running concurrently.`
 
+	mythicEnv.SetDefault("custom_rpc_timeout", 0)
+	mythicEnvInfo["custom_rpc_timeout"] = `This is an optional timeout in seconds for Mythic RPC calls that are expected to be slow but still valid. If set to 0, Mythic uses the default RPC timeout.`
+
 	// jwt configuration ---------------------------------------------
 	mythicEnv.SetDefault("jwt_secret", utils.GenerateRandomPassword(30))
 	mythicEnvInfo["jwt_secret"] = `This is the randomly generated password used to sign JWTs to ensure they're valid for this Mythic instance`

Mythic_CLI/src/cmd/internal/serviceMetadata.go

@@ -692,6 +692,7 @@ func AddMythicService(service string, removeVolume bool) {
 			"DEFAULT_OPERATION_NAME=${DEFAULT_OPERATION_NAME}",
 			"DEFAULT_OPERATION_WEBHOOK_CHANNEL=${DEFAULT_OPERATION_WEBHOOK_CHANNEL}",
 			"DEFAULT_OPERATION_WEBHOOK_URL=${DEFAULT_OPERATION_WEBHOOK_URL}",
+			"CUSTOM_RPC_TIMEOUT=${CUSTOM_RPC_TIMEOUT}",
 			"GLOBAL_SERVER_NAME=${GLOBAL_SERVER_NAME}",
 			"JWT_SECRET=${JWT_SECRET}",
 			"MYTHIC_ADMIN_USER=${MYTHIC_ADMIN_USER}",

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

@@ -28,6 +28,7 @@ set
         excluded.last_callback_display_id
     );
 
+-- +migrate StatementBegin
 do $$
 begin
     if exists (
@@ -48,6 +49,7 @@ begin
         raise exception 'Cannot add callback(operation_id, display_id) uniqueness because duplicate callback display IDs already exist';
     end if;
 end $$;
+-- +migrate StatementEnd
 
 create unique index if not exists task_operation_display_id_unique
 on "public"."task" using btree (operation_id, display_id);
@@ -55,6 +57,7 @@ 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
 create or replace function public.new_task_display_id() returns trigger
     language plpgsql
     as $$
@@ -71,7 +74,9 @@ begin
     return new;
 end;
 $$;
+-- +migrate StatementEnd
 
+-- +migrate StatementBegin
 create or replace function public.new_callback_display_id() returns trigger
     language plpgsql
     as $$
@@ -88,6 +93,7 @@ begin
     return new;
 end;
 $$;
+-- +migrate StatementEnd
 
 update "public"."task"
 set response_count = response_counts.total
@@ -102,6 +108,7 @@ from (
 where task.id = response_counts.id
 and task.response_count is distinct from response_counts.total;
 
+-- +migrate StatementBegin
 create or replace function public.update_task_response_count() returns trigger
     language plpgsql
     as $$
@@ -113,10 +120,12 @@ begin
     return new;
 end;
 $$;
+-- +migrate StatementEnd
 
 -- +migrate Down
 -- SQL in section 'Down' is executed when this migration is rolled back
 
+-- +migrate StatementBegin
 create or replace function public.update_task_response_count() returns trigger
     language plpgsql
     as $$
@@ -132,7 +141,9 @@ begin
     return new;
 end;
 $$;
+-- +migrate StatementEnd
 
+-- +migrate StatementBegin
 create or replace function public.new_callback_display_id() returns trigger
     language plpgsql
     as $$
@@ -148,7 +159,9 @@ begin
     return new;
 end;
 $$;
+-- +migrate StatementEnd
 
+-- +migrate StatementBegin
 create or replace function public.new_task_display_id() returns trigger
     language plpgsql
     as $$
@@ -164,6 +177,7 @@ begin
     return new;
 end;
 $$;
+-- +migrate StatementEnd
 
 drop index if exists "public"."callback_operation_display_id_unique";
 drop index if exists "public"."task_operation_display_id_unique";

mythic-docker/src/rabbitmq/constants.go

@@ -2,6 +2,8 @@ package rabbitmq
 
 import "time"
 
+type RPCRetryPolicy int
+
 const (
 	MYTHIC_EXCHANGE                        = "mythic_exchange"
 	MYTHIC_TOPIC_EXCHANGE                  = "mythic_topic_exchange"
@@ -13,6 +15,15 @@ const (
 	TASK_STATUS_CONTAINER_DOWN             = "Error: Container Down"
 )
 
+const (
+	// Retry response timeouts with a new RPC request, preserving Mythic's historical 3-attempt behavior.
+	RPC_RETRY_POLICY_RETRY_ON_TIMEOUT RPCRetryPolicy = iota
+	// Retry publish failures only; once a request is delivered, do not send duplicate work after response timeout.
+	RPC_RETRY_POLICY_NO_RETRY_ON_TIMEOUT
+	// Use CUSTOM_RPC_TIMEOUT for slow valid calls and do not send duplicate work after response timeout.
+	RPC_RETRY_POLICY_CUSTOM_TIMEOUT
+)
+
 // Direct fanout rabbitmq routes where Mythic is consuming messages, but others can also listen in and consume
 const (
 	// PT_SYNC_ROUTING_KEY payload routes

mythic-docker/src/rabbitmq/initialize.go

@@ -44,6 +44,16 @@ type rabbitMQConnection struct {
 	addListenerMutex sync.RWMutex
 	channelMutexMap  map[string]*channelMutex
 	channelMutex     sync.RWMutex
+	publisherMutex   sync.Mutex
+	publisherChannel *amqp.Channel
+	publisherConfirm chan amqp.Confirmation
+	publisherReturn  chan amqp.Return
+	rpcClientMutex   sync.Mutex
+	rpcChannel       *amqp.Channel
+	rpcConfirm       chan amqp.Confirmation
+	rpcReturn        chan amqp.Return
+	rpcPending       map[string]chan rpcResponse
+	rpcExchanges     map[string]bool
 	RPCQueues        []RPCQueueStruct
 	DirectQueues     []DirectQueueStruct
 }

mythic-docker/src/rabbitmq/send_auth_get_idp_metadata.go

@@ -30,6 +30,7 @@ func (r *rabbitMQConnection) SendAuthGetIDPMetadata(input GetIDPMetadataMessage)
 		GetAuthContainerGetIDPMetadataRoutingKey(input.ContainerName),
 		inputBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_RETRY_ON_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_auth_get_idp_redirect.go

@@ -36,6 +36,7 @@ func (r *rabbitMQConnection) SendAuthGetIDPRedirect(input GetIDPRedirectMessage)
 		GetAuthContainerGetIDPRedirectRoutingKey(input.ContainerName),
 		inputBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_RETRY_ON_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_auth_get_nonidp_metadata.go

@@ -30,6 +30,7 @@ func (r *rabbitMQConnection) SendAuthGetNonIDPMetadata(input GetNonIDPMetadataMe
 		GetAuthContainerGetNonIDPMetadataRoutingKey(input.ContainerName),
 		inputBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_RETRY_ON_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_auth_get_nonidp_redirect.go

@@ -30,6 +30,7 @@ func (r *rabbitMQConnection) SendAuthGetNonIDPRedirect(input GetNonIDPRedirectMe
 		GetAuthContainerGetNonIDPRedirectRoutingKey(input.ContainerName),
 		inputBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_RETRY_ON_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_auth_process_idp_response.go

@@ -35,6 +35,7 @@ func (r *rabbitMQConnection) SendAuthProcessIDPResponse(input ProcessIDPResponse
 		GetAuthContainerProcessIDPResponseRoutingKey(input.ContainerName),
 		inputBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_NO_RETRY_ON_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_auth_process_nonidp_response.go

@@ -31,6 +31,7 @@ func (r *rabbitMQConnection) SendAuthProcessNonIDPResponse(input ProcessNonIDPRe
 		GetAuthContainerProcessNonIDPResponseRoutingKey(input.ContainerName),
 		inputBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_NO_RETRY_ON_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_c2_rpc_config_check.go

@@ -32,6 +32,7 @@ func (r *rabbitMQConnection) SendC2RPCConfigCheck(configCheck C2ConfigCheckMessa
 		GetC2RPCConfigChecksRoutingKey(configCheck.Name),
 		configBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_c2_rpc_get_debug_output.go

@@ -32,6 +32,7 @@ func (r *rabbitMQConnection) SendC2RPCGetDebugOutput(getDebugOutput C2GetDebugOu
 		GetC2RPCGetServerDebugOutputRoutingKey(getDebugOutput.Name),
 		opsecBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	); err != nil {
 		logging.LogError(err, "Failed to send RPC message")
 		return &c2GetDebutOutputResponse, err

mythic-docker/src/rabbitmq/send_c2_rpc_get_ioc.go

@@ -32,6 +32,7 @@ func (r *rabbitMQConnection) SendC2RPCGetIOC(getIOC C2GetIOCMessage) (*C2GetIOCM
 		GetC2RPCGetIOCRoutingKey(getIOC.Name),
 		configBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	); err != nil {
 		logging.LogError(err, "Failed to send RPC message")
 		return nil, err

mythic-docker/src/rabbitmq/send_c2_rpc_get_redirector_rules.go

@@ -32,6 +32,7 @@ func (r *rabbitMQConnection) SendC2RPCGetRedirectorRules(redirectorRules C2GetRe
 		GetC2RPCRedirectorRulesRoutingKey(redirectorRules.Name),
 		opsecBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	); err != nil {
 		logging.LogError(err, "Failed to send RPC message")
 		return nil, err

mythic-docker/src/rabbitmq/send_c2_rpc_host_file.go

@@ -35,6 +35,7 @@ func (r *rabbitMQConnection) SendC2RPCHostFile(hostFile C2HostFileMessage) (*C2H
 		GetC2RPCHostFileRoutingKey(hostFile.Name),
 		opsecBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_c2_rpc_opsec_check.go

@@ -35,6 +35,7 @@ func (r *rabbitMQConnection) SendC2RPCOpsecCheck(opsecCheck C2OPSECMessage) (*C2
 		GetC2RPCOpsecChecksRoutingKey(opsecCheck.Name),
 		opsecBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_c2_rpc_resync.go

@@ -28,6 +28,7 @@ func (r *rabbitMQConnection) SendC2RPCReSync(resync C2RPCReSyncMessage) (*C2RPCR
 		GetC2RPCReSyncRoutingKey(resync.Name),
 		configBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	); err != nil {
 		logging.LogError(err, "Failed to send RPC message")
 		return nil, err

mythic-docker/src/rabbitmq/send_c2_rpc_sample_message.go

@@ -29,6 +29,7 @@ func (r *rabbitMQConnection) SendC2RPCSampleMessage(getSampleMessage C2SampleMes
 		GetC2RPCSampleMessageRoutingKey(getSampleMessage.Name),
 		configBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	); err != nil {
 		logging.LogError(err, "Failed to send RPC message")
 		return nil, err

mythic-docker/src/rabbitmq/send_c2_rpc_start_server.go

@@ -92,6 +92,7 @@ func (r *rabbitMQConnection) SendC2RPCStartServer(startServer C2StartServerMessa
 		GetC2RPCStartServerRoutingKey(startServer.Name),
 		opsecBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_NO_RETRY_ON_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_c2_rpc_stop_server.go

@@ -34,6 +34,7 @@ func (r *rabbitMQConnection) SendC2RPCStopServer(stopServer C2StopServerMessage)
 		GetC2RPCStopServerRoutingKey(stopServer.Name),
 		opsecBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_NO_RETRY_ON_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_container_rpc_get_file.go

@@ -28,6 +28,7 @@ func (r *rabbitMQConnection) SendContainerRPCGetFile(getFile ContainerGetFileMes
 		GetC2RPCGetFileRoutingKey(getFile.Name),
 		opsecBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	); err != nil {
 		logging.LogError(err, "Failed to send RPC message")
 		return nil, err

mythic-docker/src/rabbitmq/send_container_rpc_list_file.go

@@ -27,6 +27,7 @@ func (r *rabbitMQConnection) SendContainerRPCListFile(getFile ContainerRPCListFi
 		GetC2RPCListFileRoutingKey(getFile.Name),
 		opsecBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	); err != nil {
 		logging.LogError(err, "Failed to send RPC message")
 		return nil, err

mythic-docker/src/rabbitmq/send_container_rpc_remove_file.go

@@ -27,6 +27,7 @@ func (r *rabbitMQConnection) SendContainerRPCRemoveFile(getFile ContainerRemoveF
 		GetC2RPCRemoveFileRoutingKey(getFile.Name),
 		opsecBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	); err != nil {
 		logging.LogError(err, "Failed to send RPC message")
 		return nil, err

mythic-docker/src/rabbitmq/send_container_rpc_write_file.go

@@ -29,6 +29,7 @@ func (r *rabbitMQConnection) SendContainerRPCWriteFile(writeFile ContainerWriteF
 		GetC2RPCWriteFileRoutingKey(writeFile.Name),
 		opsecBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	); err != nil {
 		logging.LogError(err, "Failed to send RPC message")
 		return nil, err

mythic-docker/src/rabbitmq/send_pt_command_help.go

@@ -29,6 +29,7 @@ func (r *rabbitMQConnection) SendPtRPCCommandHelp(msg CommandHelpMessage) (*Comm
 		GetPtCommandHelpRoutingKey(msg.PayloadType),
 		configBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_RETRY_ON_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_pt_rpc_check_if_callbacks_alive.go

@@ -126,6 +126,7 @@ func (r *rabbitMQConnection) SendPTRPCCheckIfCallbacksAlive(sendMsg PTCheckIfCal
 		GetPtCheckIfCallbacksAliveRoutingKey(sendMsg.ContainerName),
 		opsecBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_RETRY_ON_TIMEOUT,
 	)
 	if err != nil {
 		return nil, err

mythic-docker/src/rabbitmq/send_pt_rpc_dynamic_query_build_parameter_function.go

@@ -40,6 +40,7 @@ func (r *rabbitMQConnection) SendPtRPCDynamicQueryBuildParameterFunction(dynamic
 		GetPtRPCDynamicQueryBuildParameterFunctionRoutingKey(dynamicQuery.PayloadType),
 		configBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_pt_rpc_dynamic_query_function.go

@@ -74,6 +74,7 @@ func (r *rabbitMQConnection) SendPtRPCDynamicQueryFunction(dynamicQuery PTRPCDyn
 		GetPtRPCDynamicQueryFunctionRoutingKey(dynamicQuery.CommandPayloadType),
 		configBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_pt_rpc_resync.go

@@ -30,6 +30,7 @@ func (r *rabbitMQConnection) SendPTRPCReSync(resync PTRPCReSyncMessage) (*PTRPCR
 		GetPTRPCReSyncRoutingKey(resync.Name),
 		configBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	)
 	if err != nil {
 		logging.LogError(err, "Failed to send RPC message")

mythic-docker/src/rabbitmq/send_pt_rpc_typedArray_parse.go

@@ -41,6 +41,7 @@ func (r *rabbitMQConnection) SendPtRPCTypedArrayParse(dynamicQuery PTRPCTypedArr
 		GetPtRPCTypedArrayParseRoutingKey(dynamicQuery.CommandPayloadType),
 		configBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	); err != nil {
 		logging.LogError(err, "Failed to send RPC message")
 		return nil, err

mythic-docker/src/rabbitmq/send_tr_rpc_decrypt_bytes.go

@@ -33,6 +33,7 @@ func (r *rabbitMQConnection) SendTrRPCDecryptBytes(input TrDecryptBytesMessage)
 		GetTrRPCDecryptBytesRoutingKey(input.TranslationContainerName),
 		opsecBytes,
 		!exclusiveQueue,
+		RPC_RETRY_POLICY_RETRY_ON_TIMEOUT,
 	); err != nil {
 		logging.LogError(err, "Failed to send RPC message")
 		return nil, err

mythic-docker/src/rabbitmq/send_tr_rpc_encrypt_bytes.go

@@ -34,6 +34,7 @@ func (r *rabbitMQConnection) SendTrRPCEncryptBytes(input TrEncryptBytesMessage)
 		GetTrRPCEncryptBytesRoutingKey(input.TranslationContainerName),
 		opsecBytes,
 		!exclusiveQueue,
+		RPC_RETRY_POLICY_RETRY_ON_TIMEOUT,
 	); err != nil {
 		logging.LogError(err, "Failed to send RPC message")
 		return nil, err

mythic-docker/src/rabbitmq/send_tr_rpc_resync.go

@@ -28,6 +28,7 @@ func (r *rabbitMQConnection) SendTRRPCReSync(resync TRRPCReSyncMessage) (*TRRPCR
 		GetTrRPCReSyncRoutingKey(resync.Name),
 		configBytes,
 		exclusiveQueue,
+		RPC_RETRY_POLICY_CUSTOM_TIMEOUT,
 	); err != nil {
 		logging.LogError(err, "Failed to send RPC message")
 		return nil, err

mythic-docker/src/rabbitmq/utils_rabbitmq_routing.go

@@ -14,6 +14,11 @@ import (
 	amqp "github.com/rabbitmq/amqp091-go"
 )
 
+type rpcResponse struct {
+	Body []byte
+	Err  error
+}
+
 // payload routing key functions
 func GetPtBuildRoutingKey(container string) string {
 	return fmt.Sprintf("%s_%s", container, PT_BUILD_ROUTING_KEY)
@@ -191,41 +196,62 @@ func (r *rabbitMQConnection) SendStructMessage(exchange string, queue string, co
 	}
 	return r.SendMessage(exchange, queue, correlationId, jsonBody, ignoreErrorMessage)
 }
-func (r *rabbitMQConnection) SendRPCStructMessage(exchange string, queue string, body interface{}) ([]byte, error) {
+func (r *rabbitMQConnection) SendRPCStructMessage(exchange string, queue string, body interface{}, retryPolicy RPCRetryPolicy) ([]byte, error) {
 	if inputBytes, err := json.Marshal(body); err != nil {
 		logging.LogError(err, "Failed to convert input to JSON", "input", body)
 		return nil, err
 	} else {
-		return r.SendRPCMessage(exchange, queue, inputBytes, true)
+		return r.SendRPCMessage(exchange, queue, inputBytes, true, retryPolicy)
 	}
 }
+
+func (r *rabbitMQConnection) getPublisherChannel() (*amqp.Channel, chan amqp.Confirmation, chan amqp.Return, error) {
+	if r.publisherChannel != nil && !r.publisherChannel.IsClosed() {
+		return r.publisherChannel, r.publisherConfirm, r.publisherReturn, nil
+	}
+	conn, err := r.GetConnection()
+	if err != nil {
+		return nil, nil, nil, err
+	}
+	ch, err := conn.Channel()
+	if err != nil {
+		return nil, nil, nil, err
+	}
+	if err = ch.Confirm(false); err != nil {
+		ch.Close()
+		return nil, nil, nil, err
+	}
+	r.publisherChannel = ch
+	r.publisherConfirm = ch.NotifyPublish(make(chan amqp.Confirmation, 1))
+	r.publisherReturn = ch.NotifyReturn(make(chan amqp.Return, 1))
+	return r.publisherChannel, r.publisherConfirm, r.publisherReturn, nil
+}
+
+func (r *rabbitMQConnection) resetPublisherChannel(ch *amqp.Channel) {
+	if ch != nil && !ch.IsClosed() {
+		ch.Close()
+	}
+	if r.publisherChannel == ch {
+		r.publisherChannel = nil
+		r.publisherConfirm = nil
+		r.publisherReturn = nil
+	}
+}
+
 func (r *rabbitMQConnection) SendMessage(exchange string, queue string, correlationId string, body []byte, ignoreErrormessage bool) error {
 	// to send a normal message out to a direct queue set:
 	// exchange: MYTHIC_EXCHANGE
 	// queue: which routing key is listening (this is the direct name)
 	// correlation_id: empty string
 	for attempt := 0; attempt < 3; attempt++ {
-		conn, err := r.GetConnection()
+		r.publisherMutex.Lock()
+		ch, confirmChannel, notifyReturnChannel, err := r.getPublisherChannel()
 		if err != nil {
-			logging.LogError(err, "Failed to get rabbitmq connection", "queue", queue)
+			logging.LogError(err, "Failed to get rabbitmq publisher channel", "queue", queue)
+			r.publisherMutex.Unlock()
 			time.Sleep(RPC_TIMEOUT)
 			continue
 		}
-		ch, err := conn.Channel()
-		if err != nil {
-			logging.LogError(err, "Failed to open rabbitmq channel", "queue", queue)
-			time.Sleep(RPC_TIMEOUT)
-			continue
-		}
-		err = ch.Confirm(false)
-		if err != nil {
-			logging.LogError(err, "Channel could not be put into confirm mode", "queue", queue)
-			ch.Close()
-			time.Sleep(RPC_TIMEOUT)
-			continue
-		}
-		confirmChannel := ch.NotifyPublish(make(chan amqp.Confirmation, 1))
-		notifyReturnChannel := ch.NotifyReturn(make(chan amqp.Return, 1))
 		msg := amqp.Publishing{
 			ContentType:   "application/json",
 			CorrelationId: correlationId,
@@ -240,7 +266,8 @@ func (r *rabbitMQConnection) SendMessage(exchange string, queue string, correlat
 		)
 		if err != nil {
 			logging.LogError(err, "there was an error publishing a message", "queue", queue)
-			ch.Close()
+			r.resetPublisherChannel(ch)
+			r.publisherMutex.Unlock()
 			time.Sleep(RPC_TIMEOUT)
 			continue
 		}
@@ -249,7 +276,8 @@ func (r *rabbitMQConnection) SendMessage(exchange string, queue string, correlat
 			if !ntf.Ack {
 				err = errors.New("Failed to deliver message, not ACK-ed by receiver")
 				logging.LogError(err, "failed to deliver message to exchange/queue, notifyPublish")
-				ch.Close()
+				r.resetPublisherChannel(ch)
+				r.publisherMutex.Unlock()
 				time.Sleep(RPC_TIMEOUT)
 				continue
 			}
@@ -258,16 +286,18 @@ func (r *rabbitMQConnection) SendMessage(exchange string, queue string, correlat
 			if !ignoreErrormessage {
 				logging.LogError(err, "failed to deliver message to exchange/queue, NotifyReturn", "errorCode", ret.ReplyCode, "errorText", ret.ReplyText)
 			}
-			ch.Close()
+			r.resetPublisherChannel(ch)
+			r.publisherMutex.Unlock()
 			time.Sleep(RPC_TIMEOUT)
 			continue
 		case <-time.After(RPC_TIMEOUT):
 			err = errors.New("Message delivery confirmation timed out")
 			logging.LogError(err, "message delivery confirmation to exchange/queue timed out")
-			ch.Close()
+			r.resetPublisherChannel(ch)
+			r.publisherMutex.Unlock()
 			continue
 		}
-		ch.Close()
+		r.publisherMutex.Unlock()
 		return nil
 	}
 	if !ignoreErrormessage {
@@ -275,108 +305,221 @@ func (r *rabbitMQConnection) SendMessage(exchange string, queue string, correlat
 	}
 	return errors.New(fmt.Sprintf("failed 3 times to send to queue %s", queue))
 }
-func (r *rabbitMQConnection) SendRPCMessage(exchange string, queue string, body []byte, exclusiveQueue bool) ([]byte, error) {
-	var finalError error
-	for attempt := 0; attempt < 3; attempt++ {
-		conn, err := r.GetConnection()
-		if err != nil {
-			return nil, err
+func (r *rabbitMQConnection) getRPCTimeout(retryPolicy RPCRetryPolicy) time.Duration {
+	if retryPolicy == RPC_RETRY_POLICY_CUSTOM_TIMEOUT && utils.MythicConfig.CustomRPCTimeout > 0 {
+		return utils.MythicConfig.CustomRPCTimeout
+	}
+	return RPC_TIMEOUT
+}
+
+func (r *rabbitMQConnection) getRPCClientLocked(exchange string, exclusiveQueue bool) (*amqp.Channel, chan amqp.Confirmation, chan amqp.Return, error) {
+	if r.rpcChannel != nil && !r.rpcChannel.IsClosed() {
+		if err := r.declareRPCExchangeLocked(r.rpcChannel, exchange); err != nil {
+			r.resetRPCClientLocked(r.rpcChannel, err)
+			return nil, nil, nil, err
 		}
-		ch, err := conn.Channel()
-		if err != nil {
-			logging.LogError(err, "Failed to open rabbitmq channel")
-			return nil, err
+		return r.rpcChannel, r.rpcConfirm, r.rpcReturn, nil
+	}
+	conn, err := r.GetConnection()
+	if err != nil {
+		return nil, nil, nil, err
+	}
+	ch, err := conn.Channel()
+	if err != nil {
+		return nil, nil, nil, err
+	}
+	if err = ch.Confirm(false); err != nil {
+		ch.Close()
+		return nil, nil, nil, err
+	}
+	confirmChannel := ch.NotifyPublish(make(chan amqp.Confirmation, 1))
+	notifyReturnChannel := ch.NotifyReturn(make(chan amqp.Return, 1))
+	msgs, err := ch.Consume(
+		"amq.rabbitmq.reply-to", // queue name
+		"",                      // consumer
+		true,                    // auto-ack
+		exclusiveQueue,          // exclusive
+		false,                   // no local
+		false,                   // no wait
+		nil,                     // args
+	)
+	if err != nil {
+		ch.Close()
+		return nil, nil, nil, err
+	}
+	r.rpcChannel = ch
+	r.rpcConfirm = confirmChannel
+	r.rpcReturn = notifyReturnChannel
+	if r.rpcPending == nil {
+		r.rpcPending = make(map[string]chan rpcResponse)
+	}
+	r.rpcExchanges = make(map[string]bool)
+	if err = r.declareRPCExchangeLocked(ch, exchange); err != nil {
+		r.resetRPCClientLocked(ch, err)
+		return nil, nil, nil, err
+	}
+	go r.listenForRPCReplies(ch, msgs)
+	return r.rpcChannel, r.rpcConfirm, r.rpcReturn, nil
+}
+
+func (r *rabbitMQConnection) declareRPCExchangeLocked(ch *amqp.Channel, exchange string) error {
+	if r.rpcExchanges == nil {
+		r.rpcExchanges = make(map[string]bool)
+	}
+	if r.rpcExchanges[exchange] {
+		return nil
+	}
+	err := ch.ExchangeDeclare(
+		exchange, // exchange name
+		"direct", // type of exchange, ex: topic, fanout, direct, etc
+		true,     // durable
+		true,     // auto-deleted
+		false,    // internal
+		false,    // no-wait
+		nil,      // arguments
+	)
+	if err == nil {
+		r.rpcExchanges[exchange] = true
+	}
+	return err
+}
+
+func (r *rabbitMQConnection) listenForRPCReplies(ch *amqp.Channel, msgs <-chan amqp.Delivery) {
+	for d := range msgs {
+		r.rpcClientMutex.Lock()
+		responseChannel := r.rpcPending[d.CorrelationId]
+		if responseChannel != nil {
+			delete(r.rpcPending, d.CorrelationId)
 		}
-		err = ch.Confirm(false)
-		if err != nil {
-			logging.LogError(err, "Channel could not be put into confirm mode")
-			ch.Close()
-			return nil, err
+		r.rpcClientMutex.Unlock()
+		if responseChannel != nil {
+			responseChannel <- rpcResponse{Body: d.Body}
 		}
-		err = ch.ExchangeDeclare(
-			exchange, // exchange name
-			"direct", // type of exchange, ex: topic, fanout, direct, etc
-			true,     // durable
-			true,     // auto-deleted
-			false,    // internal
-			false,    // no-wait
-			nil,      // arguments
-		)
-		if err != nil {
-			logging.LogError(err, "Failed to declare exchange", "exchange", exchange, "exchange_type", "direct", "retry_wait_time", RETRY_CONNECT_DELAY)
-			ch.Close()
-			return nil, err
-		}
-		msgs, err := ch.Consume(
-			"amq.rabbitmq.reply-to", // queue name
-			"",                      // consumer
-			true,                    // auto-ack
-			exclusiveQueue,          // exclusive
-			false,                   // no local
-			false,                   // no wait
-			nil,                     // args
-		)
-		if err != nil {
-			logging.LogError(err, "Failed to start consuming for RPC replies")
-			ch.Close()
-			return nil, err
-		}
-		confirmChannel := ch.NotifyPublish(make(chan amqp.Confirmation, 1))
-		notifyReturnChannel := ch.NotifyReturn(make(chan amqp.Return, 1))
-		msg := amqp.Publishing{
-			ContentType:   "application/json",
-			CorrelationId: uuid.NewString(),
-			Body:          body,
-			ReplyTo:       "amq.rabbitmq.reply-to",
-		}
-		err = ch.Publish(
-			exchange, // exchange
-			queue,    // routing key
-			true,     // mandatory
-			false,    // immediate
-			msg,      // publishing
-		)
-		finalError = err
-		if err != nil {
-			logging.LogError(err, "there was an error publishing an rpc message", "queue", queue)
-			ch.Close()
-			time.Sleep(RPC_TIMEOUT)
-			continue
+	}
+	r.rpcClientMutex.Lock()
+	if r.rpcChannel == ch {
+		r.resetRPCClientLocked(ch, errors.New("rpc reply consumer stopped"))
+	}
+	r.rpcClientMutex.Unlock()
+}
+
+func (r *rabbitMQConnection) resetRPCClientLocked(ch *amqp.Channel, err error) {
+	if ch != nil && !ch.IsClosed() {
+		ch.Close()
+	}
+	if r.rpcChannel == ch {
+		for correlationID, responseChannel := range r.rpcPending {
+			select {
+			case responseChannel <- rpcResponse{Err: err}:
+			default:
+			}
+			delete(r.rpcPending, correlationID)
 		}
+		r.rpcChannel = nil
+		r.rpcConfirm = nil
+		r.rpcReturn = nil
+		r.rpcExchanges = nil
+	}
+}
+
+func (r *rabbitMQConnection) removePendingRPCResponse(correlationID string) {
+	r.rpcClientMutex.Lock()
+	delete(r.rpcPending, correlationID)
+	r.rpcClientMutex.Unlock()
+}
+
+func (r *rabbitMQConnection) publishRPCMessage(exchange string, queue string, correlationID string, body []byte, exclusiveQueue bool, responseChannel chan rpcResponse) error {
+	r.rpcClientMutex.Lock()
+	defer r.rpcClientMutex.Unlock()
+	ch, confirmChannel, notifyReturnChannel, err := r.getRPCClientLocked(exchange, exclusiveQueue)
+	if err != nil {
+		logging.LogError(err, "Failed to get rabbitmq rpc channel", "queue", queue)
+		return err
+	}
+	r.rpcPending[correlationID] = responseChannel
+	msg := amqp.Publishing{
+		ContentType:   "application/json",
+		CorrelationId: correlationID,
+		Body:          body,
+		ReplyTo:       "amq.rabbitmq.reply-to",
+	}
+	err = ch.Publish(
+		exchange, // exchange
+		queue,    // routing key
+		true,     // mandatory
+		false,    // immediate
+		msg,      // publishing
+	)
+	if err != nil {
+		delete(r.rpcPending, correlationID)
+		logging.LogError(err, "there was an error publishing an rpc message", "queue", queue)
+		r.resetRPCClientLocked(ch, err)
+		return err
+	}
+	timeoutChannel := time.After(RPC_TIMEOUT)
+	for {
 		select {
 		case ntf := <-confirmChannel:
 			if !ntf.Ack {
 				err = errors.New("failed to deliver message, not ACK-ed by receiver")
-				finalError = err
+				delete(r.rpcPending, correlationID)
 				logging.LogError(err, "failed to deliver message to exchange/queue, notifyPublish", "queue", queue)
-				ch.Close()
-				time.Sleep(RPC_TIMEOUT)
+				r.resetRPCClientLocked(ch, err)
+				return err
+			}
+			return nil
+		case ret := <-notifyReturnChannel:
+			if ret.CorrelationId != correlationID {
 				continue
 			}
-		case ret := <-notifyReturnChannel:
 			err = errors.New(getMeaningfulRabbitmqError(ret))
+			delete(r.rpcPending, correlationID)
+			r.resetRPCClientLocked(ch, err)
+			return err
+		case <-timeoutChannel:
+			err = errors.New("message delivery confirmation timed out in SendRPCMessage")
+			delete(r.rpcPending, correlationID)
+			logging.LogError(err, "message delivery confirmation to exchange/queue timed out when sending", "queue", queue)
+			r.resetRPCClientLocked(ch, err)
+			return err
+		}
+	}
+}
+
+func (r *rabbitMQConnection) SendRPCMessage(exchange string, queue string, body []byte, exclusiveQueue bool, retryPolicy RPCRetryPolicy) ([]byte, error) {
+	var finalError error
+	timeout := r.getRPCTimeout(retryPolicy)
+	for attempt := 0; attempt < 3; attempt++ {
+		correlationID := uuid.NewString()
+		responseChannel := make(chan rpcResponse, 1)
+		err := r.publishRPCMessage(exchange, queue, correlationID, body, exclusiveQueue, responseChannel)
+		if err != nil {
 			finalError = err
 			time.Sleep(RPC_TIMEOUT)
-			ch.Close()
-			continue
-		case <-time.After(RPC_TIMEOUT):
-			err = errors.New("message delivery confirmation timed out in SendRPCMessage")
-			finalError = err
-			logging.LogError(err, "message delivery confirmation to exchange/queue timed out when sending", "queue", queue)
-			ch.Close()
 			continue
 		}
 		//logging.LogDebug("Sent RPC message", "queue", queue)
 		select {
-		case m := <-msgs:
+		case response := <-responseChannel:
+			if response.Err != nil {
+				finalError = response.Err
+				logging.LogError(response.Err, "rpc channel failed while waiting for response", "queue", queue)
+				if retryPolicy != RPC_RETRY_POLICY_RETRY_ON_TIMEOUT {
+					return nil, response.Err
+				}
+				time.Sleep(RPC_TIMEOUT)
+				continue
+			}
 			//logging.LogDebug("Got RPC Reply", "queue", queue)
-			ch.Close()
-			return m.Body, nil
-		case <-time.After(RPC_TIMEOUT):
-			err = errors.New("message delivery confirmation timed out")
+			return response.Body, nil
+		case <-time.After(timeout):
+			r.removePendingRPCResponse(correlationID)
+			err = errors.New("rpc response timed out")
 			finalError = err
 			logging.LogError(err, "message delivery confirmation to exchange/queue timed out when receiving", "queue", queue)
-			ch.Close()
+			if retryPolicy != RPC_RETRY_POLICY_RETRY_ON_TIMEOUT {
+				return nil, err
+			}
 			continue
 		}
 	}

mythic-docker/src/utils/config.go

@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/spf13/viper"
 )
@@ -40,6 +41,7 @@ type Config struct {
 	RabbitmqUser     string
 	RabbitmqPassword string
 	RabbitmqVHost    string
+	CustomRPCTimeout time.Duration
 
 	// postgres configuration
 	PostgresHost     string
@@ -84,6 +86,7 @@ func Initialize() {
 	mythicEnv.SetDefault("rabbitmq_user", "mythic_user")
 	mythicEnv.SetDefault("rabbitmq_password", "")
 	mythicEnv.SetDefault("rabbitmq_vhost", "mythic_vhost")
+	mythicEnv.SetDefault("custom_rpc_timeout", 0)
 	// jwt configuration
 	mythicEnv.SetDefault("jwt_secret", "")
 	// default operation configuration
@@ -189,6 +192,7 @@ func setConfigFromEnv(mythicEnv *viper.Viper) {
 	if MythicConfig.RabbitmqVHost == "" {
 		MythicConfig.RabbitmqVHost = "mythic_vhost"
 	}
+	MythicConfig.CustomRPCTimeout = time.Duration(mythicEnv.GetUint("custom_rpc_timeout")) * time.Second
 	// jwt configuration
 	MythicConfig.JWTSecret = []byte(mythicEnv.GetString("jwt_secret"))
 }