From f3c5978ee86feb599bb31d5f3a7008232fa76fe8 Mon Sep 17 00:00:00 2001 From: its-a-feature Date: Fri, 8 May 2026 20:11:48 -0500 Subject: [PATCH] better rabbitmq handling and task filter ui updates --- .../CallbacksTabsTaskingFilterDialog.js | 406 +++++++++++------- .../Callbacks/CallbacksTabsTaskingInput.js | 344 ++++++++------- .../CallbacksTabsTaskingInputTokenSelect.js | 8 +- .../pages/Callbacks/ResponseDisplay.js | 51 ++- .../Callbacks/ResponseDisplayInteractive.js | 18 +- MythicReactUI/src/themes/GlobalStyles.js | 377 ++++++++++++++++ Mythic_CLI/src/cmd/config/env.go | 5 +- .../src/cmd/internal/serviceMetadata.go | 1 + .../database/migrations/003003020_3.3.20.sql | 14 + mythic-docker/src/rabbitmq/constants.go | 11 + mythic-docker/src/rabbitmq/initialize.go | 10 + .../rabbitmq/send_auth_get_idp_metadata.go | 1 + .../rabbitmq/send_auth_get_idp_redirect.go | 1 + .../rabbitmq/send_auth_get_nonidp_metadata.go | 1 + .../rabbitmq/send_auth_get_nonidp_redirect.go | 1 + .../send_auth_process_idp_response.go | 1 + .../send_auth_process_nonidp_response.go | 1 + .../src/rabbitmq/send_c2_rpc_config_check.go | 1 + .../rabbitmq/send_c2_rpc_get_debug_output.go | 1 + .../src/rabbitmq/send_c2_rpc_get_ioc.go | 1 + .../send_c2_rpc_get_redirector_rules.go | 1 + .../src/rabbitmq/send_c2_rpc_host_file.go | 1 + .../src/rabbitmq/send_c2_rpc_opsec_check.go | 1 + .../src/rabbitmq/send_c2_rpc_resync.go | 1 + .../rabbitmq/send_c2_rpc_sample_message.go | 1 + .../src/rabbitmq/send_c2_rpc_start_server.go | 1 + .../src/rabbitmq/send_c2_rpc_stop_server.go | 1 + .../rabbitmq/send_container_rpc_get_file.go | 1 + .../rabbitmq/send_container_rpc_list_file.go | 1 + .../send_container_rpc_remove_file.go | 1 + .../rabbitmq/send_container_rpc_write_file.go | 1 + .../src/rabbitmq/send_pt_command_help.go | 1 + .../send_pt_rpc_check_if_callbacks_alive.go | 1 + ..._dynamic_query_build_parameter_function.go | 1 + .../send_pt_rpc_dynamic_query_function.go | 1 + .../src/rabbitmq/send_pt_rpc_resync.go | 1 + .../rabbitmq/send_pt_rpc_typedArray_parse.go | 1 + .../src/rabbitmq/send_tr_rpc_decrypt_bytes.go | 1 + .../src/rabbitmq/send_tr_rpc_encrypt_bytes.go | 1 + .../src/rabbitmq/send_tr_rpc_resync.go | 1 + .../src/rabbitmq/utils_rabbitmq_routing.go | 353 ++++++++++----- mythic-docker/src/utils/config.go | 4 + 42 files changed, 1198 insertions(+), 433 deletions(-) diff --git a/MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingFilterDialog.js b/MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingFilterDialog.js index 3004f53a..0890ae39 100644 --- a/MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingFilterDialog.js +++ b/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}) => ( + +) + +const MultiSelectField = ({label, value, options, onChange, emptyLabel}) => { + const renderValue = (selected) => { + if(selected.length === 0){ + return {emptyLabel} + } + return ( + + {selected.map((selectedValue) => ( + + ))} + + ); + } + + return ( + + + + ) +} + 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 ( - - Filter Which Tasks Are Visible in this Callback - - - - Only Show Tasks by the Following Operators - - - Only Show Tasks with Comments: -
- Hide Error Tasks: - - Only Show These Commands - - - {onlyCommands.length > 0 && - - } - - Do Not Show These Commands - - - {everythingBut.length > 0 && - - } - - - - - - - - + Clear exclude + + } + + + + + + + + + + + Clear all + + + Close + + + Apply filters + + + ); } - diff --git a/MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingInput.js b/MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingInput.js index 0ba440b4..c707af56 100644 --- a/MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingInput.js +++ b/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 ( + + + {label !== "" && + {label} + } + {value} + + + ) +} 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 ( -
+
{backdropOpen && } {reverseSearching && +
+ + reverse-i-search + + +
+ } + {taskingContextChips.length > 0 && +
+ {taskingContextChips.map((chip) => ( + + ))} +
+ } +
+ {tokenOptions.current.length > 0 ? ( + + ) : null} reverse-i-search: + spellCheck: false, + autoFocus: true, + startAdornment: + + + , + endAdornment: +
+ {commandPayloadType !== "" && + + + + + + } + {props.filterTasks && + + + + } + + + + +
}} /> - } - {callbackContext?.impersonation_context !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("impersonation_context") && - - - {"User: "}{callbackContext.impersonation_context} - - - - } - {callbackContext?.user !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("user") && - 2 ? " (high integrity)" : "")}> - - {"User: "}{callbackContext.user}{callbackContext.integrity_level > 2 ? "*" : ""} - - - } - {callbackContext?.cwd !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("cwd") && - - - {"Dir: "}{callbackContext.cwd} - - - } - {callbackContext?.host !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("host") && - - - {"Host: "}{callbackContext.host} - - - } - {callbackContext?.ip !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("ip") && - - - {"IP: "}{callbackContext.ip} - - - } - {callbackContext?.pid !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("pid") && - - - {"PID: "}{callbackContext.pid} - - - } - {callbackContext?.architecture !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("architecture") && - - - {"Arch: "}{callbackContext.architecture} - - - } - {callbackContext?.process_short_name !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("process_short_name") && - - - {"Process: "}{callbackContext.process_short_name} - - - } - {callbackContext?.extra_info !== "" && !hideTaskingContext.current && taskingContextFields.current.includes("extra_info") && - - - {callbackContext.extra_info} - - - } - - - - {props.filterTasks && - - } - {commandPayloadType !== "" && - - } - - , - startAdornment: - {tokenOptions.current.length > 0 ? ( - - ) : null} - - - - }} - /> +
{openFilterOptionsDialog && { diff --git a/MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingInputTokenSelect.js b/MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingInputTokenSelect.js index 9a6934eb..c4c63ef4 100644 --- a/MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingInputTokenSelect.js +++ b/MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingInputTokenSelect.js @@ -42,15 +42,16 @@ export function CallbacksTabsTaskingInputTokenSelect(props) { } } return ( - + } > Default Token {options.map( (opt) => ( @@ -60,4 +61,3 @@ export function CallbacksTabsTaskingInputTokenSelect(props) { ); } - diff --git a/MythicReactUI/src/components/pages/Callbacks/ResponseDisplay.js b/MythicReactUI/src/components/pages/Callbacks/ResponseDisplay.js index dd9cbe03..3752c401 100644 --- a/MythicReactUI/src/components/pages/Callbacks/ResponseDisplay.js +++ b/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; diff --git a/MythicReactUI/src/components/pages/Callbacks/ResponseDisplayInteractive.js b/MythicReactUI/src/components/pages/Callbacks/ResponseDisplayInteractive.js index ea58e096..3c109a8b 100644 --- a/MythicReactUI/src/components/pages/Callbacks/ResponseDisplayInteractive.js +++ b/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); } diff --git a/MythicReactUI/src/themes/GlobalStyles.js b/MythicReactUI/src/themes/GlobalStyles.js index f3aa9519..95740394 100644 --- a/MythicReactUI/src/themes/GlobalStyles.js +++ b/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; diff --git a/Mythic_CLI/src/cmd/config/env.go b/Mythic_CLI/src/cmd/config/env.go index 097e5bcb..a3091d3d 100644 --- a/Mythic_CLI/src/cmd/config/env.go +++ b/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` diff --git a/Mythic_CLI/src/cmd/internal/serviceMetadata.go b/Mythic_CLI/src/cmd/internal/serviceMetadata.go index 1992bc0b..b47e050e 100644 --- a/Mythic_CLI/src/cmd/internal/serviceMetadata.go +++ b/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}", diff --git a/mythic-docker/src/database/migrations/003003020_3.3.20.sql b/mythic-docker/src/database/migrations/003003020_3.3.20.sql index b4e05bb6..263767fe 100644 --- a/mythic-docker/src/database/migrations/003003020_3.3.20.sql +++ b/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"; diff --git a/mythic-docker/src/rabbitmq/constants.go b/mythic-docker/src/rabbitmq/constants.go index ec51021a..40c15696 100644 --- a/mythic-docker/src/rabbitmq/constants.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/initialize.go b/mythic-docker/src/rabbitmq/initialize.go index c526e3a3..811fa70b 100644 --- a/mythic-docker/src/rabbitmq/initialize.go +++ b/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 } diff --git a/mythic-docker/src/rabbitmq/send_auth_get_idp_metadata.go b/mythic-docker/src/rabbitmq/send_auth_get_idp_metadata.go index b4854a85..415e0016 100644 --- a/mythic-docker/src/rabbitmq/send_auth_get_idp_metadata.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_auth_get_idp_redirect.go b/mythic-docker/src/rabbitmq/send_auth_get_idp_redirect.go index 70229b57..36d42602 100644 --- a/mythic-docker/src/rabbitmq/send_auth_get_idp_redirect.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_auth_get_nonidp_metadata.go b/mythic-docker/src/rabbitmq/send_auth_get_nonidp_metadata.go index 0d38d761..69013db4 100644 --- a/mythic-docker/src/rabbitmq/send_auth_get_nonidp_metadata.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_auth_get_nonidp_redirect.go b/mythic-docker/src/rabbitmq/send_auth_get_nonidp_redirect.go index c3535687..b1a4beac 100644 --- a/mythic-docker/src/rabbitmq/send_auth_get_nonidp_redirect.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_auth_process_idp_response.go b/mythic-docker/src/rabbitmq/send_auth_process_idp_response.go index 35fd6952..101f978a 100644 --- a/mythic-docker/src/rabbitmq/send_auth_process_idp_response.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_auth_process_nonidp_response.go b/mythic-docker/src/rabbitmq/send_auth_process_nonidp_response.go index 04f321ce..2186218a 100644 --- a/mythic-docker/src/rabbitmq/send_auth_process_nonidp_response.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_c2_rpc_config_check.go b/mythic-docker/src/rabbitmq/send_c2_rpc_config_check.go index 78fe1abb..d8cd1c44 100644 --- a/mythic-docker/src/rabbitmq/send_c2_rpc_config_check.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_c2_rpc_get_debug_output.go b/mythic-docker/src/rabbitmq/send_c2_rpc_get_debug_output.go index 83c832f7..27ffb69b 100644 --- a/mythic-docker/src/rabbitmq/send_c2_rpc_get_debug_output.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/send_c2_rpc_get_ioc.go b/mythic-docker/src/rabbitmq/send_c2_rpc_get_ioc.go index 4e0d4b9e..ee20a6d8 100644 --- a/mythic-docker/src/rabbitmq/send_c2_rpc_get_ioc.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/send_c2_rpc_get_redirector_rules.go b/mythic-docker/src/rabbitmq/send_c2_rpc_get_redirector_rules.go index 35e1a9ff..31239f62 100644 --- a/mythic-docker/src/rabbitmq/send_c2_rpc_get_redirector_rules.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/send_c2_rpc_host_file.go b/mythic-docker/src/rabbitmq/send_c2_rpc_host_file.go index 9df436d5..66999ec8 100644 --- a/mythic-docker/src/rabbitmq/send_c2_rpc_host_file.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_c2_rpc_opsec_check.go b/mythic-docker/src/rabbitmq/send_c2_rpc_opsec_check.go index cd91cd11..1544ab53 100644 --- a/mythic-docker/src/rabbitmq/send_c2_rpc_opsec_check.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_c2_rpc_resync.go b/mythic-docker/src/rabbitmq/send_c2_rpc_resync.go index a2f3b227..8faee496 100644 --- a/mythic-docker/src/rabbitmq/send_c2_rpc_resync.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/send_c2_rpc_sample_message.go b/mythic-docker/src/rabbitmq/send_c2_rpc_sample_message.go index 09ef9b03..4155c7e1 100644 --- a/mythic-docker/src/rabbitmq/send_c2_rpc_sample_message.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/send_c2_rpc_start_server.go b/mythic-docker/src/rabbitmq/send_c2_rpc_start_server.go index e48db875..9ac09a6d 100644 --- a/mythic-docker/src/rabbitmq/send_c2_rpc_start_server.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_c2_rpc_stop_server.go b/mythic-docker/src/rabbitmq/send_c2_rpc_stop_server.go index c7f0660c..7010bdac 100644 --- a/mythic-docker/src/rabbitmq/send_c2_rpc_stop_server.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_container_rpc_get_file.go b/mythic-docker/src/rabbitmq/send_container_rpc_get_file.go index 6a2a6582..99577dc1 100644 --- a/mythic-docker/src/rabbitmq/send_container_rpc_get_file.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/send_container_rpc_list_file.go b/mythic-docker/src/rabbitmq/send_container_rpc_list_file.go index 52828921..295b3e00 100644 --- a/mythic-docker/src/rabbitmq/send_container_rpc_list_file.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/send_container_rpc_remove_file.go b/mythic-docker/src/rabbitmq/send_container_rpc_remove_file.go index c6bc1bf2..8d28a228 100644 --- a/mythic-docker/src/rabbitmq/send_container_rpc_remove_file.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/send_container_rpc_write_file.go b/mythic-docker/src/rabbitmq/send_container_rpc_write_file.go index 91558339..49ea2d16 100644 --- a/mythic-docker/src/rabbitmq/send_container_rpc_write_file.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/send_pt_command_help.go b/mythic-docker/src/rabbitmq/send_pt_command_help.go index 1cb7df9f..5511cf84 100644 --- a/mythic-docker/src/rabbitmq/send_pt_command_help.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_pt_rpc_check_if_callbacks_alive.go b/mythic-docker/src/rabbitmq/send_pt_rpc_check_if_callbacks_alive.go index e5382cbb..9589d9d3 100644 --- a/mythic-docker/src/rabbitmq/send_pt_rpc_check_if_callbacks_alive.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/send_pt_rpc_dynamic_query_build_parameter_function.go b/mythic-docker/src/rabbitmq/send_pt_rpc_dynamic_query_build_parameter_function.go index 76ba9e62..7b07c7bb 100644 --- a/mythic-docker/src/rabbitmq/send_pt_rpc_dynamic_query_build_parameter_function.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_pt_rpc_dynamic_query_function.go b/mythic-docker/src/rabbitmq/send_pt_rpc_dynamic_query_function.go index 236abe38..af15e918 100644 --- a/mythic-docker/src/rabbitmq/send_pt_rpc_dynamic_query_function.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_pt_rpc_resync.go b/mythic-docker/src/rabbitmq/send_pt_rpc_resync.go index 164da5b8..859eabbd 100644 --- a/mythic-docker/src/rabbitmq/send_pt_rpc_resync.go +++ b/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") diff --git a/mythic-docker/src/rabbitmq/send_pt_rpc_typedArray_parse.go b/mythic-docker/src/rabbitmq/send_pt_rpc_typedArray_parse.go index adcddd58..8f7a0981 100644 --- a/mythic-docker/src/rabbitmq/send_pt_rpc_typedArray_parse.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/send_tr_rpc_decrypt_bytes.go b/mythic-docker/src/rabbitmq/send_tr_rpc_decrypt_bytes.go index d3db5954..a6dcc1c2 100644 --- a/mythic-docker/src/rabbitmq/send_tr_rpc_decrypt_bytes.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/send_tr_rpc_encrypt_bytes.go b/mythic-docker/src/rabbitmq/send_tr_rpc_encrypt_bytes.go index a7afc317..a6a24ea8 100644 --- a/mythic-docker/src/rabbitmq/send_tr_rpc_encrypt_bytes.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/send_tr_rpc_resync.go b/mythic-docker/src/rabbitmq/send_tr_rpc_resync.go index 358541af..e8da8bc0 100644 --- a/mythic-docker/src/rabbitmq/send_tr_rpc_resync.go +++ b/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 diff --git a/mythic-docker/src/rabbitmq/utils_rabbitmq_routing.go b/mythic-docker/src/rabbitmq/utils_rabbitmq_routing.go index aef0677d..54095ed2 100644 --- a/mythic-docker/src/rabbitmq/utils_rabbitmq_routing.go +++ b/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 } } diff --git a/mythic-docker/src/utils/config.go b/mythic-docker/src/utils/config.go index f6d225c7..8fad7e3c 100644 --- a/mythic-docker/src/utils/config.go +++ b/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")) }