admin/its-a-feature-Mythic
1
0
Fork 0
mirror of https://github.com/its-a-feature/Mythic synced 2026-06-08 14:55:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

added file browser insert/delete performance updates

Browse Source
This commit is contained in:
its-a-feature
2026-05-10 13:52:31 -05:00
parent 468decd0bc
commit 4c2f5e3756
4 changed files with 864 additions and 440 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mythic-docker/src/database/migrations/003003020_3.3.20.sql
+16
View File
@@ -76,6 +76,18 @@ 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;
create index if not exists mythictree_operation_tree_host_parent_callback_idx
on "public"."mythictree" using btree (operation_id, tree_type, host, parent_path, callback_id);
create index if not exists mythictree_operation_tree_host_full_callback_idx
on "public"."mythictree" using btree (operation_id, tree_type, host, full_path, callback_id);
create index if not exists mythictree_operation_tree_timestamp_idx
on "public"."mythictree" using btree (operation_id, tree_type, "timestamp");
create index if not exists mythictree_tree_deleted_timestamp_idx
on "public"."mythictree" using btree (tree_type, deleted, "timestamp");
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,
@@ -206,6 +218,10 @@ $$;
drop index if exists "public"."callback_operation_display_id_unique";
drop index if exists "public"."task_operation_display_id_unique";
drop index if exists "public"."response_task_sequence_number_unique";
drop index if exists "public"."mythictree_tree_deleted_timestamp_idx";
drop index if exists "public"."mythictree_operation_tree_timestamp_idx";
drop index if exists "public"."mythictree_operation_tree_host_full_callback_idx";
drop index if exists "public"."mythictree_operation_tree_host_parent_callback_idx";
alter table "public"."task"
drop column if exists completed_callback_function_started,
drop column if exists group_callback_function_started,
mythic-docker/src/rabbitmq/util_agent_message_actions_post_response.go
+113 -440
View File
@@ -624,7 +624,7 @@ func handleAgentMessagePostResponse(incoming *map[string]interface{}, uUIDInfo *
currentTask.Stdout += *agentMessage.Responses[i].Stdout
}
if agentMessage.Responses[i].Stderr != nil {
currentTask.Stderr = *agentMessage.Responses[i].Stderr
currentTask.Stderr += *agentMessage.Responses[i].Stderr
}
if agentMessage.Responses[i].FileBrowser != nil {
// do it in the background - the agent doesn't need the result of this directly
@@ -1858,7 +1858,9 @@ func associateFileMetaWithMythicTree(pathData utils.AnalyzedPath, fileMeta datab
newTree.CallbackID.Int64 = int64(task.Callback.ID)
// remove this filename from parent creation
pathData.PathPieces = pathData.PathPieces[:len(pathData.PathPieces)-1]
resolveAndCreateParentPathsForTreeNode(pathData, utils.AnalyzedPath{}, task, databaseStructs.TREE_TYPE_FILE)
if err := upsertMythicTreeNodes(buildMythicTreeParentNodes(pathData, utils.AnalyzedPath{}, task, databaseStructs.TREE_TYPE_FILE)); err != nil {
logging.LogError(err, "Failed to create file browser parent paths")
}
createTreeNode(&newTree)
if newTree.ID == 0 {
logging.LogError(nil, "Failed to create new tree entry")
@@ -1986,7 +1988,7 @@ func HandleAgentMessagePostResponseFileBrowser(task databaseStructs.Task, fileBr
if fileBrowser.Host != "" {
pathData.Host = strings.ToUpper(fileBrowser.Host)
}
resolveAndCreateParentPathsForTreeNode(pathData, utils.AnalyzedPath{}, task, databaseStructs.TREE_TYPE_FILE)
treeNodesToUpsert := buildMythicTreeParentNodes(pathData, utils.AnalyzedPath{}, task, databaseStructs.TREE_TYPE_FILE)
// now that the parents and all ancestors are resolved, process the current path and all children
realParentPath := strings.Join(pathData.PathPieces, pathData.PathSeparator)
// check for the instance of // as a leading path
@@ -2033,7 +2035,17 @@ func HandleAgentMessagePostResponseFileBrowser(task databaseStructs.Task, fileBr
newTree.APITokensID.Valid = true
newTree.APITokensID.Int64 = int64(apitokensId)
}
createTreeNode(&newTree)
treeNodesToUpsert = append(treeNodesToUpsert, &newTree)
if fileBrowser.Files != nil {
// Always upsert reported children immediately so the first listing is
// visible without waiting for update_deleted cleanup to flush.
for _, newEntry := range *fileBrowser.Files {
treeNodesToUpsert = append(treeNodesToUpsert, buildFileBrowserChildMythicTreeNode(task, pathData.Host, fullPath, pathData.PathSeparator, newEntry, newTree.Os, apitokensId))
}
}
if err := upsertMythicTreeNodes(treeNodesToUpsert); err != nil {
logging.LogError(err, "Failed to create file browser MythicTree entries")
}
//logging.LogInfo("checking update deleted", "update_deleted", fileBrowser.UpdateDeleted)
if fileBrowser.UpdateDeleted != nil && *fileBrowser.UpdateDeleted {
fileBrowserUpdateDeletedChannel <- FileBrowserChannelMessage{
@@ -2041,32 +2053,6 @@ func HandleAgentMessagePostResponseFileBrowser(task databaseStructs.Task, fileBr
NewFileBrowserData: fileBrowser,
APITokenID: apitokensId,
}
} else if fileBrowser.Files != nil {
// we're not automatically updating deleted children, so just iterate over the files and insert/update them
for _, newEntry := range *fileBrowser.Files {
newTreeChild := databaseStructs.MythicTree{
Host: pathData.Host,
TaskID: task.ID,
OperationID: task.OperationID,
Name: []byte(newEntry.Name),
ParentPath: fullPath,
TreeType: databaseStructs.TREE_TYPE_FILE,
CanHaveChildren: !newEntry.IsFile,
Deleted: false,
Os: newTree.Os,
DisplayPath: []byte{},
}
newTreeChild.FullPath = treeNodeGetFullPath(fullPath, []byte(newEntry.Name), []byte(pathData.PathSeparator), databaseStructs.TREE_TYPE_FILE)
fileMetaData = addChildFilePermissions(&newEntry)
newTreeChild.Metadata = GetMythicJSONTextFromStruct(fileMetaData)
newTreeChild.CallbackID.Valid = true
newTreeChild.CallbackID.Int64 = int64(task.Callback.ID)
if apitokensId > 0 {
newTree.APITokensID.Valid = true
newTree.APITokensID.Int64 = int64(apitokensId)
}
createTreeNode(&newTreeChild)
}
}
return nil
}
@@ -2094,120 +2080,13 @@ func listenForFileBrowserUpdateDeleted() {
// only process it all once this is nil
continue
}
// check if we actually need to do anything, if this was kicked off via task completion
//logging.LogInfo("processing update deleted data now", "task", updateMessage.Task.Completed, "cache data", len(fileBrowserUpdateDeletedCache[updateMessage.Task.ID]))
// now process all of it
// we need to iterate over the children for this entry and potentially remove any that the database know of but that aren't in our `files` list
pathData, err := utils.SplitFilePathGetHost(fileBrowserUpdateDeletedCache[updateMessage.Task.ID][0].ParentPath, fileBrowserUpdateDeletedCache[updateMessage.Task.ID][0].Name, []string{})
groups, err := groupFileBrowserUpdateDeletedEntries(updateMessage.Task, fileBrowserUpdateDeletedCache[updateMessage.Task.ID])
if err != nil {
logging.LogError(err, "failed to split path to get path data for updating paths")
delete(fileBrowserUpdateDeletedCache, updateMessage.Task.ID)
continue
}
if pathData.Host == "" {
pathData.Host = strings.ToUpper(updateMessage.Task.Callback.Host)
}
if fileBrowserUpdateDeletedCache[updateMessage.Task.ID][0].Host != "" {
pathData.Host = strings.ToUpper(fileBrowserUpdateDeletedCache[updateMessage.Task.ID][0].Host)
}
// now that the parents and all ancestors are resolved, process the current path and all children
realParentPath := strings.Join(pathData.PathPieces, pathData.PathSeparator)
// check for the instance of // as a leading path
if len(realParentPath) > 2 {
if realParentPath[0] == '/' && realParentPath[1] == '/' {
realParentPath = realParentPath[1:]
}
}
if fileBrowserUpdateDeletedCache[updateMessage.Task.ID][0].Name == "" {
logging.LogError(nil, "Can't create file browser entry with empty name")
delete(fileBrowserUpdateDeletedCache, updateMessage.Task.ID)
continue
}
fullPath := treeNodeGetFullPath(
[]byte(realParentPath),
[]byte(fileBrowserUpdateDeletedCache[updateMessage.Task.ID][0].Name),
[]byte(pathData.PathSeparator),
databaseStructs.TREE_TYPE_FILE)
var existingTreeEntries []databaseStructs.MythicTree
err = database.DB.Select(&existingTreeEntries, `SELECT
id, "name", success, full_path, parent_path, operation_id, host, tree_type, callback_id
FROM mythictree WHERE
parent_path=$1 AND operation_id=$2 AND host=$3 AND tree_type=$4`,
fullPath, updateMessage.Task.OperationID, pathData.Host, databaseStructs.TREE_TYPE_FILE)
if err != nil {
logging.LogError(err, "Failed to fetch existing children")
delete(fileBrowserUpdateDeletedCache, updateMessage.Task.ID)
continue
}
var namesToDeleteAndUpdate []string // will get existing database IDs for things that aren't in the files list
for _, existingEntry := range existingTreeEntries {
existingEntryStillExists := false
for _, fileBrowser := range fileBrowserUpdateDeletedCache[updateMessage.Task.ID] {
//logging.LogInfo("trying to loop through cached file browser entry", "cached entry", fileBrowser)
if fileBrowser.Files != nil {
//logging.LogInfo("checking for file existing", "name", existingEntry.Name)
for _, newEntry := range *fileBrowser.Files {
//logging.LogInfo("checking for existing file", "existing", existingEntry.Name, "reported", newEntry.Name)
if bytes.Equal([]byte(newEntry.Name), existingEntry.Name) {
//logging.LogInfo("[+] found a match in existing data and new data", "name", newEntry.Name, "id", existingEntry.ID, "callback", existingEntry.CallbackID)
namesToDeleteAndUpdate = append(namesToDeleteAndUpdate, newEntry.Name)
existingEntryStillExists = true
// update the entry in the database
newTreeChild := databaseStructs.MythicTree{
Host: pathData.Host,
TaskID: updateMessage.Task.ID,
OperationID: updateMessage.Task.OperationID,
Name: []byte(newEntry.Name),
ParentPath: existingEntry.ParentPath,
FullPath: existingEntry.FullPath,
TreeType: databaseStructs.TREE_TYPE_FILE,
CanHaveChildren: !newEntry.IsFile,
Deleted: false,
Success: existingEntry.Success,
ID: existingEntry.ID,
}
fileMetaData := addChildFilePermissions(&newEntry)
newTreeChild.Metadata = GetMythicJSONTextFromStruct(fileMetaData)
newTreeChild.CallbackID.Valid = true
newTreeChild.CallbackID.Int64 = int64(updateMessage.Task.Callback.ID)
updateTreeNode(newTreeChild)
}
}
}
}
if !existingEntryStillExists {
//logging.LogError(nil, "failed to find match, marking as deleted", "name", existingEntry.Name)
namesToDeleteAndUpdate = append(namesToDeleteAndUpdate, string(existingEntry.Name))
existingEntry.Deleted = true
existingEntry.TaskID = updateMessage.Task.ID
deleteTreeNode(existingEntry, true)
}
}
// now all existing ones have been updated or deleted, so it's time to add new ones
for _, fileBrowser := range fileBrowserUpdateDeletedCache[updateMessage.Task.ID] {
if fileBrowser.Files != nil {
for _, newEntry := range *fileBrowser.Files {
if !utils.SliceContains(namesToDeleteAndUpdate, newEntry.Name) {
// this isn't marked as updated or deleted, so let's create it
newTreeChild := databaseStructs.MythicTree{
Host: pathData.Host,
TaskID: updateMessage.Task.ID,
OperationID: updateMessage.Task.OperationID,
Name: []byte(newEntry.Name),
ParentPath: fullPath,
TreeType: databaseStructs.TREE_TYPE_FILE,
CanHaveChildren: !newEntry.IsFile,
Deleted: false,
DisplayPath: []byte{},
}
newTreeChild.FullPath = treeNodeGetFullPath(fullPath, []byte(newEntry.Name), []byte(pathData.PathSeparator), databaseStructs.TREE_TYPE_FILE)
fileMetaData := addChildFilePermissions(&newEntry)
newTreeChild.Metadata = GetMythicJSONTextFromStruct(fileMetaData)
newTreeChild.CallbackID.Valid = true
newTreeChild.CallbackID.Int64 = int64(updateMessage.Task.Callback.ID)
createTreeNode(&newTreeChild)
}
logging.LogError(err, "Failed to group file browser update_deleted entries")
} else {
for _, group := range groups {
if err := reconcileFileBrowserUpdateDeletedGroup(updateMessage.Task, group, updateMessage.APITokenID); err != nil {
logging.LogError(err, "Failed to update file browser children")
}
}
}
@@ -2249,139 +2128,34 @@ func HandleAgentMessagePostResponseProcesses(task databaseStructs.Task, processe
logging.LogError(err, "Failed to fetch existing children")
return err
}
// we have all the current processes for that host, track which ones need updating/deleting
var namesToDeleteAndUpdate []string // will get existing database IDs for things that aren't in the files list
incomingProcessesByMatchKey := make(map[string]agentMessagePostResponseProcesses, len(*processes))
for _, newEntry := range *processes {
incomingProcessesByMatchKey[getProcessMatchKey(newEntry)] = newEntry
}
namesToDeleteAndUpdate := make(map[string]struct{}, len(existingTreeEntries)+len(*processes))
treeNodesToUpsert := make([]*databaseStructs.MythicTree, 0, len(*processes))
for _, existingEntry := range existingTreeEntries {
existingEntryStillExists := false
for _, newEntry := range *processes {
if newEntry.Name == "" {
newEntry.Name = "unknown"
}
parentPath := strconv.Itoa(newEntry.ParentProcessID)
if newEntry.ParentProcessID <= 0 {
parentPath = ""
}
if strconv.Itoa(newEntry.ProcessID) == string(existingEntry.FullPath) &&
newEntry.Name == string(existingEntry.Name) &&
parentPath == string(existingEntry.ParentPath) {
namesToDeleteAndUpdate = append(namesToDeleteAndUpdate, strconv.Itoa(newEntry.ProcessID))
existingEntryStillExists = true
// update the entry in the database
fullPath := treeNodeGetFullPath(
[]byte(parentPath),
[]byte(strconv.Itoa(newEntry.ProcessID)),
[]byte("/"),
databaseStructs.TREE_TYPE_PROCESS)
newTree := databaseStructs.MythicTree{
Host: host,
TaskID: task.ID,
OperationID: task.OperationID,
Name: []byte(newEntry.Name),
ParentPath: []byte(parentPath),
FullPath: fullPath,
TreeType: databaseStructs.TREE_TYPE_PROCESS,
CanHaveChildren: true,
Deleted: false,
DisplayPath: []byte{},
}
if newEntry.OS != nil {
newTree.Os = *newEntry.OS
} else {
newTree.Os = task.Callback.Payload.Os
}
metadata := map[string]interface{}{
"process_id": newEntry.ProcessID,
"parent_process_id": newEntry.ParentProcessID,
"architecture": newEntry.Architecture,
"bin_path": newEntry.BinPath,
"name": newEntry.Name,
"user": newEntry.User,
"command_line": newEntry.CommandLine,
"integrity_level": newEntry.IntegrityLevel,
"start_time": newEntry.StartTime,
"description": newEntry.Description,
"signer": newEntry.Signer,
"protected_process_level": newEntry.ProtectionProcessLevel,
}
reflectBackOtherKeys(&metadata, &newEntry.Other)
newTree.Metadata = GetMythicJSONTextFromStruct(metadata)
newTree.CallbackID.Valid = true
newTree.CallbackID.Int64 = int64(task.Callback.ID)
if apitokensId > 0 {
newTree.APITokensID.Valid = true
newTree.APITokensID.Int64 = int64(apitokensId)
}
createTreeNode(&newTree)
}
}
if !existingEntryStillExists {
if newEntry, ok := incomingProcessesByMatchKey[getExistingProcessMatchKey(existingEntry)]; ok {
namesToDeleteAndUpdate[strconv.Itoa(newEntry.ProcessID)] = struct{}{}
treeNodesToUpsert = append(treeNodesToUpsert, buildProcessMythicTreeNode(task, host, newEntry, apitokensId))
} else {
// full path is just the string of the PID
namesToDeleteAndUpdate = append(namesToDeleteAndUpdate, string(existingEntry.FullPath))
namesToDeleteAndUpdate[string(existingEntry.FullPath)] = struct{}{}
existingEntry.Deleted = true
existingEntry.TaskID = task.ID
//logging.LogInfo("found process to delete", "name", string(existingEntry.Name), "pid", string(existingEntry.FullPath))
idsToDelete = append(idsToDelete, existingEntry.ID)
deleteTreeNode(existingEntry, false)
}
}
// now all existing ones have been updated or deleted, so it's time to add new ones
for _, newEntry := range *processes {
if newEntry.Name == "" {
newEntry.Name = "unknown"
}
if !utils.SliceContains(namesToDeleteAndUpdate, strconv.Itoa(newEntry.ProcessID)) {
// this isn't marked as updated or deleted, so let's create it
parentPath := strconv.Itoa(newEntry.ParentProcessID)
if newEntry.ParentProcessID <= 0 {
parentPath = ""
}
fullPath := treeNodeGetFullPath(
[]byte(parentPath),
[]byte(strconv.Itoa(newEntry.ProcessID)),
[]byte("/"),
databaseStructs.TREE_TYPE_PROCESS)
newTree := databaseStructs.MythicTree{
Host: host,
TaskID: task.ID,
OperationID: task.OperationID,
Name: []byte(newEntry.Name),
ParentPath: []byte(parentPath),
FullPath: fullPath,
TreeType: databaseStructs.TREE_TYPE_PROCESS,
CanHaveChildren: true,
Deleted: false,
DisplayPath: []byte{},
}
if newEntry.OS != nil {
newTree.Os = *newEntry.OS
} else {
newTree.Os = task.Callback.Payload.Os
}
metadata := map[string]interface{}{
"process_id": newEntry.ProcessID,
"parent_process_id": newEntry.ParentProcessID,
"architecture": newEntry.Architecture,
"bin_path": newEntry.BinPath,
"name": newEntry.Name,
"user": newEntry.User,
"command_line": newEntry.CommandLine,
"integrity_level": newEntry.IntegrityLevel,
"start_time": newEntry.StartTime,
"description": newEntry.Description,
"signer": newEntry.Signer,
"protected_process_level": newEntry.ProtectionProcessLevel,
}
reflectBackOtherKeys(&metadata, &newEntry.Other)
newTree.Metadata = GetMythicJSONTextFromStruct(metadata)
newTree.CallbackID.Valid = true
newTree.CallbackID.Int64 = int64(task.Callback.ID)
if apitokensId > 0 {
newTree.APITokensID.Valid = true
newTree.APITokensID.Int64 = int64(apitokensId)
}
createTreeNode(&newTree)
if _, ok := namesToDeleteAndUpdate[strconv.Itoa(newEntry.ProcessID)]; !ok {
treeNodesToUpsert = append(treeNodesToUpsert, buildProcessMythicTreeNode(task, host, newEntry, apitokensId))
}
}
if err := upsertMythicTreeNodes(treeNodesToUpsert); err != nil {
logging.LogError(err, "Failed to upsert process MythicTree entries")
return err
}
// delete all the ids marked for deletion
if len(idsToDelete) > 0 {
deleteQuery, args, err := sqlx.Named(`UPDATE mythictree SET
@@ -2407,95 +2181,25 @@ func HandleAgentMessagePostResponseProcesses(task databaseStructs.Task, processe
}
}
} else {
treeNodesToUpsert := make([]*databaseStructs.MythicTree, 0, len(*processes))
for _, newEntry := range *processes {
if newEntry.Name == "" {
newEntry.Name = "unknown"
}
host = task.Callback.Host
if newEntry.Host != nil && *newEntry.Host != "" {
host = strings.ToUpper(*newEntry.Host)
}
// can't pre-create parent entries since that info isn't known
parentPath := strconv.Itoa(newEntry.ParentProcessID)
if newEntry.ParentProcessID <= 0 {
parentPath = ""
}
fullPath := treeNodeGetFullPath(
[]byte(parentPath),
[]byte(strconv.Itoa(newEntry.ProcessID)),
[]byte("/"),
databaseStructs.TREE_TYPE_PROCESS)
newTree := databaseStructs.MythicTree{
Host: host,
TaskID: task.ID,
OperationID: task.OperationID,
Name: []byte(newEntry.Name),
ParentPath: []byte(parentPath),
FullPath: fullPath,
TreeType: databaseStructs.TREE_TYPE_PROCESS,
CanHaveChildren: true,
Deleted: false,
DisplayPath: []byte{},
}
if newEntry.OS != nil {
newTree.Os = *newEntry.OS
} else {
newTree.Os = task.Callback.Payload.Os
}
metadata := map[string]interface{}{
"process_id": newEntry.ProcessID,
"parent_process_id": newEntry.ParentProcessID,
"architecture": newEntry.Architecture,
"bin_path": newEntry.BinPath,
"name": newEntry.Name,
"user": newEntry.User,
"command_line": newEntry.CommandLine,
"integrity_level": newEntry.IntegrityLevel,
"start_time": newEntry.StartTime,
"description": newEntry.Description,
"signer": newEntry.Signer,
"protected_process_level": newEntry.ProtectionProcessLevel,
}
reflectBackOtherKeys(&metadata, &newEntry.Other)
newTree.Metadata = GetMythicJSONTextFromStruct(metadata)
newTree.CallbackID.Valid = true
newTree.CallbackID.Int64 = int64(task.Callback.ID)
if apitokensId > 0 {
newTree.APITokensID.Valid = true
newTree.APITokensID.Int64 = int64(apitokensId)
}
createTreeNode(&newTree)
treeNodesToUpsert = append(treeNodesToUpsert, buildProcessMythicTreeNode(task, host, newEntry, apitokensId))
}
if err := upsertMythicTreeNodes(treeNodesToUpsert); err != nil {
logging.LogError(err, "Failed to upsert process MythicTree entries")
return err
}
}
return nil
}
func resolveAndCreateParentPathsForTreeNode(pathData utils.AnalyzedPath, displayPathData utils.AnalyzedPath, task databaseStructs.Task, treeType string) {
for i, _ := range pathData.PathPieces {
parentPath, fullPath, name := getParentPathFullPathName(pathData, i, treeType)
if parentPath == "" && fullPath == "" && name == "" {
continue
}
displayParentPath, displayFullPath, displayName := getParentPathFullPathName(displayPathData, i, treeType)
newTree := databaseStructs.MythicTree{
Host: pathData.Host,
TaskID: task.ID,
OperationID: task.OperationID,
Name: []byte(name),
FullPath: []byte(fullPath),
ParentPath: []byte(parentPath),
TreeType: treeType,
CanHaveChildren: true,
Deleted: false,
HasChildren: true,
DisplayPath: []byte(displayFullPath),
}
newTree.Metadata = GetMythicJSONTextFromStruct(nil)
newTree.Os = getOSTypeBasedOnPathSeparator(pathData.PathSeparator, treeType)
newTree.CallbackID.Valid = true
newTree.CallbackID.Int64 = int64(task.Callback.ID)
logging.LogInfo("creating parent", "name", name, "parentPath", parentPath, "fullPath", fullPath)
logging.LogInfo("display path data for parent path", "displayParentPath", displayParentPath, "displayFullPath", displayFullPath, "displayName", displayName)
createTreeNode(&newTree)
parentNodes := buildMythicTreeParentNodes(pathData, displayPathData, task, treeType)
if err := upsertMythicTreeNodes(parentNodes); err != nil {
logging.LogError(err, "Failed to create MythicTree parent paths")
}
}
func getOSTypeBasedOnPathSeparator(pathSeparator string, treeType string) string {
@@ -2519,7 +2223,7 @@ func getOSTypeBasedOnPathSeparator(pathSeparator string, treeType string) string
func treeNodeGetFullPath(parentPath []byte, name []byte, pathSeparator []byte, treeType string) []byte {
switch treeType {
case databaseStructs.TREE_TYPE_FILE:
fullPath := parentPath
fullPath := append([]byte(nil), parentPath...)
// if parent path is empty, then we don't want to add an extra instance of the path separator
if len(fullPath) > 0 {
if fullPath[len(fullPath)-1] == pathSeparator[len(pathSeparator)-1] {
@@ -2537,11 +2241,11 @@ func treeNodeGetFullPath(parentPath []byte, name []byte, pathSeparator []byte, t
return fullPath
case databaseStructs.TREE_TYPE_PROCESS:
// full path is just the process_id of the current process
return name
return append([]byte(nil), name...)
default:
custom, ok := getCustomBrowser(treeType)
if ok && custom.Type == databaseStructs.TREE_TYPE_FILE {
fullPath := parentPath
fullPath := append([]byte(nil), parentPath...)
// if parent path is empty, then we don't want to add an extra instance of the path separator
if len(fullPath) > 0 {
if fullPath[len(fullPath)-1] == pathSeparator[len(pathSeparator)-1] {
@@ -2671,29 +2375,7 @@ func deleteTreeNode(treeNode databaseStructs.MythicTree, cascade bool) {
}
}
func createTreeNode(treeNode *databaseStructs.MythicTree) {
if len(treeNode.Name) == 0 {
logging.LogError(nil, "Can't create file browser entry with empty name", "tree", treeNode)
return
}
statement, err := database.DB.PrepareNamed(`INSERT INTO mythictree
(host, task_id, operation_id, "name", full_path, parent_path, tree_type, can_have_children, success, metadata, os, callback_id, apitokens_id, has_children, display_path)
VALUES
(:host, :task_id, :operation_id, :name, :full_path, :parent_path, :tree_type, :can_have_children, :success, :metadata, :os, :callback_id, :apitokens_id, :has_children, :display_path)
ON CONFLICT (host, operation_id, full_path, tree_type, callback_id)
DO UPDATE SET
task_id=:task_id, "name"=:name, parent_path=:parent_path,
can_have_children=(mythictree.has_children OR :has_children OR :can_have_children),
has_children=(mythictree.has_children OR :has_children),
metadata=mythictree.metadata || :metadata, os=:os, "timestamp"=now(), deleted=false, display_path=:display_path,
success=(mythictree.success OR :success)
RETURNING id`)
if err != nil {
logging.LogError(err, "Failed to create or update mythictree statement")
return
}
err = statement.Get(&treeNode.ID, treeNode)
if err != nil {
if err := upsertMythicTreeNode(treeNode); err != nil {
logging.LogError(err, "Failed to create or update mythictree entry")
}
}
@@ -2930,8 +2612,7 @@ func handleAgentMessagePostResponseCustomBrowser(task databaseStructs.Task, agen
displayData.PathPieces = displayData.PathPieces[1:]
displayData.ReverseCombine = true
}
logging.LogInfo("SplitCustomPathResponse", "pathData", pathData)
resolveAndCreateParentPathsForTreeNode(pathData, displayData, task, customBrowserData.Name)
treeNodesToUpsert := buildMythicTreeParentNodes(pathData, displayData, task, customBrowserData.Name)
// now that the parents and all ancestors are resolved, process the current path and all children
realParentPath := strings.Join(pathData.PathPieces, pathData.PathSeparator)
if entry.Name == "" {
@@ -2945,7 +2626,6 @@ func handleAgentMessagePostResponseCustomBrowser(task databaseStructs.Task, agen
customBrowserData.Name)
parentPath := treeNodeGetFullPath([]byte(realParentPath), []byte(""), []byte(pathData.PathSeparator), customBrowserData.Name)
name := treeNodeGetFullPath([]byte(""), []byte(entry.Name), []byte(pathData.PathSeparator), customBrowserData.Name)
logging.LogInfo("creating entry", "name", entry.Name, "parentPath", realParentPath, "fullPath", fullPath)
newTree := databaseStructs.MythicTree{
Host: pathData.Host,
TaskID: task.ID,
@@ -2969,13 +2649,13 @@ func handleAgentMessagePostResponseCustomBrowser(task databaseStructs.Task, agen
newTree.Metadata = GetMythicJSONTextFromStruct(entry.Metadata)
newTree.CallbackID.Valid = true
newTree.CallbackID.Int64 = int64(task.Callback.ID)
createTreeNode(&newTree)
treeNodesToUpsert = append(treeNodesToUpsert, &newTree)
if agentCustomBrowser.UpdateDeleted {
if customBrowserData.Type == databaseStructs.TREE_TYPE_FILE {
// we need to iterate over the children for this entry and potentially remove any that the database know of but that aren't in our `files` list
var existingTreeEntries []databaseStructs.MythicTree
err = database.DB.Select(&existingTreeEntries, `SELECT
id, "name", success, full_path, parent_path, operation_id, host, tree_type
err = database.DB.Select(&existingTreeEntries, `SELECT
id, "name", success, full_path, parent_path, operation_id, host, tree_type, callback_id
FROM mythictree WHERE
parent_path=$1 AND operation_id=$2 AND host=$3 AND tree_type=$4`,
fullPath, task.OperationID, pathData.Host, customBrowserData.Name)
@@ -2983,71 +2663,62 @@ func handleAgentMessagePostResponseCustomBrowser(task databaseStructs.Task, agen
logging.LogError(err, "Failed to fetch existing children when trying to update deleted")
continue
}
var namesToDeleteAndUpdate []string // will get existing database IDs for things that aren't in the files list
for _, existingEntry := range existingTreeEntries {
if entry.Children != nil {
existingEntryStillExists := false
//logging.LogInfo("checking for file existing", "name", existingEntry.Name)
for _, newEntry := range *entry.Children {
if bytes.Equal([]byte(newEntry.Name), existingEntry.Name) {
//logging.LogInfo("[+] found a match in existing data and new data", "name", newEntry.Name, "id", existingEntry.ID)
namesToDeleteAndUpdate = append(namesToDeleteAndUpdate, newEntry.Name)
existingEntryStillExists = true
// update the entry in the database
newTreeChild := databaseStructs.MythicTree{
Host: pathData.Host,
TaskID: task.ID,
OperationID: task.OperationID,
Name: []byte(newEntry.Name),
ParentPath: existingEntry.ParentPath,
FullPath: existingEntry.FullPath,
TreeType: customBrowserData.Name,
CanHaveChildren: newEntry.CanHaveChildren,
DisplayPath: []byte(newEntry.DisplayPath),
Deleted: false,
Success: existingEntry.Success,
ID: existingEntry.ID,
Os: newTree.Os,
}
newTreeChild.Metadata = GetMythicJSONTextFromStruct(newEntry.Metadata)
newTreeChild.CallbackID.Valid = true
newTreeChild.CallbackID.Int64 = int64(task.Callback.ID)
updateTreeNode(newTreeChild)
}
}
if !existingEntryStillExists {
//logging.LogError(nil, "failed to find match, marking as deleted", "name", existingEntry.Name)
namesToDeleteAndUpdate = append(namesToDeleteAndUpdate, string(existingEntry.Name))
existingEntry.Deleted = true
existingEntry.TaskID = task.ID
deleteTreeNode(existingEntry, true)
}
}
}
// now all existing ones have been updated or deleted, so it's time to add new ones
incomingChildrenByName := make(map[string]agentMessagePostResponseCustomBrowserChildren)
if entry.Children != nil {
for _, newEntry := range *entry.Children {
if !utils.SliceContains(namesToDeleteAndUpdate, newEntry.Name) {
// this isn't marked as updated or deleted, so let's create it
newTreeChild := databaseStructs.MythicTree{
Host: pathData.Host,
TaskID: task.ID,
OperationID: task.OperationID,
Name: []byte(newEntry.Name),
DisplayPath: []byte(newEntry.DisplayPath),
ParentPath: fullPath,
TreeType: customBrowserData.Name,
CanHaveChildren: newEntry.CanHaveChildren,
Deleted: false,
Os: newTree.Os,
}
newTreeChild.FullPath = treeNodeGetFullPath(fullPath, []byte(newEntry.Name), []byte(pathData.PathSeparator), customBrowserData.Name)
newTreeChild.Metadata = GetMythicJSONTextFromStruct(newEntry.Metadata)
newTreeChild.CallbackID.Valid = true
newTreeChild.CallbackID.Int64 = int64(task.Callback.ID)
createTreeNode(&newTreeChild)
incomingChildrenByName[newEntry.Name] = newEntry
}
}
namesToDeleteAndUpdate := make(map[string]struct{}, len(existingTreeEntries)+len(incomingChildrenByName))
for _, existingEntry := range existingTreeEntries {
if newEntry, ok := incomingChildrenByName[string(existingEntry.Name)]; ok {
namesToDeleteAndUpdate[newEntry.Name] = struct{}{}
newTreeChild := databaseStructs.MythicTree{
Host: pathData.Host,
TaskID: task.ID,
OperationID: task.OperationID,
Name: []byte(newEntry.Name),
ParentPath: existingEntry.ParentPath,
FullPath: existingEntry.FullPath,
TreeType: customBrowserData.Name,
CanHaveChildren: newEntry.CanHaveChildren,
DisplayPath: []byte(newEntry.DisplayPath),
Deleted: false,
Success: existingEntry.Success,
ID: existingEntry.ID,
Os: newTree.Os,
}
newTreeChild.Metadata = GetMythicJSONTextFromStruct(newEntry.Metadata)
newTreeChild.CallbackID.Valid = true
newTreeChild.CallbackID.Int64 = int64(task.Callback.ID)
treeNodesToUpsert = append(treeNodesToUpsert, &newTreeChild)
} else {
namesToDeleteAndUpdate[string(existingEntry.Name)] = struct{}{}
existingEntry.Deleted = true
existingEntry.TaskID = task.ID
deleteTreeNode(existingEntry, true)
}
}
// now all existing ones have been updated or deleted, so it's time to add new ones
for name, newEntry := range incomingChildrenByName {
if _, ok := namesToDeleteAndUpdate[name]; !ok {
newTreeChild := databaseStructs.MythicTree{
Host: pathData.Host,
TaskID: task.ID,
OperationID: task.OperationID,
Name: []byte(newEntry.Name),
DisplayPath: []byte(newEntry.DisplayPath),
ParentPath: fullPath,
TreeType: customBrowserData.Name,
CanHaveChildren: newEntry.CanHaveChildren,
Deleted: false,
Os: newTree.Os,
}
newTreeChild.FullPath = treeNodeGetFullPath(fullPath, []byte(newEntry.Name), []byte(pathData.PathSeparator), customBrowserData.Name)
newTreeChild.Metadata = GetMythicJSONTextFromStruct(newEntry.Metadata)
newTreeChild.CallbackID.Valid = true
newTreeChild.CallbackID.Int64 = int64(task.Callback.ID)
treeNodesToUpsert = append(treeNodesToUpsert, &newTreeChild)
}
}
}
@@ -3071,9 +2742,11 @@ func handleAgentMessagePostResponseCustomBrowser(task databaseStructs.Task, agen
newTreeChild.Metadata = GetMythicJSONTextFromStruct(newEntry.Metadata)
newTreeChild.CallbackID.Valid = true
newTreeChild.CallbackID.Int64 = int64(task.Callback.ID)
logging.LogInfo("creating entry child", "name", newEntry.Name, "parentPath", fullPath)
createTreeNode(&newTreeChild)
treeNodesToUpsert = append(treeNodesToUpsert, &newTreeChild)
}
}
if err := upsertMythicTreeNodes(treeNodesToUpsert); err != nil {
logging.LogError(err, "Failed to upsert custom browser MythicTree entries")
}
}
}
mythic-docker/src/rabbitmq/utils_mythictree.go
+516
View File
@@ -0,0 +1,516 @@
package rabbitmq
import (
"database/sql"
"fmt"
"strconv"
"strings"
"github.com/its-a-feature/Mythic/database"
databaseStructs "github.com/its-a-feature/Mythic/database/structs"
"github.com/its-a-feature/Mythic/logging"
"github.com/its-a-feature/Mythic/utils"
)
// MythicTree ingest can happen on busy callback paths for file browser,
// process, and custom browser data. Keep the shared helpers here so the
// message handlers can focus on request-specific normalization while this file
// owns the DB write shape, duplicate handling, and process node construction.
// upsertMythicTreeNodeQuery mirrors the previous single-row createTreeNode
// behavior. It intentionally preserves metadata merge and bool-OR semantics so
// batching rows does not change what an ON CONFLICT update means.
const upsertMythicTreeNodeQuery = `INSERT INTO mythictree
(host, task_id, operation_id, "name", full_path, parent_path, tree_type, can_have_children, success, metadata, os, callback_id, apitokens_id, has_children, display_path)
VALUES
(:host, :task_id, :operation_id, :name, :full_path, :parent_path, :tree_type, :can_have_children, :success, :metadata, :os, :callback_id, :apitokens_id, :has_children, :display_path)
ON CONFLICT (host, operation_id, full_path, tree_type, callback_id)
DO UPDATE SET
task_id=:task_id, "name"=:name, parent_path=:parent_path,
can_have_children=(mythictree.has_children OR :has_children OR :can_have_children),
has_children=(mythictree.has_children OR :has_children),
metadata=mythictree.metadata || :metadata, os=:os, "timestamp"=now(), deleted=false, display_path=:display_path,
success=(mythictree.success OR :success)
RETURNING id`
// mythicTreeNodeKey represents the database uniqueness key used by the
// MythicTree upsert. The callback validity flag is included so an invalid
// callback id cannot accidentally collapse with callback_id=0.
type mythicTreeNodeKey struct {
Host string
OperationID int
FullPath string
TreeType string
CallbackID int64
CallbackValid bool
}
type fileBrowserUpdateDeletedGroup struct {
PathData utils.AnalyzedPath
FullPath []byte
ChildrenByName map[string]agentMessagePostResponseFileBrowserChildren
}
// upsertMythicTreeNode is the compatibility wrapper for callers that still need
// the created/updated ID immediately, such as filemeta association.
func upsertMythicTreeNode(treeNode *databaseStructs.MythicTree) error {
if treeNode == nil {
return nil
}
if err := normalizeMythicTreeNodeForUpsert(treeNode); err != nil {
return err
}
statement, err := database.DB.PrepareNamed(upsertMythicTreeNodeQuery)
if err != nil {
return err
}
defer statement.Close()
if err = statement.Get(&treeNode.ID, treeNode); err != nil {
return fmt.Errorf("failed to create or update MythicTree entry %q: %w", string(treeNode.FullPath), err)
}
return nil
}
// upsertMythicTreeNodes writes a batch of MythicTree nodes through one
// transaction and one prepared statement. This avoids per-node prepare/autocommit
// overhead while keeping RETURNING id behavior for every node in the batch.
func upsertMythicTreeNodes(treeNodes []*databaseStructs.MythicTree) error {
treeNodes = normalizeMythicTreeNodesForBatch(treeNodes)
treeNodes = dedupeMythicTreeNodes(treeNodes)
if len(treeNodes) == 0 {
return nil
}
transaction, err := database.DB.Beginx()
if err != nil {
return err
}
committed := false
defer func() {
if !committed {
_ = transaction.Rollback()
}
}()
statement, err := transaction.PrepareNamed(upsertMythicTreeNodeQuery)
if err != nil {
return err
}
defer statement.Close()
for _, treeNode := range treeNodes {
if err = statement.Get(&treeNode.ID, treeNode); err != nil {
return fmt.Errorf("failed to create or update MythicTree entry %q: %w", string(treeNode.FullPath), err)
}
}
if err = transaction.Commit(); err != nil {
return err
}
committed = true
return nil
}
// normalizeMythicTreeNodesForBatch applies required-field defaults while
// preserving the old per-row behavior for invalid entries: skip the bad node and
// continue writing the rest of the batch.
func normalizeMythicTreeNodesForBatch(treeNodes []*databaseStructs.MythicTree) []*databaseStructs.MythicTree {
normalized := make([]*databaseStructs.MythicTree, 0, len(treeNodes))
for _, treeNode := range treeNodes {
if treeNode == nil {
continue
}
if len(treeNode.Name) == 0 {
logging.LogError(nil, "Skipping MythicTree entry with empty name", "full_path", string(treeNode.FullPath), "tree_type", treeNode.TreeType)
continue
}
if len(treeNode.DisplayPath) == 0 {
treeNode.DisplayPath = append([]byte(nil), treeNode.Name...)
}
cloneMythicTreeNodeBytes(treeNode)
normalized = append(normalized, treeNode)
}
return normalized
}
// normalizeMythicTreeNodeForUpsert fills required database fields that callers
// may omit. Agents do not always provide a display path, so default it to the
// node name before insert/update to avoid null display_path writes and keep the
// UI label meaningful.
func normalizeMythicTreeNodeForUpsert(treeNode *databaseStructs.MythicTree) error {
if treeNode == nil {
return nil
}
if len(treeNode.Name) == 0 {
return fmt.Errorf("can't create MythicTree entry with empty name")
}
if len(treeNode.DisplayPath) == 0 {
treeNode.DisplayPath = append([]byte(nil), treeNode.Name...)
}
cloneMythicTreeNodeBytes(treeNode)
return nil
}
// cloneMythicTreeNodeBytes gives queued nodes ownership of mutable byte slices.
// Batches hold many MythicTree structs in memory before DB writes, so none of
// their path/name fields should alias shared builder buffers.
func cloneMythicTreeNodeBytes(treeNode *databaseStructs.MythicTree) {
treeNode.Name = cloneByteSliceForDatabase(treeNode.Name)
treeNode.FullPath = cloneByteSliceForDatabase(treeNode.FullPath)
treeNode.ParentPath = cloneByteSliceForDatabase(treeNode.ParentPath)
treeNode.DisplayPath = cloneByteSliceForDatabase(treeNode.DisplayPath)
}
// cloneByteSliceForDatabase gives callers a private slice while preserving
// empty-but-non-null bytea values. PostgreSQL treats nil []byte as NULL, which
// is different from the empty path MythicTree uses for root parent_path values.
func cloneByteSliceForDatabase(input []byte) []byte {
output := make([]byte, len(input))
copy(output, input)
return output
}
// dedupeMythicTreeNodes removes duplicate rows that would hit the same
// ON CONFLICT key inside a single batch. PostgreSQL rejects a multi-source
// statement that updates the same target row twice, and even though this code
// loops prepared statements, deduping saves extra DB work while preserving the
// effective "last write wins plus merged metadata" behavior.
func dedupeMythicTreeNodes(treeNodes []*databaseStructs.MythicTree) []*databaseStructs.MythicTree {
if len(treeNodes) == 0 {
return treeNodes
}
deduped := make([]*databaseStructs.MythicTree, 0, len(treeNodes))
seen := make(map[mythicTreeNodeKey]int, len(treeNodes))
for _, treeNode := range treeNodes {
if treeNode == nil || len(treeNode.Name) == 0 {
continue
}
key := getMythicTreeNodeKey(treeNode)
if index, ok := seen[key]; ok {
deduped[index] = mergeMythicTreeNodes(deduped[index], treeNode)
continue
}
seen[key] = len(deduped)
deduped = append(deduped, treeNode)
}
return deduped
}
// getMythicTreeNodeKey builds the conflict identity used before the DB write so
// parent/current/child nodes in the same response can be merged safely.
func getMythicTreeNodeKey(treeNode *databaseStructs.MythicTree) mythicTreeNodeKey {
return mythicTreeNodeKey{
Host: treeNode.Host,
OperationID: treeNode.OperationID,
FullPath: string(treeNode.FullPath),
TreeType: treeNode.TreeType,
CallbackID: treeNode.CallbackID.Int64,
CallbackValid: treeNode.CallbackID.Valid,
}
}
// mergeMythicTreeNodes combines two queued writes for the same MythicTree row.
// It mirrors the SQL upsert rules: preserve child state, OR success, merge JSON
// metadata, and keep an existing API token if the incoming node omitted one.
func mergeMythicTreeNodes(existing *databaseStructs.MythicTree, incoming *databaseStructs.MythicTree) *databaseStructs.MythicTree {
incoming.CanHaveChildren = existing.CanHaveChildren || existing.HasChildren || incoming.CanHaveChildren
incoming.HasChildren = existing.HasChildren || incoming.HasChildren
incoming.Success = mythicTreeSQLBoolOr(existing.Success, incoming.Success)
incoming.Metadata = mergeMythicTreeMetadata(existing.Metadata, incoming.Metadata)
if !incoming.APITokensID.Valid && existing.APITokensID.Valid {
incoming.APITokensID = existing.APITokensID
}
return incoming
}
// mythicTreeSQLBoolOr matches PostgreSQL's nullable boolean OR behavior closely
// enough for pre-DB merges: true wins, false only survives when both sides are
// explicitly valid, and unknown remains unknown otherwise.
func mythicTreeSQLBoolOr(left sql.NullBool, right sql.NullBool) sql.NullBool {
if (left.Valid && left.Bool) || (right.Valid && right.Bool) {
return sql.NullBool{Valid: true, Bool: true}
}
if left.Valid && right.Valid {
return sql.NullBool{Valid: true, Bool: false}
}
return sql.NullBool{}
}
// mergeMythicTreeMetadata applies the same right-side override behavior as the
// jsonb concatenation used in the database upsert.
func mergeMythicTreeMetadata(left databaseStructs.MythicJSONText, right databaseStructs.MythicJSONText) databaseStructs.MythicJSONText {
merged := left.StructValue()
for key, value := range right.StructValue() {
merged[key] = value
}
return GetMythicJSONTextFromStruct(merged)
}
// buildMythicTreeParentNodes constructs all ancestor nodes for a path without
// writing them immediately. Callers can append the current node and children,
// then flush everything in one batch.
func buildMythicTreeParentNodes(pathData utils.AnalyzedPath, displayPathData utils.AnalyzedPath, task databaseStructs.Task, treeType string) []*databaseStructs.MythicTree {
parentNodes := make([]*databaseStructs.MythicTree, 0, len(pathData.PathPieces))
for i := range pathData.PathPieces {
parentPath, fullPath, name := getParentPathFullPathName(pathData, i, treeType)
if parentPath == "" && fullPath == "" && name == "" {
continue
}
_, displayFullPath, _ := getParentPathFullPathName(displayPathData, i, treeType)
newTree := databaseStructs.MythicTree{
Host: pathData.Host,
TaskID: task.ID,
OperationID: task.OperationID,
Name: []byte(name),
FullPath: []byte(fullPath),
ParentPath: []byte(parentPath),
TreeType: treeType,
CanHaveChildren: true,
Deleted: false,
HasChildren: true,
DisplayPath: []byte(displayFullPath),
}
newTree.Metadata = GetMythicJSONTextFromStruct(nil)
newTree.Os = getOSTypeBasedOnPathSeparator(pathData.PathSeparator, treeType)
newTree.CallbackID.Valid = true
newTree.CallbackID.Int64 = int64(task.Callback.ID)
parentNodes = append(parentNodes, &newTree)
}
return parentNodes
}
// buildFileBrowserChildMythicTreeNode converts a reported child entry into the
// exact MythicTree row shape used by both normal file browser handling and
// update_deleted reconciliation.
func buildFileBrowserChildMythicTreeNode(task databaseStructs.Task, host string, parentFullPath []byte, pathSeparator string, child agentMessagePostResponseFileBrowserChildren, os string, apitokensId int) *databaseStructs.MythicTree {
newTreeChild := databaseStructs.MythicTree{
Host: host,
TaskID: task.ID,
OperationID: task.OperationID,
Name: []byte(child.Name),
ParentPath: parentFullPath,
TreeType: databaseStructs.TREE_TYPE_FILE,
CanHaveChildren: !child.IsFile,
Deleted: false,
Os: os,
DisplayPath: []byte{},
}
newTreeChild.FullPath = treeNodeGetFullPath(parentFullPath, []byte(child.Name), []byte(pathSeparator), databaseStructs.TREE_TYPE_FILE)
fileMetaData := addChildFilePermissions(&child)
newTreeChild.Metadata = GetMythicJSONTextFromStruct(fileMetaData)
newTreeChild.CallbackID.Valid = true
newTreeChild.CallbackID.Int64 = int64(task.Callback.ID)
if apitokensId > 0 {
newTreeChild.APITokensID.Valid = true
newTreeChild.APITokensID.Int64 = int64(apitokensId)
}
return &newTreeChild
}
// groupFileBrowserUpdateDeletedEntries keeps chunked listings for the same
// directory together while preventing entries for different directories in the
// same task from being reconciled against the first path only.
func groupFileBrowserUpdateDeletedEntries(task *databaseStructs.Task, fileBrowsers []*agentMessagePostResponseFileBrowser) ([]fileBrowserUpdateDeletedGroup, error) {
groups := make([]fileBrowserUpdateDeletedGroup, 0, len(fileBrowsers))
groupsByPath := make(map[string]int, len(fileBrowsers))
for _, fileBrowser := range fileBrowsers {
if fileBrowser == nil {
continue
}
if fileBrowser.Name == "" {
return nil, fmt.Errorf("can't reconcile file browser update_deleted entry with empty name")
}
pathData, fullPath, err := getFileBrowserUpdateDeletedPath(task, fileBrowser)
if err != nil {
return nil, err
}
groupKey := pathData.Host + "\x00" + string(fullPath)
groupIndex, ok := groupsByPath[groupKey]
if !ok {
groupIndex = len(groups)
groupsByPath[groupKey] = groupIndex
groups = append(groups, fileBrowserUpdateDeletedGroup{
PathData: pathData,
FullPath: fullPath,
ChildrenByName: make(map[string]agentMessagePostResponseFileBrowserChildren),
})
}
if fileBrowser.Files == nil {
continue
}
for _, newEntry := range *fileBrowser.Files {
groups[groupIndex].ChildrenByName[newEntry.Name] = newEntry
}
}
return groups, nil
}
// getFileBrowserUpdateDeletedPath derives the database parent identity for one
// update_deleted response. This mirrors the current-node full_path logic in the
// main file browser handler so reconciliation queries the correct child set.
func getFileBrowserUpdateDeletedPath(task *databaseStructs.Task, fileBrowser *agentMessagePostResponseFileBrowser) (utils.AnalyzedPath, []byte, error) {
pathData, err := utils.SplitFilePathGetHost(fileBrowser.ParentPath, fileBrowser.Name, []string{})
if err != nil {
return utils.AnalyzedPath{}, nil, err
}
if pathData.Host == "" {
pathData.Host = strings.ToUpper(task.Callback.Host)
}
if fileBrowser.Host != "" {
pathData.Host = strings.ToUpper(fileBrowser.Host)
}
realParentPath := strings.Join(pathData.PathPieces, pathData.PathSeparator)
if len(realParentPath) > 2 && realParentPath[0] == '/' && realParentPath[1] == '/' {
realParentPath = realParentPath[1:]
}
fullPath := treeNodeGetFullPath(
[]byte(realParentPath),
[]byte(fileBrowser.Name),
[]byte(pathData.PathSeparator),
databaseStructs.TREE_TYPE_FILE)
return pathData, fullPath, nil
}
// reconcileFileBrowserUpdateDeletedGroup applies update_deleted semantics for a
// single listed directory: update reported children, insert new children, and
// mark DB-only children deleted. Each group is independent so a task can report
// multiple directories without cross-contaminating their child sets.
func reconcileFileBrowserUpdateDeletedGroup(task *databaseStructs.Task, group fileBrowserUpdateDeletedGroup, apitokensId int) error {
var existingTreeEntries []databaseStructs.MythicTree
err := database.DB.Select(&existingTreeEntries, `SELECT
id, "name", success, full_path, parent_path, operation_id, host, tree_type, callback_id, os, display_path
FROM mythictree WHERE
parent_path=$1 AND operation_id=$2 AND host=$3 AND tree_type=$4`,
group.FullPath, task.OperationID, group.PathData.Host, databaseStructs.TREE_TYPE_FILE)
if err != nil {
return err
}
namesToDeleteAndUpdate := make(map[string]struct{}, len(existingTreeEntries)+len(group.ChildrenByName))
treeNodesToUpsert := make([]*databaseStructs.MythicTree, 0, len(group.ChildrenByName))
for _, existingEntry := range existingTreeEntries {
if newEntry, ok := group.ChildrenByName[string(existingEntry.Name)]; ok {
namesToDeleteAndUpdate[newEntry.Name] = struct{}{}
newTreeChild := databaseStructs.MythicTree{
Host: group.PathData.Host,
TaskID: task.ID,
OperationID: task.OperationID,
Name: []byte(newEntry.Name),
ParentPath: existingEntry.ParentPath,
FullPath: existingEntry.FullPath,
TreeType: databaseStructs.TREE_TYPE_FILE,
CanHaveChildren: !newEntry.IsFile,
Deleted: false,
Success: existingEntry.Success,
ID: existingEntry.ID,
Os: existingEntry.Os,
DisplayPath: existingEntry.DisplayPath,
CallbackID: existingEntry.CallbackID,
}
fileMetaData := addChildFilePermissions(&newEntry)
newTreeChild.Metadata = GetMythicJSONTextFromStruct(fileMetaData)
if !newTreeChild.CallbackID.Valid {
newTreeChild.CallbackID.Valid = true
newTreeChild.CallbackID.Int64 = int64(task.Callback.ID)
}
treeNodesToUpsert = append(treeNodesToUpsert, &newTreeChild)
} else {
namesToDeleteAndUpdate[string(existingEntry.Name)] = struct{}{}
existingEntry.Deleted = true
existingEntry.TaskID = task.ID
deleteTreeNode(existingEntry, true)
}
}
for name, newEntry := range group.ChildrenByName {
if _, ok := namesToDeleteAndUpdate[name]; !ok {
treeNodesToUpsert = append(treeNodesToUpsert, buildFileBrowserChildMythicTreeNode(
*task,
group.PathData.Host,
group.FullPath,
group.PathData.PathSeparator,
newEntry,
getOSTypeBasedOnPathSeparator(group.PathData.PathSeparator, databaseStructs.TREE_TYPE_FILE),
apitokensId))
}
}
return upsertMythicTreeNodes(treeNodesToUpsert)
}
// buildProcessMythicTreeNode centralizes process-to-MythicTree normalization so
// update_deleted and normal process responses use identical path, OS, metadata,
// callback, and API token handling.
func buildProcessMythicTreeNode(task databaseStructs.Task, host string, process agentMessagePostResponseProcesses, apitokensId int) *databaseStructs.MythicTree {
process.Name = normalizeProcessName(process.Name)
parentPath := getProcessParentPath(process.ParentProcessID)
fullPath := treeNodeGetFullPath(
[]byte(parentPath),
[]byte(strconv.Itoa(process.ProcessID)),
[]byte("/"),
databaseStructs.TREE_TYPE_PROCESS)
newTree := databaseStructs.MythicTree{
Host: host,
TaskID: task.ID,
OperationID: task.OperationID,
Name: []byte(process.Name),
ParentPath: []byte(parentPath),
FullPath: fullPath,
TreeType: databaseStructs.TREE_TYPE_PROCESS,
CanHaveChildren: true,
Deleted: false,
DisplayPath: []byte{},
}
if process.OS != nil {
newTree.Os = *process.OS
} else {
newTree.Os = task.Callback.Payload.Os
}
metadata := map[string]interface{}{
"process_id": process.ProcessID,
"parent_process_id": process.ParentProcessID,
"architecture": process.Architecture,
"bin_path": process.BinPath,
"name": process.Name,
"user": process.User,
"command_line": process.CommandLine,
"integrity_level": process.IntegrityLevel,
"start_time": process.StartTime,
"description": process.Description,
"signer": process.Signer,
"protected_process_level": process.ProtectionProcessLevel,
}
reflectBackOtherKeys(&metadata, &process.Other)
newTree.Metadata = GetMythicJSONTextFromStruct(metadata)
newTree.CallbackID.Valid = true
newTree.CallbackID.Int64 = int64(task.Callback.ID)
if apitokensId > 0 {
newTree.APITokensID.Valid = true
newTree.APITokensID.Int64 = int64(apitokensId)
}
return &newTree
}
// normalizeProcessName preserves the existing MythicTree behavior for unnamed
// processes so process matching and metadata generation stay consistent.
func normalizeProcessName(name string) string {
if name == "" {
return "unknown"
}
return name
}
// getProcessParentPath keeps root/no-parent processes represented as an empty
// parent path, matching the previous process tree layout.
func getProcessParentPath(parentProcessID int) string {
if parentProcessID <= 0 {
return ""
}
return strconv.Itoa(parentProcessID)
}
// getProcessMatchKey creates the in-memory identity for reported process rows.
// PID alone is not enough here because the existing behavior treats name and
// parent PID as part of whether a process entry still matches.
func getProcessMatchKey(process agentMessagePostResponseProcesses) string {
return strconv.Itoa(process.ProcessID) + "\x00" + normalizeProcessName(process.Name) + "\x00" + getProcessParentPath(process.ParentProcessID)
}
// getExistingProcessMatchKey creates the same identity from database rows so
// update_deleted can compare existing and incoming process sets in O(n).
func getExistingProcessMatchKey(existingEntry databaseStructs.MythicTree) string {
return string(existingEntry.FullPath) + "\x00" + string(existingEntry.Name) + "\x00" + string(existingEntry.ParentPath)
}
mythic-docker/src/rabbitmq/utils_mythictree_test.go
+219
View File
@@ -0,0 +1,219 @@
package rabbitmq
import (
"database/sql"
"testing"
databaseStructs "github.com/its-a-feature/Mythic/database/structs"
)
func TestDedupeMythicTreeNodesMergesRepeatedKeys(t *testing.T) {
first := &databaseStructs.MythicTree{
Host: "HOST",
OperationID: 1,
FullPath: []byte("C:\\Windows"),
ParentPath: []byte("C:"),
Name: []byte("Windows"),
TreeType: databaseStructs.TREE_TYPE_FILE,
CanHaveChildren: true,
HasChildren: true,
Metadata: GetMythicJSONTextFromStruct(map[string]interface{}{"first": "value", "replace": "old"}),
}
first.CallbackID.Valid = true
first.CallbackID.Int64 = 10
second := &databaseStructs.MythicTree{
Host: "HOST",
OperationID: 1,
FullPath: []byte("C:\\Windows"),
ParentPath: []byte("C:"),
Name: []byte("Windows"),
TreeType: databaseStructs.TREE_TYPE_FILE,
Success: sql.NullBool{Valid: true, Bool: true},
Metadata: GetMythicJSONTextFromStruct(map[string]interface{}{"second": "value", "replace": "new"}),
}
second.CallbackID.Valid = true
second.CallbackID.Int64 = 10
deduped := dedupeMythicTreeNodes([]*databaseStructs.MythicTree{first, second})
if len(deduped) != 1 {
t.Fatalf("expected one deduped MythicTree node, got %d", len(deduped))
}
if !deduped[0].CanHaveChildren {
t.Fatalf("expected can_have_children to preserve prior child state")
}
if !deduped[0].HasChildren {
t.Fatalf("expected has_children to preserve prior child state")
}
if !deduped[0].Success.Valid || !deduped[0].Success.Bool {
t.Fatalf("expected success to merge with SQL OR semantics, got %#v", deduped[0].Success)
}
metadata := deduped[0].Metadata.StructValue()
if metadata["first"] != "value" || metadata["second"] != "value" || metadata["replace"] != "new" {
t.Fatalf("unexpected merged metadata: %#v", metadata)
}
}
func TestNormalizeMythicTreeNodeForUpsertDefaultsDisplayPathToName(t *testing.T) {
node := &databaseStructs.MythicTree{
Name: []byte("poseidon"),
}
if err := normalizeMythicTreeNodeForUpsert(node); err != nil {
t.Fatalf("expected display path normalization to succeed: %v", err)
}
if string(node.DisplayPath) != "poseidon" {
t.Fatalf("expected display path to default to name, got %q", string(node.DisplayPath))
}
}
func TestNormalizeMythicTreeNodeForUpsertKeepsProvidedDisplayPath(t *testing.T) {
node := &databaseStructs.MythicTree{
Name: []byte("poseidon"),
DisplayPath: []byte("Poseidon"),
}
if err := normalizeMythicTreeNodeForUpsert(node); err != nil {
t.Fatalf("expected display path normalization to succeed: %v", err)
}
if string(node.DisplayPath) != "Poseidon" {
t.Fatalf("expected provided display path to be preserved, got %q", string(node.DisplayPath))
}
}
func TestNormalizeMythicTreeNodesForBatchSkipsInvalidEntries(t *testing.T) {
valid := &databaseStructs.MythicTree{Name: []byte("valid")}
invalid := &databaseStructs.MythicTree{FullPath: []byte("/tmp/invalid")}
normalized := normalizeMythicTreeNodesForBatch([]*databaseStructs.MythicTree{nil, invalid, valid})
if len(normalized) != 1 {
t.Fatalf("expected only the valid node to remain, got %d nodes", len(normalized))
}
if normalized[0] != valid {
t.Fatalf("expected the valid node to be preserved")
}
if string(normalized[0].DisplayPath) != "valid" {
t.Fatalf("expected display path to default during batch normalization, got %q", string(normalized[0].DisplayPath))
}
}
func TestNormalizeMythicTreeNodeForUpsertKeepsEmptyParentPathNonNil(t *testing.T) {
node := &databaseStructs.MythicTree{
Name: []byte("/"),
FullPath: []byte("/"),
ParentPath: []byte{},
}
if err := normalizeMythicTreeNodeForUpsert(node); err != nil {
t.Fatalf("expected root node normalization to succeed: %v", err)
}
if node.ParentPath == nil {
t.Fatalf("expected empty parent path to remain non-nil for database bytea writes")
}
if len(node.ParentPath) != 0 {
t.Fatalf("expected empty parent path, got %q", string(node.ParentPath))
}
}
func TestBuildFileBrowserChildMythicTreeNodeDoesNotAliasSiblingFullPaths(t *testing.T) {
task := databaseStructs.Task{
ID: 5,
OperationID: 6,
}
task.Callback.ID = 7
parentPath := make([]byte, len("/tmp/listing"), len("/tmp/listing")+64)
copy(parentPath, "/tmp/listing")
first := buildFileBrowserChildMythicTreeNode(task, "HOST", parentPath, "/", agentMessagePostResponseFileBrowserChildren{Name: "alpha"}, "linux", 0)
second := buildFileBrowserChildMythicTreeNode(task, "HOST", parentPath, "/", agentMessagePostResponseFileBrowserChildren{Name: "bravo"}, "linux", 0)
if string(first.FullPath) != "/tmp/listing/alpha" {
t.Fatalf("expected first child full path to remain tied to its name, got %q", string(first.FullPath))
}
if string(second.FullPath) != "/tmp/listing/bravo" {
t.Fatalf("expected second child full path to match its name, got %q", string(second.FullPath))
}
if string(parentPath) != "/tmp/listing" {
t.Fatalf("expected parent path to remain unchanged, got %q", string(parentPath))
}
}
func TestGroupFileBrowserUpdateDeletedEntriesMergesChunksForSameDirectory(t *testing.T) {
task := databaseStructs.Task{}
task.Callback.Host = "HOST"
firstChildren := []agentMessagePostResponseFileBrowserChildren{{Name: "alpha"}}
secondChildren := []agentMessagePostResponseFileBrowserChildren{{Name: "bravo"}}
groups, err := groupFileBrowserUpdateDeletedEntries(&task, []*agentMessagePostResponseFileBrowser{
{ParentPath: "/tmp", Name: "listing", Files: &firstChildren},
{ParentPath: "/tmp", Name: "listing", Files: &secondChildren},
})
if err != nil {
t.Fatalf("expected grouping to succeed: %v", err)
}
if len(groups) != 1 {
t.Fatalf("expected chunks for the same directory to merge into one group, got %d", len(groups))
}
if _, ok := groups[0].ChildrenByName["alpha"]; !ok {
t.Fatalf("expected first chunk child to be present")
}
if _, ok := groups[0].ChildrenByName["bravo"]; !ok {
t.Fatalf("expected second chunk child to be present")
}
}
func TestGroupFileBrowserUpdateDeletedEntriesSeparatesDifferentDirectories(t *testing.T) {
task := databaseStructs.Task{}
task.Callback.Host = "HOST"
firstChildren := []agentMessagePostResponseFileBrowserChildren{{Name: "alpha"}}
secondChildren := []agentMessagePostResponseFileBrowserChildren{{Name: "bravo"}}
groups, err := groupFileBrowserUpdateDeletedEntries(&task, []*agentMessagePostResponseFileBrowser{
{ParentPath: "/tmp", Name: "first", Files: &firstChildren},
{ParentPath: "/tmp", Name: "second", Files: &secondChildren},
})
if err != nil {
t.Fatalf("expected grouping to succeed: %v", err)
}
if len(groups) != 2 {
t.Fatalf("expected different directories to remain separate, got %d groups", len(groups))
}
if string(groups[0].FullPath) == string(groups[1].FullPath) {
t.Fatalf("expected different directories to have different full paths")
}
if _, ok := groups[0].ChildrenByName["bravo"]; ok {
t.Fatalf("did not expect first directory to include second directory's child")
}
if _, ok := groups[1].ChildrenByName["alpha"]; ok {
t.Fatalf("did not expect second directory to include first directory's child")
}
}
func TestBuildProcessMythicTreeNodeNormalizesProcessData(t *testing.T) {
task := databaseStructs.Task{
ID: 5,
OperationID: 6,
}
task.Callback.ID = 7
task.Callback.Payload.Os = "linux"
process := agentMessagePostResponseProcesses{
ProcessID: 123,
ParentProcessID: 1,
}
node := buildProcessMythicTreeNode(task, "HOST", process, 0)
if string(node.Name) != "unknown" {
t.Fatalf("expected empty process name to normalize to unknown, got %q", string(node.Name))
}
if string(node.FullPath) != "123" {
t.Fatalf("expected process full path to be pid, got %q", string(node.FullPath))
}
if string(node.ParentPath) != "1" {
t.Fatalf("expected process parent path to be parent pid, got %q", string(node.ParentPath))
}
if !node.CallbackID.Valid || node.CallbackID.Int64 != 7 {
t.Fatalf("expected callback id 7, got %#v", node.CallbackID)
}
if node.Os != "linux" {
t.Fatalf("expected process OS to default from callback payload, got %q", node.Os)
}
}