diff --git a/MythicReactUI/src/components/MythicComponents/MythicTabPanel.js b/MythicReactUI/src/components/MythicComponents/MythicTabPanel.js index 7122b24d..fe06c689 100644 --- a/MythicReactUI/src/components/MythicComponents/MythicTabPanel.js +++ b/MythicReactUI/src/components/MythicComponents/MythicTabPanel.js @@ -172,10 +172,9 @@ export function MythicSearchTabLabel(props) { return ( - {label} -
+ {iconComponent} + {label} } {...a11yProps(index)} diff --git a/MythicReactUI/src/components/MythicComponents/MythicTableToolbar.js b/MythicReactUI/src/components/MythicComponents/MythicTableToolbar.js index faeef724..09bf07c2 100644 --- a/MythicReactUI/src/components/MythicComponents/MythicTableToolbar.js +++ b/MythicReactUI/src/components/MythicComponents/MythicTableToolbar.js @@ -17,9 +17,12 @@ export const MythicTableToolbar = ({children, className = "", style = {}, varian ); }; -export const MythicTableToolbarGroup = ({children, grow = false, className = "", style = {}}) => { +export const MythicTableToolbarGroup = ({children, grow = false, label, className = "", style = {}}) => { return ( + {label && + {label} + } {children} ); diff --git a/MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingInput.js b/MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingInput.js index 01a81b18..f7709df6 100644 --- a/MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingInput.js +++ b/MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingInput.js @@ -226,9 +226,11 @@ const getTaskingParameterTooltip = (parameter) => { const description = parameter.description ? `\n${parameter.description}` : ""; return `${label} ยท ${parameter.parameter_type}${requiredText}${displayName}${internalName}${description}`; } -const TaskingParameterPreviewChip = ({parameter, required=false, active=false}) => ( +const TaskingParameterPreviewChip = ({parameter, required=false, active=false, onClick}) => ( @@ -239,6 +241,14 @@ const TaskingParameterPreviewChip = ({parameter, required=false, active=false}) {parameter.parameter_type} } + onClick={onClick ? (event) => { + event.preventDefault(); + event.stopPropagation(); + onClick(parameter); + } : undefined} + onMouseDown={onClick ? (event) => { + event.preventDefault(); + } : undefined} size="small" /> @@ -394,6 +404,7 @@ export function CallbacksTabsTaskingInputPreMemo(props){ const toastId = "tasking-toast-message"; const theme = useTheme(); const inputRef = React.useRef(null); + const pendingCursorPosition = React.useRef(null); const snackMessageStyles = {position:"bottom-left", autoClose: 1000, toastId: toastId, style: {marginBottom: "30px"}}; const snackReverseSearchMessageStyles = {position:"bottom-left", autoClose: 1000, toastId: toastId, style: {marginBottom: "70px"}}; const [commandPayloadType, setCommandPayloadType] = React.useState(""); @@ -641,6 +652,50 @@ export function CallbacksTabsTaskingInputPreMemo(props){ setUnmodifiedHistoryValue("parsed_cli"); } } + React.useEffect( () => { + if(pendingCursorPosition.current === null){ + return; + } + const cursorPosition = pendingCursorPosition.current; + pendingCursorPosition.current = null; + window.requestAnimationFrame(() => { + if(inputRef.current){ + inputRef.current.focus(); + if(inputRef.current.setSelectionRange){ + const adjustedCursorPosition = Math.min(cursorPosition, inputRef.current.value.length); + inputRef.current.setSelectionRange(adjustedCursorPosition, adjustedCursorPosition); + } + } + }); + }, [message]) + const insertParameterIntoCommandLine = (parameter) => { + if(!parameter?.cli_name){ + return; + } + const input = inputRef.current; + let selectionStart = typeof input?.selectionStart === "number" ? input.selectionStart : message.length; + let selectionEnd = typeof input?.selectionEnd === "number" ? input.selectionEnd : selectionStart; + if(selectionStart === selectionEnd){ + const beforeCursor = message.slice(0, selectionStart); + const partialParameterMatch = beforeCursor.match(/(^|\s)(-\S*)$/); + if(partialParameterMatch){ + selectionStart -= partialParameterMatch[2].length; + } + } + const before = message.slice(0, selectionStart); + let after = message.slice(selectionEnd); + const parameterText = `-${parameter.cli_name}`; + const insertText = `${before.length > 0 && !/\s$/.test(before) ? " " : ""}${parameterText} `; + if(after.length > 0 && !/^\s/.test(after)){ + after = ` ${after}`; + } + tabOptions.current = []; + tabOptionsType.current = "param_name"; + tabOptionsIndex.current = 0; + setUnmodifiedHistoryValue("parsed_cli"); + pendingCursorPosition.current = before.length + insertText.length; + setMessage(`${before}${insertText}${after}`); + } const onKeyDown = (event) => { if(event.key === "Enter" && (event.ctrlKey || event.metaKey)){ setMessage(message + "\n"); @@ -2117,7 +2172,7 @@ export function CallbacksTabsTaskingInputPreMemo(props){ } {commandParameterPreview.requiredParameters.slice(0, commandParameterPreview.activeParameter ? 5 : 6).map((parameter) => ( - + ))} {commandParameterPreview.requiredParameters.length > (commandParameterPreview.activeParameter ? 5 : 6) && @@ -2125,7 +2180,7 @@ export function CallbacksTabsTaskingInputPreMemo(props){ } {commandParameterPreview.optionalParameters.slice(0, commandParameterPreview.requiredParameters.length > 0 ? 4 : (commandParameterPreview.activeParameter ? 5 : 6)).map((parameter) => ( - + ))} {commandParameterPreview.optionalParameters.length > (commandParameterPreview.requiredParameters.length > 0 ? 4 : (commandParameterPreview.activeParameter ? 5 : 6)) && diff --git a/MythicReactUI/src/components/pages/Search/Search.js b/MythicReactUI/src/components/pages/Search/Search.js index 7e35b8b0..50169001 100644 --- a/MythicReactUI/src/components/pages/Search/Search.js +++ b/MythicReactUI/src/components/pages/Search/Search.js @@ -97,7 +97,7 @@ export function Search(props){ ({ backgroundColor: theme.surfaces?.muted || theme.palette.background.paper, border: `1px solid ${theme.table?.borderSoft || theme.borderColor}`, @@ -106,6 +106,7 @@ export function Search(props){ })} > { } }, [props.value, props.index]) return ( - - + + - + { )) } + + { } - + {createArtifactDialogOpen && {setCreateArtifactDialogOpen(false);}} diff --git a/MythicReactUI/src/components/pages/Search/SearchTabCallbacks.js b/MythicReactUI/src/components/pages/Search/SearchTabCallbacks.js index da563a34..7080b277 100644 --- a/MythicReactUI/src/components/pages/Search/SearchTabCallbacks.js +++ b/MythicReactUI/src/components/pages/Search/SearchTabCallbacks.js @@ -231,10 +231,10 @@ const SearchTabCallbacksSearchPanel = (props) => { }, [props.value, props.index]) return ( - + - + { }, [props.value, props.index]) return ( - + - + { } - + {createCredentialDialogOpen && {setCreateCredentialDialogOpen(false);}} diff --git a/MythicReactUI/src/components/pages/Search/SearchTabCustomBrowsers.js b/MythicReactUI/src/components/pages/Search/SearchTabCustomBrowsers.js index 01c7375f..9ea71be9 100644 --- a/MythicReactUI/src/components/pages/Search/SearchTabCustomBrowsers.js +++ b/MythicReactUI/src/components/pages/Search/SearchTabCustomBrowsers.js @@ -195,8 +195,8 @@ const SearchTabCustomBrowserSearchPanel = (props) => { } }, [props.value, props.index]); return ( - - + + { } - + - + - + { }, [props.value, props.index]); return ( - + - + - + { } - + { } - + }> Files diff --git a/MythicReactUI/src/components/pages/Search/SearchTabKeylog.js b/MythicReactUI/src/components/pages/Search/SearchTabKeylog.js index d718924e..46f35bdb 100644 --- a/MythicReactUI/src/components/pages/Search/SearchTabKeylog.js +++ b/MythicReactUI/src/components/pages/Search/SearchTabKeylog.js @@ -172,11 +172,11 @@ const SearchTabKeylogsSearchPanel = (props) => { } }, [props.value, props.index]) return ( - - + + - + { }, []); return ( - + - + { } - + { - + { } - + { } }, [props.value, props.index]); return ( - - + + - + - + { const [callbackData, setCallbackData] = React.useState([]); const [showDeleted, setShowDeleted] = React.useState(false); + const [search, setSearch] = React.useState(""); + const handleSearchValueChange = (name, value, error) => { + setSearch(value); + } + const submitSearch = () => { + props.changeSearchParam("search", search); + } + const handleToggleShowDeleted = () => { + const nextShowDeleted = !showDeleted; + setShowDeleted(nextShowDeleted); + props.changeSearchParam("showStopped", nextShowDeleted ? "true" : "false"); + } + React.useEffect(() => { + if(props.value === props.index){ + const queryParams = new URLSearchParams(window.location.search); + if(queryParams.has("search")){ + setSearch(queryParams.get("search")); + } + if(queryParams.has("showStopped")){ + setShowDeleted(queryParams.get("showStopped") === "true"); + } + } + }, [props.value, props.index]); useSubscription(callbackPortsSub, { fetchPolicy: "no-cache", onData: ({data}) => { @@ -79,14 +102,48 @@ export const SearchTabSocksPanel = (props) =>{ console.log(data); } }); + const filteredCallbacks = React.useMemo(() => { + const normalizedSearch = search.trim().toLowerCase(); + if(normalizedSearch === ""){ + return callbackData; + } + return callbackData.filter((proxy) => { + const searchableText = [ + proxy.callback?.user, + proxy.callback?.host, + proxy.callback?.description, + proxy.callback?.domain, + proxy.callback?.display_id ? `C-${proxy.callback.display_id}` : "", + proxy.task?.display_id ? `T-${proxy.task.display_id}` : "", + proxy.local_port, + proxy.remote_ip, + proxy.remote_port, + proxy.port_type, + proxy.username, + ].join(" ").toLowerCase(); + return searchableText.includes(normalizedSearch); + }); + }, [callbackData, search]); + const visibleCallbacks = React.useMemo(() => { + return filteredCallbacks.filter((proxy) => showDeleted || !proxy.deleted); + }, [filteredCallbacks, showDeleted]); return ( - - + + + + + setShowDeleted(!showDeleted)} + onClick={handleToggleShowDeleted} label="Stopped" activeIcon={} inactiveIcon={} @@ -94,11 +151,13 @@ export const SearchTabSocksPanel = (props) =>{
- {callbackData.length > 0 ? ( - ) : ( + {visibleCallbacks.length > 0 ? ( + ) : ( )} diff --git a/MythicReactUI/src/components/pages/Search/SearchTabTags.js b/MythicReactUI/src/components/pages/Search/SearchTabTags.js index ca74f2ae..8914208f 100644 --- a/MythicReactUI/src/components/pages/Search/SearchTabTags.js +++ b/MythicReactUI/src/components/pages/Search/SearchTabTags.js @@ -195,11 +195,11 @@ const SearchTabTagsSearchPanel = (props) => { } }, [props.value, props.index]); return ( - - + + - + { }, [props.value, props.index]); return ( - + - + - + { } - + - + { } }, [props.value, props.index]) return ( - - + + - + alpha(props.theme.palette.info.main, props.theme.palette.mode === "dark" ? 0.12 : 0.07)}; + border-color: ${(props) => alpha(props.theme.palette.info.main, props.theme.palette.mode === "dark" ? 0.36 : 0.24)}; + color: ${(props) => props.theme.palette.info.main}; +} .mythic-tasking-parameter-preview-chip-required.MuiChip-root { background-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.14 : 0.08)}; border-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.36 : 0.24)}; color: ${(props) => props.theme.palette.warning.main}; } +.mythic-tasking-parameter-preview-chip-required.MuiChip-clickable:hover { + background-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.18 : 0.1)}; + border-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.48 : 0.32)}; + color: ${(props) => props.theme.palette.warning.main}; +} .mythic-tasking-parameter-preview-chip-active.MuiChip-root { background-color: ${(props) => alpha(props.theme.palette.info.main, props.theme.palette.mode === "dark" ? 0.14 : 0.08)}; border-color: ${(props) => alpha(props.theme.palette.info.main, props.theme.palette.mode === "dark" ? 0.42 : 0.28)}; @@ -2627,10 +2637,11 @@ tspan { .mythic-table-toolbar-search { align-items: stretch; background-color: ${(props) => props.theme.palette.mode === "dark" ? "rgba(255,255,255,0.028)" : "rgba(0,0,0,0.012)"}; - gap: 0.45rem; - padding: 0.45rem; + gap: 0.55rem; + padding: 0.55rem; } .mythic-table-toolbar-group { + align-content: center; align-items: center; display: flex; flex: 0 1 auto; @@ -2642,6 +2653,14 @@ tspan { .mythic-table-toolbar-search .mythic-table-toolbar-group { gap: 0.35rem; } +.mythic-table-toolbar-group-label { + color: ${(props) => props.theme.palette.text.secondary}; + flex: 0 0 100%; + font-size: 0.66rem; + font-weight: 800; + line-height: 1; + padding-left: 0.1rem; +} .mythic-table-toolbar-group-grow { flex: 1 1 20rem; min-width: min(100%, 16rem); @@ -8443,6 +8462,30 @@ tspan { .MuiTabs-scrollButtons.Mui-disabled { opacity: 0.3; } +.mythic-search-tabs-bar.MuiPaper-root { + background-image: ${getSectionHeaderGradient}; +} +.mythic-search-tabs.MuiTabs-root { + min-height: 42px; +} +.mythic-search-tabs .MuiTabs-flexContainer { + flex-wrap: nowrap; +} +.mythic-search-tabs .MuiTab-root { + min-height: 42px; + padding: 0.45rem 0.7rem; +} +.mythic-search-tab-label { + align-items: center; + display: inline-flex; + gap: 0.42rem; + white-space: nowrap; +} +.mythic-search-tab-label .MuiSvgIcon-root, +.mythic-search-tab-label .svg-inline--fa { + color: inherit; + font-size: 1rem; +} .ace_editor{ background-color: ${(props) => props.theme.outputBackgroundColor + "20"}; } 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 26a26245..5293d921 100644 --- a/mythic-docker/src/database/migrations/003003020_3.3.20.sql +++ b/mythic-docker/src/database/migrations/003003020_3.3.20.sql @@ -76,6 +76,11 @@ create unique index if not exists response_task_sequence_number_unique on "public"."response" using btree (task_id, sequence_number) where sequence_number is not null; +alter table "public"."task" + add column if not exists subtask_callback_function_started boolean not null default false, + add column if not exists group_callback_function_started boolean not null default false, + add column if not exists completed_callback_function_started boolean not null default false; + -- +migrate StatementBegin create or replace function public.new_task_display_id() returns trigger language plpgsql @@ -201,4 +206,8 @@ $$; drop index if exists "public"."callback_operation_display_id_unique"; drop index if exists "public"."task_operation_display_id_unique"; drop index if exists "public"."response_task_sequence_number_unique"; +alter table "public"."task" + drop column if exists completed_callback_function_started, + drop column if exists group_callback_function_started, + drop column if exists subtask_callback_function_started; drop table if exists "public"."operation_display_counters"; diff --git a/mythic-docker/src/database/structs/Task.go b/mythic-docker/src/database/structs/Task.go index 6dd83a1a..557c5c58 100644 --- a/mythic-docker/src/database/structs/Task.go +++ b/mythic-docker/src/database/structs/Task.go @@ -47,10 +47,13 @@ type Task struct { ParentTaskID structs.NullInt64 `db:"parent_task_id" json:"parent_task_id"` SubtaskCallbackFunction string `db:"subtask_callback_function" json:"subtask_callback_function"` SubtaskCallbackFunctionCompleted bool `db:"subtask_callback_function_completed" json:"subtask_callback_function_completed"` + SubtaskCallbackFunctionStarted bool `db:"subtask_callback_function_started" json:"subtask_callback_function_started"` GroupCallbackFunction string `db:"group_callback_function" json:"group_callback_function"` GroupCallbackFunctionCompleted bool `db:"group_callback_function_completed" json:"group_callback_function_completed"` + GroupCallbackFunctionStarted bool `db:"group_callback_function_started" json:"group_callback_function_started"` CompletedCallbackFunction string `db:"completed_callback_function" json:"completed_callback_function"` CompletedCallbackFunctionCompleted bool `db:"completed_callback_function_completed" json:"completed_callback_function_completed"` + CompletedCallbackFunctionStarted bool `db:"completed_callback_function_started" json:"completed_callback_function_started"` SubtaskGroupName string `db:"subtask_group_name" json:"subtask_group_name"` TaskingLocation string `db:"tasking_location" json:"tasking_location"` ParameterGroupName string `db:"parameter_group_name" json:"parameter_group_name"` diff --git a/mythic-docker/src/rabbitmq/util_agent_message_actions_post_response.go b/mythic-docker/src/rabbitmq/util_agent_message_actions_post_response.go index 6e0ca090..bee3d9f7 100644 --- a/mythic-docker/src/rabbitmq/util_agent_message_actions_post_response.go +++ b/mythic-docker/src/rabbitmq/util_agent_message_actions_post_response.go @@ -294,14 +294,14 @@ type writeDownloadChunkToDisk struct { var writeDownloadChunkToDiskChan = make(chan writeDownloadChunkToDisk) -type agentAgentMessagePostResponseChannelMessage struct { +type agentMessagePostResponseUserOutputChannelMessage struct { Response string SequenceNum *int64 Task databaseStructs.Task } -var asyncAgentMessagePostResponseWorkerChannels []chan agentAgentMessagePostResponseChannelMessage -var asyncAgentMessagePostResponseWorkersOnce sync.Once +var asyncAgentMessagePostResponseUserOutputWorkerChannels []chan agentMessagePostResponseUserOutputChannelMessage +var asyncAgentMessagePostResponseUserOutputWorkersOnce sync.Once var responseInterceptMapOperationIDToEventGroupID = make(map[int]int) var responseInterceptMapLock = sync.RWMutex{} @@ -320,25 +320,25 @@ func getAgentMessagePostResponseQueueSize() int { return defaultAgentMessagePostResponseQueueSize } -func initializeAsyncAgentMessagePostResponseWorkers() { +func initializeAsyncAgentMessagePostResponseUserOutputWorkers() { UpdateCachedResponseIntercept() workerCount := getAgentMessagePostResponseWorkerCount() queueSize := getAgentMessagePostResponseQueueSize() - asyncAgentMessagePostResponseWorkerChannels = make([]chan agentAgentMessagePostResponseChannelMessage, workerCount) + asyncAgentMessagePostResponseUserOutputWorkerChannels = make([]chan agentMessagePostResponseUserOutputChannelMessage, workerCount) for i := 0; i < workerCount; i++ { - asyncAgentMessagePostResponseWorkerChannels[i] = make(chan agentAgentMessagePostResponseChannelMessage, queueSize) - go listenForAsyncAgentMessagePostResponseWorker(i, asyncAgentMessagePostResponseWorkerChannels[i]) + asyncAgentMessagePostResponseUserOutputWorkerChannels[i] = make(chan agentMessagePostResponseUserOutputChannelMessage, queueSize) + go listenForAsyncAgentMessagePostResponseUserOutputWorker(i, asyncAgentMessagePostResponseUserOutputWorkerChannels[i]) } } -func enqueueAsyncAgentMessagePostResponse(msg agentAgentMessagePostResponseChannelMessage) { - asyncAgentMessagePostResponseWorkersOnce.Do(initializeAsyncAgentMessagePostResponseWorkers) - workerCount := len(asyncAgentMessagePostResponseWorkerChannels) +func enqueueAsyncAgentMessagePostResponse(msg agentMessagePostResponseUserOutputChannelMessage) { + asyncAgentMessagePostResponseUserOutputWorkersOnce.Do(initializeAsyncAgentMessagePostResponseUserOutputWorkers) + workerCount := len(asyncAgentMessagePostResponseUserOutputWorkerChannels) workerID := msg.Task.ID % workerCount if workerID < 0 { workerID = 0 } - asyncAgentMessagePostResponseWorkerChannels[workerID] <- msg + asyncAgentMessagePostResponseUserOutputWorkerChannels[workerID] <- msg } func UpdateCachedResponseIntercept() { @@ -374,10 +374,10 @@ func getCachedResponseInterceptEventGroupID(operationID int) (int, bool) { } func listenForAsyncAgentMessagePostResponseContent() { - asyncAgentMessagePostResponseWorkersOnce.Do(initializeAsyncAgentMessagePostResponseWorkers) + asyncAgentMessagePostResponseUserOutputWorkersOnce.Do(initializeAsyncAgentMessagePostResponseUserOutputWorkers) } -func listenForAsyncAgentMessagePostResponseWorker(workerID int, input <-chan agentAgentMessagePostResponseChannelMessage) { +func listenForAsyncAgentMessagePostResponseUserOutputWorker(workerID int, input <-chan agentMessagePostResponseUserOutputChannelMessage) { insertStatement, err := database.DB.Preparex(insertAgentMessagePostResponseUserOutputQuery) if err != nil { logging.LogError(err, "Failed to prepare user_output insert statement for post_response worker", "worker_id", workerID) @@ -386,11 +386,11 @@ func listenForAsyncAgentMessagePostResponseWorker(workerID int, input <-chan age defer insertStatement.Close() } for msg := range input { - processAsyncAgentMessagePostResponseContent(msg, insertStatement) + processAsyncAgentMessagePostResponseUserOutput(msg, insertStatement) } } -func processAsyncAgentMessagePostResponseContent(msg agentAgentMessagePostResponseChannelMessage, insertStatement *sqlx.Stmt) { +func processAsyncAgentMessagePostResponseUserOutput(msg agentMessagePostResponseUserOutputChannelMessage, insertStatement *sqlx.Stmt) { // force chunking user_output can be useful, but might also affect command's browserscripts //chunkSize := 1024 * 1024 //chunks := int(len(msg.Response)/chunkSize) + 1 // 1MB chunks @@ -451,6 +451,44 @@ func getAgentMessagePostResponseTasks(responses []agentMessagePostResponse) (map return tasksByAgentTaskID, nil } +func updateAgentMessagePostResponseTask(task databaseStructs.Task) (bool, error) { + tx, err := database.DB.Beginx() + if err != nil { + return false, err + } + defer tx.Rollback() + + wasCompleted := false + if err = tx.Get(&wasCompleted, `SELECT completed FROM task WHERE id=$1 FOR UPDATE`, task.ID); err != nil { + return false, err + } + + updatedCompleted := false + err = tx.Get(&updatedCompleted, `UPDATE task SET + status=CASE WHEN completed THEN status ELSE $2 END, + completed=(completed OR $3), + status_timestamp_processed=coalesce(status_timestamp_processed, $4), + "timestamp"=$5, + stdout=$6, + stderr=$7 + WHERE id=$1 + RETURNING completed`, + task.ID, + task.Status, + task.Completed, + task.StatusTimestampProcessed, + task.Timestamp, + task.Stdout, + task.Stderr) + if err != nil { + return false, err + } + if err = tx.Commit(); err != nil { + return false, err + } + return !wasCompleted && updatedCompleted, nil +} + func handleAgentMessagePostResponse(incoming *map[string]interface{}, uUIDInfo *cachedUUIDInfo) (map[string]interface{}, error) { // got message: /* @@ -576,7 +614,7 @@ func handleAgentMessagePostResponse(incoming *map[string]interface{}, uUIDInfo * if agentMessage.Responses[i].UserOutput != nil && *agentMessage.Responses[i].UserOutput != "" { // do it in the background - the agent doesn't need the result of this directly //handleAgentMessagePostResponseUserOutput(currentTask, agentResponse, true) - enqueueAsyncAgentMessagePostResponse(agentAgentMessagePostResponseChannelMessage{ + enqueueAsyncAgentMessagePostResponse(agentMessagePostResponseUserOutputChannelMessage{ Task: currentTask, Response: *agentMessage.Responses[i].UserOutput, SequenceNum: agentMessage.Responses[i].SequenceNumber, @@ -658,15 +696,11 @@ func handleAgentMessagePostResponse(incoming *map[string]interface{}, uUIDInfo * delete(*incoming, "responses") for _, currentTask := range tasksToUpdate { // always updating at least the timestamp for the last thing that happened - _, err = database.DB.NamedExec(`UPDATE task SET - status=:status, completed=:completed, status_timestamp_processed=:status_timestamp_processed, "timestamp"=:timestamp, - stdout=:stdout, stderr=:stderr - WHERE id=:id`, currentTask) + transitionedToCompleted, err := updateAgentMessagePostResponseTask(currentTask) if err != nil { logging.LogError(err, "Failed to update task from agent response") } - if currentTask.Completed { - // use updatedToCompleted to try to make sure we only do this once per task + if transitionedToCompleted { go CheckAndProcessTaskCompletionHandlers(currentTask.ID) FileBrowserChannel <- FileBrowserChannelMessage{ Task: ¤tTask, @@ -1933,7 +1967,7 @@ func HandleAgentMessagePostResponseFileBrowser(task databaseStructs.Task, fileBr logging.LogError(err, "failed to marshal filebrowser data to JSON") return } - enqueueAsyncAgentMessagePostResponse(agentAgentMessagePostResponseChannelMessage{ + enqueueAsyncAgentMessagePostResponse(agentMessagePostResponseUserOutputChannelMessage{ Task: task, Response: string(outputBytes), SequenceNum: nil, @@ -2859,7 +2893,7 @@ func handleAgentMessagePostResponseCustomBrowser(task databaseStructs.Task, agen logging.LogError(err, "failed to marshal filebrowser data to JSON") return } - enqueueAsyncAgentMessagePostResponse(agentAgentMessagePostResponseChannelMessage{ + enqueueAsyncAgentMessagePostResponse(agentMessagePostResponseUserOutputChannelMessage{ Task: task, Response: string(outputBytes), SequenceNum: nil, diff --git a/mythic-docker/src/rabbitmq/utils.go b/mythic-docker/src/rabbitmq/utils.go index 275ebd65..19c5a12e 100644 --- a/mythic-docker/src/rabbitmq/utils.go +++ b/mythic-docker/src/rabbitmq/utils.go @@ -767,6 +767,62 @@ func GetCallbackCommandInformation(callback databaseStructs.Callback) []string { var completionMutex sync.Mutex +func claimTaskCompletionFunctionStart(taskID int, startedColumn string, completedColumn string, functionColumn string, runningStatus PT_TASK_FUNCTION_STATUS) bool { + query := fmt.Sprintf(`UPDATE task + SET %s=true, status=$2 + WHERE id=$1 + AND completed=true + AND %s != '' + AND %s=false + AND %s=false`, startedColumn, functionColumn, startedColumn, completedColumn) + result, err := database.DB.Exec(query, taskID, runningStatus) + if err != nil { + logging.LogError(err, "Failed to claim task completion function start", "task_id", taskID, "started_column", startedColumn) + return false + } + rowsAffected, err := result.RowsAffected() + if err != nil { + logging.LogError(err, "Failed to check task completion function claim result", "task_id", taskID, "started_column", startedColumn) + return false + } + return rowsAffected == 1 +} + +func claimGroupCompletionFunctionStart(parentTaskID int64, groupName string, groupCallbackFunction string) bool { + result, err := database.DB.Exec(`UPDATE task + SET group_callback_function_started=true, status=$4 + WHERE parent_task_id=$1 + AND subtask_group_name=$2 + AND group_callback_function=$3 + AND group_callback_function != '' + AND completed=true + AND group_callback_function_started=false + AND group_callback_function_completed=false + AND NOT EXISTS ( + SELECT 1 + FROM task existing_group_task + WHERE existing_group_task.parent_task_id=$1 + AND existing_group_task.subtask_group_name=$2 + AND existing_group_task.group_callback_function=$3 + AND ( + existing_group_task.completed=false + OR existing_group_task.group_callback_function_started=true + OR existing_group_task.group_callback_function_completed=true + ) + )`, + parentTaskID, groupName, groupCallbackFunction, PT_TASK_FUNCTION_STATUS_GROUP_COMPLETED_FUNCTION) + if err != nil { + logging.LogError(err, "Failed to claim group completion function start", "parent_task_id", parentTaskID, "group_name", groupName) + return false + } + rowsAffected, err := result.RowsAffected() + if err != nil { + logging.LogError(err, "Failed to check group completion function claim result", "parent_task_id", parentTaskID, "group_name", groupName) + return false + } + return rowsAffected > 0 +} + // Helper functions for getting information for sending Task data to a container func CheckAndProcessTaskCompletionHandlers(taskId int) { // check if this task has a completion function @@ -783,8 +839,11 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) { err := database.DB.Get(&task, `SELECT task.parent_task_id, task.operator_id, task.completed, task.subtask_callback_function, task.subtask_callback_function_completed, + task.subtask_callback_function_started, task.group_callback_function, task.group_callback_function_completed, task.completed_callback_function, - task.completed_callback_function_completed, task.subtask_group_name, task.id, task.status, task.eventstepinstance_id + task.group_callback_function_started, + task.completed_callback_function_completed, task.completed_callback_function_started, + task.subtask_group_name, task.id, task.status, task.eventstepinstance_id FROM task WHERE task.id=$1`, taskId) if err != nil { @@ -799,18 +858,23 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) { } if task.ParentTaskID.Valid { err = database.DB.Get(&parentTask, `SELECT - task.id, task.status, task.completed, task.eventstepinstance_id, task.completed_callback_function_completed, - task.subtask_callback_function, - c.script_only "command.script_only" - from task - LEFT OUTER JOIN command c on task.command_id = c.id - WHERE task.id=$1`, task.ParentTaskID.Int64) + task.id, task.status, task.completed, task.eventstepinstance_id, task.completed_callback_function, task.completed_callback_function_completed, + task.completed_callback_function_started, + task.subtask_callback_function, + c.script_only "command.script_only" + from task + LEFT OUTER JOIN command c on task.command_id = c.id + WHERE task.id=$1`, task.ParentTaskID.Int64) if err != nil { logging.LogError(err, "Failed to get parent task information") } } // now we have info for both the current task that finished something and the parent task (if there is one) if task.CompletedCallbackFunction != "" && !task.CompletedCallbackFunctionCompleted { + if !claimTaskCompletionFunctionStart(task.ID, "completed_callback_function_started", "completed_callback_function_completed", + "completed_callback_function", PT_TASK_FUNCTION_STATUS_COMPLETION_FUNCTION) { + return + } // this task has a completion function set, and it hasn't been executed, so run it // ex in create_tasking: completionFunctionName := "shellCompleted" // response.CompletionFunctionName = &completionFunctionName @@ -820,9 +884,7 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) { CompletionFunctionName: task.CompletedCallbackFunction, } task.Status = PT_TASK_FUNCTION_STATUS_COMPLETION_FUNCTION - if _, err := database.DB.NamedExec(`UPDATE task SET status=:status WHERE id=:id`, task); err != nil { - logging.LogError(err, "Failed to update status to completion task running") - } else if err = RabbitMQConnection.SendPtTaskCompletionFunction(taskMessage); err != nil { + if err = RabbitMQConnection.SendPtTaskCompletionFunction(taskMessage); err != nil { logging.LogError(err, "Failed to send task completion function message to container") task.Status = PT_TASK_FUNCTION_STATUS_COMPLETION_FUNCTION_ERROR if _, err = database.DB.NamedExec(`UPDATE task SET status=:status WHERE id=:id`, task); err != nil { @@ -830,6 +892,14 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) { } } } else if task.SubtaskCallbackFunction != "" && !task.SubtaskCallbackFunctionCompleted { + if !task.ParentTaskID.Valid { + logging.LogError(nil, "Task has subtask callback function but no parent task", "task_id", task.ID) + return + } + if !claimTaskCompletionFunctionStart(task.ID, "subtask_callback_function_started", "subtask_callback_function_completed", + "subtask_callback_function", PT_TASK_FUNCTION_STATUS_SUBTASK_COMPLETED_FUNCTION) { + return + } // this task has a subtask callback function, and that subtask function hasn't completed successfully // ex in create_tasking: // completionFunctionName := "pwdCompleted" @@ -854,6 +924,10 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) { } } } else if task.SubtaskGroupName != "" && task.GroupCallbackFunction != "" && !task.GroupCallbackFunctionCompleted { + if !task.ParentTaskID.Valid { + logging.LogError(nil, "Task has group callback function but no parent task", "task_id", task.ID) + return + } // we have a subtask group name, we have a group callback function defined, and that group callback function is done // need to check if we're the last one in the group to finish - if so, we need to call the function, if not do nothing // this function is executed for the PARENT_TASK @@ -872,6 +946,9 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) { task.ParentTaskID.Int64, task.SubtaskGroupName); err != nil { logging.LogError(err, "Failed to fetch other group tasks") } else if len(groupTasks) == 0 { + if !claimGroupCompletionFunctionStart(task.ParentTaskID.Int64, task.SubtaskGroupName, task.GroupCallbackFunction) { + return + } logging.LogInfo("subtask completed, parent task completion function going to run", "taskId", task.ParentTaskID.Int64, "subtaskid", task.ID) err = RabbitMQConnection.SendPtTaskCompletionFunction(taskMessage) if err != nil { @@ -901,6 +978,10 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) { logging.LogError(err, "Failed to update parent task information to submitted") } if parentTask.CompletedCallbackFunction != "" && !parentTask.CompletedCallbackFunctionCompleted { + if !claimTaskCompletionFunctionStart(parentTask.ID, "completed_callback_function_started", "completed_callback_function_completed", + "completed_callback_function", PT_TASK_FUNCTION_STATUS_COMPLETION_FUNCTION) { + return + } // this task has a completion function set, and it hasn't been executed, so run it // ex in create_tasking: completionFunctionName := "shellCompleted" // response.CompletionFunctionName = &completionFunctionName @@ -912,8 +993,8 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) { err = RabbitMQConnection.SendPtTaskCompletionFunction(taskMessage) if err != nil { logging.LogError(err, "Failed to send task completion function message to container") - task.Status = PT_TASK_FUNCTION_STATUS_COMPLETION_FUNCTION_ERROR - if _, err = database.DB.NamedExec(`UPDATE task SET status=:status WHERE id=:id`, task); err != nil { + parentTask.Status = PT_TASK_FUNCTION_STATUS_COMPLETION_FUNCTION_ERROR + if _, err = database.DB.NamedExec(`UPDATE task SET status=:status WHERE id=:id`, parentTask); err != nil { logging.LogError(err, "Failed to update task status for completion function error") } } diff --git a/mythic-docker/src/rabbitmq/utils_proxy_traffic.go b/mythic-docker/src/rabbitmq/utils_proxy_traffic.go index 9e6a43a5..76dfb419 100644 --- a/mythic-docker/src/rabbitmq/utils_proxy_traffic.go +++ b/mythic-docker/src/rabbitmq/utils_proxy_traffic.go @@ -13,6 +13,7 @@ import ( "sort" "strconv" "sync" + "sync/atomic" "time" "github.com/its-a-feature/Mythic/database" @@ -28,6 +29,7 @@ const ( CALLBACK_PORT_TYPE_SOCKS CallbackPortType = "socks" CALLBACK_PORT_TYPE_RPORTFWD = "rpfwd" CALLBACK_PORT_TYPE_INTERACTIVE = "interactive" + callbackPortByteFlushInterval = 20 * time.Second ) type proxyToAgentMessage struct { @@ -53,37 +55,28 @@ type acceptedConnection struct { TaskUUID *string AgentClosedConnection bool } -type bytesSentToAgentMessage struct { - CallbackPortID int `json:"id"` - ByteCount int64 - Initial bool -} -type bytesReceivedFromAgentMessage struct { - CallbackPortID int `json:"id"` - ByteCount int64 - Initial bool -} - type callbackPortUsage struct { - CallbackPortID int `json:"id" db:"id"` - CallbackID int `json:"callback_id" db:"callback_id"` - CallbackDisplayID int `json:"callback_display_id" db:"callback_display_id"` - TaskID int `json:"task_id" db:"task_id"` - LocalPort int `json:"local_port" db:"local_port"` - RemotePort int `json:"remote_port" db:"remote_port"` - RemoteIP string `json:"remote_ip" db:"remote_ip"` - OperationID int `json:"operation_id" db:"operation_id"` - Username string `json:"username" db:"username"` - Password string `json:"password" db:"password"` - PortType CallbackPortType `json:"port_type" db:"port_type"` - listener *net.Listener - bytesReceivedFromAgentChan chan bytesReceivedFromAgentMessage - bytesSentToAgentChan chan bytesSentToAgentMessage - acceptedConnections *[]*acceptedConnection - messagesToAgent chan proxyToAgentMessage - interactiveMessagesToAgent chan agentMessagePostResponseInteractive - newConnectionChannel chan *acceptedConnection - removeConnectionsChannel chan *acceptedConnection + CallbackPortID int `json:"id" db:"id"` + CallbackID int `json:"callback_id" db:"callback_id"` + CallbackDisplayID int `json:"callback_display_id" db:"callback_display_id"` + TaskID int `json:"task_id" db:"task_id"` + LocalPort int `json:"local_port" db:"local_port"` + RemotePort int `json:"remote_port" db:"remote_port"` + RemoteIP string `json:"remote_ip" db:"remote_ip"` + OperationID int `json:"operation_id" db:"operation_id"` + Username string `json:"username" db:"username"` + Password string `json:"password" db:"password"` + PortType CallbackPortType `json:"port_type" db:"port_type"` + listener *net.Listener + bytesReceivedFromAgent atomic.Int64 + bytesSentToAgent atomic.Int64 + lastFlushedBytesReceivedFromAgent atomic.Int64 + lastFlushedBytesSentToAgent atomic.Int64 + acceptedConnections *[]*acceptedConnection + messagesToAgent chan proxyToAgentMessage + interactiveMessagesToAgent chan agentMessagePostResponseInteractive + newConnectionChannel chan *acceptedConnection + removeConnectionsChannel chan *acceptedConnection // messagesFromAgent - these get parsed by manageConnections and passed to the right connection's messagesFromAgent messagesFromAgent chan proxyFromAgentMessage interactiveMessagesFromAgent chan agentMessagePostResponseInteractive @@ -96,8 +89,6 @@ type callbackPortsInUse struct { portsByCallbackIDAndType map[int]map[CallbackPortType][]*callbackPortUsage callbacksWithPorts []int proxyFromAgentMessageChannel chan ProxyFromAgentMessageForMythic - bytesReceivedFromAgentChan chan bytesReceivedFromAgentMessage - bytesSentToAgentChan chan bytesSentToAgentMessage sync.RWMutex } @@ -227,8 +218,6 @@ func (c *callbackPortsInUse) Initialize() { c.ports = make([]*callbackPortUsage, 0) c.resetPortIndexesLocked() c.proxyFromAgentMessageChannel = make(chan ProxyFromAgentMessageForMythic, 2000) - c.bytesReceivedFromAgentChan = make(chan bytesReceivedFromAgentMessage, 2000) - c.bytesSentToAgentChan = make(chan bytesSentToAgentMessage, 2000) c.Unlock() go c.ListenForProxyFromAgentMessage() go c.ListenForNewByteTransferUpdates() @@ -236,7 +225,7 @@ func (c *callbackPortsInUse) Initialize() { callbackport.id, callbackport.callback_id, callbackport.task_id, callbackport.local_port, callbackport.remote_port, callbackport.remote_ip, callbackport.port_type, callbackport.operation_id, callbackport.username, - callbackport.password, + callbackport.password, callbackport.bytes_sent, callbackport.bytes_received, callback.display_id "callback.display_id" FROM callbackport JOIN callback ON callbackport.callback_id=callback.id @@ -244,16 +233,6 @@ func (c *callbackPortsInUse) Initialize() { logging.LogError(err, "Failed to load callback ports from database") } else { for _, proxy := range callbackPorts { - c.bytesSentToAgentChan <- bytesSentToAgentMessage{ - ByteCount: proxy.BytesSent, - CallbackPortID: proxy.ID, - Initial: true, - } - c.bytesReceivedFromAgentChan <- bytesReceivedFromAgentMessage{ - ByteCount: proxy.BytesReceived, - CallbackPortID: proxy.ID, - Initial: true, - } newPort := callbackPortUsage{ CallbackPortID: proxy.ID, CallbackID: proxy.CallbackID, @@ -266,8 +245,6 @@ func (c *callbackPortsInUse) Initialize() { OperationID: proxy.OperationID, Username: proxy.Username, Password: proxy.Password, - bytesSentToAgentChan: c.bytesSentToAgentChan, - bytesReceivedFromAgentChan: c.bytesReceivedFromAgentChan, messagesToAgent: make(chan proxyToAgentMessage, 1000), newConnectionChannel: make(chan *acceptedConnection, 1000), removeConnectionsChannel: make(chan *acceptedConnection, 1000), @@ -276,6 +253,7 @@ func (c *callbackPortsInUse) Initialize() { interactiveMessagesFromAgent: make(chan agentMessagePostResponseInteractive, 1000), stopAllConnections: make(chan bool), } + newPort.initializeByteCounters(proxy.BytesSent, proxy.BytesReceived) acceptedConnections := make([]*acceptedConnection, 0) newPort.acceptedConnections = &acceptedConnections if err := newPort.Start(); err != nil { @@ -374,6 +352,107 @@ func cloneCallbackPortUsages(ports []*callbackPortUsage) []*callbackPortUsage { return append([]*callbackPortUsage(nil), ports...) } +func clampCallbackPortByteCount(value int64) int64 { + if value < 0 { + return 0 + } + if value >= POSTGRES_MAX_BIGINT { + return POSTGRES_MAX_BIGINT - 1 + } + return value +} + +func addCallbackPortByteCounts(value int64, delta int64) int64 { + value = clampCallbackPortByteCount(value) + if delta <= 0 { + return value + } + if delta >= POSTGRES_MAX_BIGINT-value { + return POSTGRES_MAX_BIGINT - 1 + } + return value + delta +} + +func addClampedAtomicInt64(counter *atomic.Int64, delta int64) { + if delta <= 0 { + return + } + for { + current := counter.Load() + next := addCallbackPortByteCounts(current, delta) + if next == current || counter.CompareAndSwap(current, next) { + return + } + } +} + +func (p *callbackPortUsage) initializeByteCounters(bytesSentToAgent int64, bytesReceivedFromAgent int64) { + bytesSentToAgent = clampCallbackPortByteCount(bytesSentToAgent) + bytesReceivedFromAgent = clampCallbackPortByteCount(bytesReceivedFromAgent) + p.bytesSentToAgent.Store(bytesSentToAgent) + p.lastFlushedBytesSentToAgent.Store(bytesSentToAgent) + p.bytesReceivedFromAgent.Store(bytesReceivedFromAgent) + p.lastFlushedBytesReceivedFromAgent.Store(bytesReceivedFromAgent) +} + +func (p *callbackPortUsage) applyPersistedByteCounters(bytesSentToAgent int64, bytesReceivedFromAgent int64) { + applyPersistedByteCounter(&p.bytesSentToAgent, &p.lastFlushedBytesSentToAgent, bytesSentToAgent) + applyPersistedByteCounter(&p.bytesReceivedFromAgent, &p.lastFlushedBytesReceivedFromAgent, bytesReceivedFromAgent) +} + +func applyPersistedByteCounter(counter *atomic.Int64, lastFlushed *atomic.Int64, persistedValue int64) { + persistedValue = clampCallbackPortByteCount(persistedValue) + for { + currentDelta := counter.Load() + next := addCallbackPortByteCounts(persistedValue, currentDelta) + if counter.CompareAndSwap(currentDelta, next) { + lastFlushed.Store(persistedValue) + return + } + } +} + +func (p *callbackPortUsage) addBytesSentToAgent(byteCount int64) { + addClampedAtomicInt64(&p.bytesSentToAgent, byteCount) +} + +func (p *callbackPortUsage) addBytesReceivedFromAgent(byteCount int64) { + addClampedAtomicInt64(&p.bytesReceivedFromAgent, byteCount) +} + +func (p *callbackPortUsage) flushByteCounters() { + if p.CallbackPortID <= 0 { + return + } + currentBytesSentToAgent := p.bytesSentToAgent.Load() + bytesSentToAgent := clampCallbackPortByteCount(currentBytesSentToAgent) + if bytesSentToAgent != currentBytesSentToAgent { + p.bytesSentToAgent.CompareAndSwap(currentBytesSentToAgent, bytesSentToAgent) + } + if bytesSentToAgent != p.lastFlushedBytesSentToAgent.Load() && + updateCallbackPortStats("bytes_sent", bytesSentToAgent, p.CallbackPortID) { + p.lastFlushedBytesSentToAgent.Store(bytesSentToAgent) + } + currentBytesReceivedFromAgent := p.bytesReceivedFromAgent.Load() + bytesReceivedFromAgent := clampCallbackPortByteCount(currentBytesReceivedFromAgent) + if bytesReceivedFromAgent != currentBytesReceivedFromAgent { + p.bytesReceivedFromAgent.CompareAndSwap(currentBytesReceivedFromAgent, bytesReceivedFromAgent) + } + if bytesReceivedFromAgent != p.lastFlushedBytesReceivedFromAgent.Load() && + updateCallbackPortStats("bytes_received", bytesReceivedFromAgent, p.CallbackPortID) { + p.lastFlushedBytesReceivedFromAgent.Store(bytesReceivedFromAgent) + } +} + +func (c *callbackPortsInUse) flushChangedByteCounters() { + c.RLock() + ports := cloneCallbackPortUsages(c.ports) + c.RUnlock() + for _, port := range ports { + port.flushByteCounters() + } +} + func (c *callbackPortsInUse) getPortsForCallbackAndType(callbackID int, portType CallbackPortType) []*callbackPortUsage { c.RLock() defer c.RUnlock() @@ -431,83 +510,28 @@ func (c *callbackPortsInUse) findPortForRemoval(callbackId int, portType Callbac return nil } func (c *callbackPortsInUse) ListenForNewByteTransferUpdates() { - oldByteValues := make(map[int]map[string]int64) - currentByteValues := make(map[int]map[string]int64) - updateValuesChan := make(chan bool) - go func() { - // every 30s send a signal to make sure we update our current byte values if they've changed - for { - select { - case <-time.After(20 * time.Second): - updateValuesChan <- true - } - } - }() + ticker := time.NewTicker(callbackPortByteFlushInterval) + defer ticker.Stop() defer func() { logging.LogError(nil, "no longer listening for new byte transfer updates") }() - for { - select { - case <-updateValuesChan: - for callbackPortID, _ := range currentByteValues { - if currentByteValues[callbackPortID]["sent"] != oldByteValues[callbackPortID]["sent"] { - go updateCallbackPortStats("bytes_sent", currentByteValues[callbackPortID]["sent"], callbackPortID) - oldByteValues[callbackPortID]["sent"] = currentByteValues[callbackPortID]["sent"] - } - if currentByteValues[callbackPortID]["received"] != oldByteValues[callbackPortID]["received"] { - go updateCallbackPortStats("bytes_received", currentByteValues[callbackPortID]["received"], callbackPortID) - oldByteValues[callbackPortID]["received"] = currentByteValues[callbackPortID]["received"] - } - } - - case bytesFromAgentMsg := <-c.bytesReceivedFromAgentChan: - if _, ok := currentByteValues[bytesFromAgentMsg.CallbackPortID]; !ok { - currentByteValues[bytesFromAgentMsg.CallbackPortID] = map[string]int64{ - "received": 0, - "sent": 0, - } - oldByteValues[bytesFromAgentMsg.CallbackPortID] = map[string]int64{ - "received": 0, - "sent": 0, - } - } - if bytesFromAgentMsg.Initial { - oldByteValues[bytesFromAgentMsg.CallbackPortID]["received"] = bytesFromAgentMsg.ByteCount - currentByteValues[bytesFromAgentMsg.CallbackPortID]["received"] = bytesFromAgentMsg.ByteCount - } else { - currentByteValues[bytesFromAgentMsg.CallbackPortID]["received"] += bytesFromAgentMsg.ByteCount - } - case bytesSentToAgentMsg := <-c.bytesSentToAgentChan: - if _, ok := currentByteValues[bytesSentToAgentMsg.CallbackPortID]; !ok { - currentByteValues[bytesSentToAgentMsg.CallbackPortID] = map[string]int64{ - "received": 0, - "sent": 0, - } - oldByteValues[bytesSentToAgentMsg.CallbackPortID] = map[string]int64{ - "received": 0, - "sent": 0, - } - } - if bytesSentToAgentMsg.Initial { - oldByteValues[bytesSentToAgentMsg.CallbackPortID]["sent"] = bytesSentToAgentMsg.ByteCount - currentByteValues[bytesSentToAgentMsg.CallbackPortID]["sent"] = bytesSentToAgentMsg.ByteCount - } else { - currentByteValues[bytesSentToAgentMsg.CallbackPortID]["sent"] += bytesSentToAgentMsg.ByteCount - } - } + for range ticker.C { + c.flushChangedByteCounters() } } -func updateCallbackPortStats(field string, value int64, callbackPortID int) { - updatedValue := value - if updatedValue > POSTGRES_MAX_BIGINT { - updatedValue = POSTGRES_MAX_BIGINT - 1 +func updateCallbackPortStats(field string, value int64, callbackPortID int) bool { + if callbackPortID <= 0 { + return false } + updatedValue := clampCallbackPortByteCount(value) _, err := database.DB.Exec(fmt.Sprintf("UPDATE callbackport SET %s=$1 WHERE id=$2", - field), value, callbackPortID) + field), updatedValue, callbackPortID) if err != nil { logging.LogError(err, "Failed to update callback port stats") + return false } + return true } func (c *callbackPortsInUse) ListenForProxyFromAgentMessage() { for { @@ -712,18 +736,6 @@ func (c *callbackPortsInUse) GetDataForCallbackId(callbackId int, portType strin func (c *callbackPortsInUse) Add(callbackId int, portType CallbackPortType, localPort int, remotePort int, remoteIP string, taskId int, operationId int, bytesSentToAgent int64, bytesReceivedFromAgent int64, callbackPortID int, username string, password string) error { - if callbackPortID > 0 { - c.bytesSentToAgentChan <- bytesSentToAgentMessage{ - ByteCount: bytesSentToAgent, - CallbackPortID: callbackPortID, - Initial: true, - } - c.bytesReceivedFromAgentChan <- bytesReceivedFromAgentMessage{ - ByteCount: bytesReceivedFromAgent, - CallbackPortID: callbackPortID, - Initial: true, - } - } newPort := callbackPortUsage{ CallbackPortID: callbackPortID, CallbackID: callbackId, @@ -735,8 +747,6 @@ func (c *callbackPortsInUse) Add(callbackId int, portType CallbackPortType, loca PortType: portType, Username: username, Password: password, - bytesSentToAgentChan: c.bytesSentToAgentChan, - bytesReceivedFromAgentChan: c.bytesReceivedFromAgentChan, messagesToAgent: make(chan proxyToAgentMessage, 1000), newConnectionChannel: make(chan *acceptedConnection, 1000), removeConnectionsChannel: make(chan *acceptedConnection, 1000), @@ -745,6 +755,11 @@ func (c *callbackPortsInUse) Add(callbackId int, portType CallbackPortType, loca interactiveMessagesFromAgent: make(chan agentMessagePostResponseInteractive, 1000), stopAllConnections: make(chan bool), } + if callbackPortID > 0 { + newPort.initializeByteCounters(bytesSentToAgent, bytesReceivedFromAgent) + } else { + newPort.initializeByteCounters(0, 0) + } acceptedConnections := make([]*acceptedConnection, 0) newPort.acceptedConnections = &acceptedConnections callbackPort := databaseStructs.Callbackport{} @@ -759,7 +774,7 @@ func (c *callbackPortsInUse) Add(callbackId int, portType CallbackPortType, loca logging.LogError(err, "Failed to start new proxy port") return err } - err = database.DB.Get(&callbackPort, `SELECT id FROM callbackport WHERE + err = database.DB.Get(&callbackPort, `SELECT id, bytes_sent, bytes_received FROM callbackport WHERE operation_id=$1 AND callback_id=$2 AND local_port=$3 AND port_type=$4 AND remote_ip=$5 AND remote_port=$6 AND username=$7 AND password=$8`, operationId, callbackId, localPort, portType, remoteIP, remotePort, username, password) @@ -775,7 +790,7 @@ func (c *callbackPortsInUse) Add(callbackId int, portType CallbackPortType, loca } return err } - err = statement.Get(&newPort.CallbackPortID, newPort) + err = statement.Get(&newPort.CallbackPortID, &newPort) if err != nil { logging.LogError(err, "Failed to insert new callback port") if err := newPort.Stop(); err != nil { @@ -786,6 +801,9 @@ func (c *callbackPortsInUse) Add(callbackId int, portType CallbackPortType, loca } else if err == nil { _, err = database.DB.NamedExec(`UPDATE callbackport SET deleted=false WHERE id=:id`, callbackPort) newPort.CallbackPortID = callbackPort.ID + if callbackPortID <= 0 { + newPort.applyPersistedByteCounters(callbackPort.BytesSent, callbackPort.BytesReceived) + } } else { logging.LogError(err, "Failed to create new callback port mapping") if err := newPort.Stop(); err != nil { @@ -829,6 +847,7 @@ func (c *callbackPortsInUse) Remove(callbackId int, portType CallbackPortType, l logging.LogError(err, "Failed to stop proxy") return err } + port.flushByteCounters() queryString := `UPDATE callbackport SET deleted=true WHERE id=$1` queryArgs := []interface{}{port.CallbackPortID} _, err = database.DB.Exec(queryString, queryArgs...) @@ -1412,12 +1431,7 @@ func (p *callbackPortUsage) readFromAgentSocksConn(newConnection *acceptedConnec p.removeConnectionsChannel <- newConnection return } - // non-blocking send stats update - go func(byteCount int64, callbackPortID int) { - select { - case p.bytesReceivedFromAgentChan <- bytesReceivedFromAgentMessage{ByteCount: byteCount, CallbackPortID: callbackPortID, Initial: false}: - } - }(int64(len(dataBytes)), p.CallbackPortID) + p.addBytesReceivedFromAgent(int64(len(dataBytes))) if agentMessage.IsExit { logging.LogDebug("got message from agent isExit", "server_id", newConnection.ServerID) @@ -1746,12 +1760,7 @@ func (p *callbackPortUsage) handleSocksConnections() { CallbackID: p.CallbackID, ProxyType: p.PortType, } - // non-blocking send stats update - go func(byteCount int64, callbackPortID int) { - select { - case p.bytesSentToAgentChan <- bytesSentToAgentMessage{ByteCount: byteCount, CallbackPortID: callbackPortID, Initial: false}: - } - }(int64(newUDPConnectionBufLength), p.CallbackPortID) + p.addBytesSentToAgent(int64(newUDPConnectionBufLength)) go func(remoteAddr net.Addr) { // work on this section!! // handle the read from agent and write back out to socket here so we can use the same addr @@ -1774,12 +1783,7 @@ func (p *callbackPortUsage) handleSocksConnections() { p.removeConnectionsChannel <- &newUDPConnection return } - // non-blocking send stats update - go func(byteCount int64, callbackPortID int) { - select { - case p.bytesReceivedFromAgentChan <- bytesReceivedFromAgentMessage{ByteCount: byteCount, CallbackPortID: callbackPortID, Initial: false}: - } - }(int64(len(dataBytes)), p.CallbackPortID) + p.addBytesReceivedFromAgent(int64(len(dataBytes))) if agentMessage.IsExit { //logging.LogDebug("got message from agent isExit", "server_id", newConnection.ServerID) @@ -1821,12 +1825,7 @@ func (p *callbackPortUsage) handleSocksConnections() { CallbackID: p.CallbackID, ProxyType: p.PortType, } - // non-blocking send stats update - go func(byteCount int64, callbackPortID int) { - select { - case p.bytesSentToAgentChan <- bytesSentToAgentMessage{ByteCount: byteCount, CallbackPortID: callbackPortID, Initial: false}: - } - }(int64(length), p.CallbackPortID) + p.addBytesSentToAgent(int64(length)) } if err != nil { logging.LogError(err, "closing tcp connection") @@ -1866,11 +1865,7 @@ func (p *callbackPortUsage) handleRpfwdConnections(newConnection *acceptedConnec p.removeConnectionsChannel <- newConnection return } - // non-blocking send stats update - select { - case p.bytesReceivedFromAgentChan <- bytesReceivedFromAgentMessage{ByteCount: int64(len(dataBytes)), CallbackPortID: p.CallbackPortID}: - default: - } + p.addBytesReceivedFromAgent(int64(len(dataBytes))) if agentMessage.IsExit { //logging.LogDebug("got message isExit", "server_id", newConnection.ServerID) // cleanup the connection data, but don't send an exit back to the agent @@ -1906,11 +1901,7 @@ func (p *callbackPortUsage) handleRpfwdConnections(newConnection *acceptedConnec ProxyType: p.PortType, CallbackID: p.CallbackID, } - // non-blocking send stats update - select { - case p.bytesSentToAgentChan <- bytesSentToAgentMessage{ByteCount: int64(length), CallbackPortID: p.CallbackPortID}: - default: - } + p.addBytesSentToAgent(int64(length)) //fmt.Printf("Message sent to p.messagesToAgent channel for chan %d\n", newConnection.ServerID) } if err != nil { @@ -1978,11 +1969,7 @@ func (p *callbackPortUsage) handleInteractiveConnections() { p.removeConnectionsChannel <- &newConnection return } - // non-blocking send stats update - select { - case p.bytesReceivedFromAgentChan <- bytesReceivedFromAgentMessage{ByteCount: int64(len(dataBytes)), CallbackPortID: p.CallbackPortID}: - default: - } + p.addBytesReceivedFromAgent(int64(len(dataBytes))) if agentMessage.MessageType == InteractiveTask.Exit { //logging.LogDebug("got message isExit", "server_id", newConnection.ServerID) // cleanup the connection data @@ -2018,11 +2005,7 @@ func (p *callbackPortUsage) handleInteractiveConnections() { CallbackID: p.CallbackID, ProxyType: p.PortType, } - // non-blocking send stats update - select { - case p.bytesSentToAgentChan <- bytesSentToAgentMessage{ByteCount: int64(length), CallbackPortID: p.CallbackPortID}: - default: - } + p.addBytesSentToAgent(int64(length)) //fmt.Printf("Message sent to p.messagesToAgent channel for chan %d\n", newConnection.ServerID) } diff --git a/mythic-docker/src/rabbitmq/utils_proxy_traffic_test.go b/mythic-docker/src/rabbitmq/utils_proxy_traffic_test.go index acc28df1..c40b1cb2 100644 --- a/mythic-docker/src/rabbitmq/utils_proxy_traffic_test.go +++ b/mythic-docker/src/rabbitmq/utils_proxy_traffic_test.go @@ -2,6 +2,7 @@ package rabbitmq import ( "slices" + "sync" "testing" ) @@ -92,3 +93,90 @@ func TestCallbackPortsInUseCombinedDataFetchUsesIndexedPorts(t *testing.T) { t.Fatalf("expected proxy data to be drained, got %#v", proxyData) } } + +func TestCallbackPortUsageByteCountersAccumulateWithoutChannels(t *testing.T) { + port := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 1) + port.initializeByteCounters(10, 20) + + port.addBytesSentToAgent(5) + port.addBytesReceivedFromAgent(7) + + if got := port.bytesSentToAgent.Load(); got != 15 { + t.Fatalf("expected bytes sent to agent to be 15, got %d", got) + } + if got := port.bytesReceivedFromAgent.Load(); got != 27 { + t.Fatalf("expected bytes received from agent to be 27, got %d", got) + } + if got := port.lastFlushedBytesSentToAgent.Load(); got != 10 { + t.Fatalf("expected last flushed bytes sent to agent to remain 10, got %d", got) + } + if got := port.lastFlushedBytesReceivedFromAgent.Load(); got != 20 { + t.Fatalf("expected last flushed bytes received from agent to remain 20, got %d", got) + } +} + +func TestCallbackPortUsageByteCountersApplyPersistedBaseline(t *testing.T) { + port := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 1) + port.initializeByteCounters(0, 0) + port.addBytesSentToAgent(3) + port.addBytesReceivedFromAgent(4) + + port.applyPersistedByteCounters(100, 200) + + if got := port.bytesSentToAgent.Load(); got != 103 { + t.Fatalf("expected bytes sent to agent to include persisted baseline and pending delta, got %d", got) + } + if got := port.bytesReceivedFromAgent.Load(); got != 204 { + t.Fatalf("expected bytes received from agent to include persisted baseline and pending delta, got %d", got) + } + if got := port.lastFlushedBytesSentToAgent.Load(); got != 100 { + t.Fatalf("expected last flushed bytes sent to agent to track persisted baseline, got %d", got) + } + if got := port.lastFlushedBytesReceivedFromAgent.Load(); got != 200 { + t.Fatalf("expected last flushed bytes received from agent to track persisted baseline, got %d", got) + } +} + +func TestCallbackPortUsageByteCountersClampAtPostgresBigint(t *testing.T) { + port := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 1) + port.initializeByteCounters(POSTGRES_MAX_BIGINT-2, POSTGRES_MAX_BIGINT-3) + + port.addBytesSentToAgent(10) + port.addBytesReceivedFromAgent(10) + + if got := port.bytesSentToAgent.Load(); got != POSTGRES_MAX_BIGINT-1 { + t.Fatalf("expected bytes sent to agent to clamp to %d, got %d", POSTGRES_MAX_BIGINT-1, got) + } + if got := port.bytesReceivedFromAgent.Load(); got != POSTGRES_MAX_BIGINT-1 { + t.Fatalf("expected bytes received from agent to clamp to %d, got %d", POSTGRES_MAX_BIGINT-1, got) + } +} + +func TestCallbackPortUsageByteCountersHandleConcurrentAdds(t *testing.T) { + port := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 1) + port.initializeByteCounters(5, 7) + + const goroutineCount = 32 + const incrementsPerGoroutine = 1000 + var wg sync.WaitGroup + wg.Add(goroutineCount) + for i := 0; i < goroutineCount; i++ { + go func() { + defer wg.Done() + for j := 0; j < incrementsPerGoroutine; j++ { + port.addBytesSentToAgent(1) + port.addBytesReceivedFromAgent(2) + } + }() + } + wg.Wait() + + expectedSent := int64(5 + goroutineCount*incrementsPerGoroutine) + expectedReceived := int64(7 + goroutineCount*incrementsPerGoroutine*2) + if got := port.bytesSentToAgent.Load(); got != expectedSent { + t.Fatalf("expected bytes sent to agent to be %d, got %d", expectedSent, got) + } + if got := port.bytesReceivedFromAgent.Load(); got != expectedReceived { + t.Fatalf("expected bytes received from agent to be %d, got %d", expectedReceived, got) + } +}