fixing sql queries and container down resolver

mythic-docker/src/rabbitmq/check_container_status.go

@@ -31,6 +31,22 @@ var customBrowsersToCheck = map[string]databaseStructs.CustomBrowser{}
 var containerOnStartLock sync.Mutex
 var containerOnStartInFlight = map[string]bool{}
 
+func claimContainerOnStart(containerName string) bool {
+	containerOnStartLock.Lock()
+	defer containerOnStartLock.Unlock()
+	if containerOnStartInFlight[containerName] {
+		return false
+	}
+	containerOnStartInFlight[containerName] = true
+	return true
+}
+
+func clearContainerOnStart(containerName string) {
+	containerOnStartLock.Lock()
+	delete(containerOnStartInFlight, containerName)
+	containerOnStartLock.Unlock()
+}
+
 func checkContainerStatusAddPT() {
 	for {
 		pt := <-checkContainerStatusAddPtChannel
@@ -126,14 +142,11 @@ type rabbitmqAPIQuery struct {
 }
 
 func CreateGraphQLSpectatorAPITokenAndSendOnStartMessage(containerName string) {
-	containerOnStartLock.Lock()
-	if containerOnStartInFlight[containerName] {
-		containerOnStartLock.Unlock()
+	if !claimContainerOnStart(containerName) {
 		logging.LogInfo("Skipping duplicate container on start run already in progress", "container", containerName)
 		return
 	}
-	containerOnStartInFlight[containerName] = true
-	containerOnStartLock.Unlock()
+	defer clearContainerOnStart(containerName)
 	operations := []databaseStructs.Operation{}
 	err := database.DB.Select(&operations, `SELECT id FROM operation WHERE deleted=false and complete=false`)
 	if err != nil {
@@ -191,6 +204,7 @@ func CreateGraphQLSpectatorAPITokenAndSendOnStartMessage(containerName string) {
 			mythicjwt.SCOPE_PAYLOAD_READ,
 			mythicjwt.SCOPE_FILE_WRITE,
 			mythicjwt.SCOPE_TAG_WRITE,
+			mythicjwt.SCOPE_CALLBACK_WRITE,
 		}
 		statement, err := database.DB.PrepareNamed(`INSERT INTO apitokens
 		(token_value, operator_id, token_type, active, "name", created_by, task_id, callback_id, scopes)
@@ -209,16 +223,18 @@ func CreateGraphQLSpectatorAPITokenAndSendOnStartMessage(containerName string) {
 		plainAPITokenValue, storedAPITokenValue, err := mythicjwt.GenerateOpaqueAPIToken()
 		if err != nil {
 			logging.LogError(err, "failed to generate new API token")
+			expireAPIToken(apiToken.ID)
 			continue
 		}
 		apiToken.TokenValue = storedAPITokenValue
 		_, err = database.DB.Exec(`UPDATE apitokens SET token_value=$1 WHERE id=$2`, apiToken.TokenValue, apiToken.ID)
 		if err != nil {
 			logging.LogError(err, "Failed to update apitoken with value")
+			expireAPIToken(apiToken.ID)
 			continue
 		}
 		onStartMessage.APIToken = plainAPITokenValue
-		go updateAPITokenAfter5Minutes(apiToken.ID, containerName)
+		go expireAPITokenAfterShortLivedTTL(apiToken.ID)
 		err = RabbitMQConnection.SendContainerOnStart(onStartMessage)
 		if err != nil {
 			logging.LogError(err, "Failed to send container on start")
@@ -227,24 +243,15 @@ func CreateGraphQLSpectatorAPITokenAndSendOnStartMessage(containerName string) {
 		time.Sleep(5 * time.Second)
 	}
 }
-func updateAPITokenAfter5Minutes(apitokenID int, containerName string) {
-	<-time.After(5 * time.Minute)
-	_, err := database.DB.Exec(`UPDATE apitokens SET active=false, deleted=true WHERE id=$1`, apitokenID)
-	if err != nil {
-		logging.LogError(err, "failed to mark apitoken as deleted")
-	}
-	containerOnStartLock.Lock()
-	delete(containerOnStartInFlight, containerName)
-	containerOnStartLock.Unlock()
-}
 func checkContainerStatus() {
 	// get all queues from rabbitmq
 	// http://rabbitmq_user:rabbitmq_password@rabbitmq_host:15672/rabbitmq/api/queues/mythic_vhost
 	rabbitmqReqURL := fmt.Sprintf("http://%s:%s@%s:15672/api/queues/%s?use_regex=true&page=1&page_size=500&name=%s",
 		utils.MythicConfig.RabbitmqUser, utils.MythicConfig.RabbitmqPassword, utils.MythicConfig.RabbitmqHost,
 		utils.MythicConfig.RabbitmqVHost,
-		fmt.Sprintf("(.%%2A_%s|.%%2A_%s|.%%2A_%s)",
-			PT_BUILD_ROUTING_KEY, C2_RPC_START_SERVER_ROUTING_KEY, CONSUMING_CONTAINER_RESYNC_ROUTING_KEY))
+		fmt.Sprintf("(.%%2A_%s|.%%2A_%s|.%%2A_%s|.%%2A_%s)",
+			PT_BUILD_ROUTING_KEY, C2_RPC_START_SERVER_ROUTING_KEY, CONSUMING_CONTAINER_RESYNC_ROUTING_KEY,
+			CUSTOMBROWSER_EXPORT_FUNCTION))
 	go checkContainerStatusAddPT()
 	go checkContainerStatusAddC2()
 	go checkContainerStatusAddTR()
@@ -423,20 +430,20 @@ func checkContainerStatus() {
 		for container := range customBrowsersToCheck {
 			// check that a container is online
 			//logging.LogDebug("checking container", "container", container)
-			running := utils.SliceContains(existingQueues, GetCustomBrowserExportFunctionRoutingKey(consumingContainersToCheck[container].Name))
+			running := utils.SliceContains(existingQueues, GetCustomBrowserExportFunctionRoutingKey(customBrowsersToCheck[container].Name))
 			//logging.LogInfo("checking container running", "container", container, "running", running, "current_running", c2profilesToCheck[container].ContainerRunning)
-			if running != consumingContainersToCheck[container].ContainerRunning {
-				if entry, ok := consumingContainersToCheck[container]; ok {
+			if running != customBrowsersToCheck[container].ContainerRunning {
+				if entry, ok := customBrowsersToCheck[container]; ok {
 					entry.ContainerRunning = running
 					_, err = database.DB.NamedExec(`UPDATE custombrowser SET
 							container_running=:container_running, deleted=false
 							WHERE id=:id`, entry,
 					)
 					if err != nil {
-						logging.LogError(err, "Failed to set container running status", "container_running", consumingContainersToCheck[container].ContainerRunning, "container", container)
+						logging.LogError(err, "Failed to set container running status", "container_running", customBrowsersToCheck[container].ContainerRunning, "container", container)
 						continue
 					}
-					consumingContainersToCheck[container] = entry
+					customBrowsersToCheck[container] = entry
 					if !running {
 						SendAllOperationsMessage(
 							getDownContainerMessage(container),

mythic-docker/src/rabbitmq/eventing.go

@@ -1444,7 +1444,7 @@ func startEventStepInstanceActionInterceptResponse(eventStepInstance databaseStr
 }
 func restartFailedJobs(eventgroupInstanceID int) error {
 	_, err := database.DB.Exec(`UPDATE eventstepinstance 
-		SET status=$1, end_timestamp=$2 WHERE eventgroupinstance_id=$3 AND status=$4 OR status=$5`,
+		SET status=$1, end_timestamp=$2 WHERE eventgroupinstance_id=$3 AND status IN ($4, $5)`,
 		eventing.EventGroupInstanceStatusQueued, nil, eventgroupInstanceID,
 		eventing.EventGroupInstanceStatusError, eventing.EventGroupInstanceStatusCancelled)
 	if err != nil {

mythic-docker/src/rabbitmq/initialize.go

@@ -97,6 +97,7 @@ var pushC2StreamingDisconnectNotification = make(chan int, 100)
 func Initialize() {
 	RabbitMQConnection.channelMutexMap = make(map[string]*channelMutex)
 	invalidateAllNonUserAPITokens()
+	go startShortLivedAPITokenCleanup()
 	InvalidateOperationEventLogCacheMap()
 	go listenForOperationsMessages()
 	go listenForWriteDownloadChunkToLocalDisk()

mythic-docker/src/rabbitmq/recv_mythic_rpc_apitoken_create.go

@@ -175,19 +175,21 @@ func MythicRPCAPITokenCreate(input MythicRPCAPITokenCreateMessage) MythicRPCAPIT
 	accessToken, storedAPITokenValue, err := mythicjwt.GenerateOpaqueAPIToken()
 	if err != nil {
 		response.Error = err.Error()
+		expireAPIToken(apiToken.ID)
 		return response
 	}
 	apiToken.TokenValue = storedAPITokenValue
 	_, err = database.DB.Exec(`UPDATE apitokens SET token_value=$1 WHERE id=$2`, apiToken.TokenValue, apiToken.ID)
 	if err != nil {
 		response.Error = err.Error()
+		expireAPIToken(apiToken.ID)
 		return response
 	}
 	response.Success = true
 	response.APIToken = accessToken
 	if apiToken.TokenType == mythicjwt.AUTH_METHOD_CALLBACK || apiToken.TokenType == mythicjwt.AUTH_METHOD_PAYLOAD {
 		// deactivate the token after 5 min (should be a short-lived use)
-		go updateAPITokenAfter5Minutes(apiToken.ID, response.APIToken)
+		go expireAPITokenAfterShortLivedTTL(apiToken.ID)
 	}
 	return response
 }

mythic-docker/src/rabbitmq/recv_mythic_rpc_callback_search.go

@@ -121,7 +121,7 @@ func MythicRPCCallbackSearch(input MythicRPCCallbackSearchMessage) MythicRPCCall
 	if input.SearchCallbackDisplayID != nil || input.SearchCallbackID != nil || input.SearchCallbackUUID != nil {
 		searchString += ` AND (callback.id=:id OR
 			callback.agent_callback_id=:agent_callback_id OR
-			callback.display_id=:display_id)`
+			callback.display_id=:display_id) `
 	}
 	if input.SearchCallbackUser != nil {
 		targetCallback.User = *input.SearchCallbackUser
@@ -168,11 +168,11 @@ func MythicRPCCallbackSearch(input MythicRPCCallbackSearchMessage) MythicRPCCall
 	}
 	if input.SearchCallbackArchitecture != nil {
 		targetCallback.Architecture = *input.SearchCallbackArchitecture
-		searchString += `AND architecture ILIKE :architecture`
+		searchString += `AND architecture ILIKE :architecture `
 	}
 	if input.SearchCallbackDescription != nil {
 		targetCallback.Description = *input.SearchCallbackDescription
-		searchString += `AND callback.description ILIKE :description`
+		searchString += `AND callback.description ILIKE :description `
 	}
 	searchString += " ORDER BY callback.id DESC"
 	rows, err := database.DB.NamedQuery(searchString, targetCallback)

mythic-docker/src/rabbitmq/recv_mythic_rpc_custombrowser_search.go

@@ -115,13 +115,17 @@ func MythicRPCCustomBrowserSearch(input MythicRPCCustomBrowserSearchMessage) Myt
 
 	}
 	if input.SearchCustomBrowser.MetadataKey != nil {
-		searchString += "AND metadata->:metadata_key "
 		paramDict["metadata_key"] = *input.SearchCustomBrowser.MetadataKey
 		if input.SearchCustomBrowser.MetadataValue != nil {
-			searchString += " = :metadata_value "
-			paramDict["metadata_value"] = input.SearchCustomBrowser.MetadataValue
+			metadataValueBytes, err := json.Marshal(input.SearchCustomBrowser.MetadataValue)
+			if err != nil {
+				response.Error = err.Error()
+				return response
+			}
+			searchString += "AND metadata->:metadata_key = CAST(:metadata_value AS jsonb) "
+			paramDict["metadata_value"] = string(metadataValueBytes)
 		} else {
-			searchString += " is not null "
+			searchString += "AND metadata->:metadata_key IS NOT NULL "
 		}
 	}
 	if input.SearchCustomBrowser.CallbackGroup != nil {

mythic-docker/src/rabbitmq/recv_mythic_rpc_process_search.go

@@ -65,7 +65,7 @@ func MythicRPCProcessSearch(input MythicRPCProcessSearchMessage) MythicRPCProces
 			searchString += "AND host ILIKE :host "
 		}
 		if input.SearchProcess.ProcessID != nil {
-			paramDict["process_id"] = *input.SearchProcess.ProcessID
+			paramDict["process_id"] = fmt.Sprintf("%d", *input.SearchProcess.ProcessID)
 			searchString += "AND metadata->>'process_id'=:process_id "
 		}
 		if input.SearchProcess.Architecture != nil {
@@ -73,7 +73,7 @@ func MythicRPCProcessSearch(input MythicRPCProcessSearchMessage) MythicRPCProces
 			searchString += "AND metadata->>'architecture'=:architecture "
 		}
 		if input.SearchProcess.ParentProcessID != nil {
-			paramDict["parent_process_id"] = *input.SearchProcess.ParentProcessID
+			paramDict["parent_process_id"] = fmt.Sprintf("%d", *input.SearchProcess.ParentProcessID)
 			searchString += "AND metadata->>'parent_process_id'=:parent_process_id "
 		}
 		if input.SearchProcess.BinPath != nil {
@@ -86,14 +86,14 @@ func MythicRPCProcessSearch(input MythicRPCProcessSearchMessage) MythicRPCProces
 		}
 		if input.SearchProcess.User != nil {
 			paramDict["user"] = fmt.Sprintf("%%%s%%", *input.SearchProcess.User)
-			searchString += "AND metadata->>\"user\" ILIKE :user "
+			searchString += "AND metadata->>'user' ILIKE :user "
 		}
 		if input.SearchProcess.CommandLine != nil {
 			paramDict["command_line"] = fmt.Sprintf("%%%s%%", *input.SearchProcess.CommandLine)
 			searchString += "AND metadata->>'command_line' ILIKE :command_line "
 		}
 		if input.SearchProcess.IntegrityLevel != nil {
-			paramDict["integrity_level"] = *input.SearchProcess.IntegrityLevel
+			paramDict["integrity_level"] = fmt.Sprintf("%d", *input.SearchProcess.IntegrityLevel)
 			searchString += "AND metadata->>'integrity_level'=:integrity_level "
 		}
 		if input.SearchProcess.Description != nil {

mythic-docker/src/rabbitmq/recv_pt_payload_build_response.go

@@ -89,10 +89,7 @@ func processPayloadBuildResponse(msg amqp.Delivery) {
 		logging.LogError(err, "Failed to get payload from the database")
 		return
 	}
-	_, err = database.DB.Exec(`UPDATE apitokens SET deleted=true AND active=false WHERE payload_id=$1`, databasePayload.ID)
-	if err != nil {
-		logging.LogError(err, "Failed to update the apitokens to set to deleted")
-	}
+	expireAPITokensForPayload(databasePayload.ID)
 	databasePayload.BuildMessage += payloadBuildResponse.BuildMessage
 	databasePayload.BuildStderr += payloadBuildResponse.BuildStdErr
 	databasePayload.BuildStdout += payloadBuildResponse.BuildStdOut

mythic-docker/src/rabbitmq/recv_pt_task_create_response.go

@@ -3,10 +3,11 @@ package rabbitmq
 import (
 	"encoding/json"
 	"fmt"
-	"github.com/its-a-feature/Mythic/eventing"
 	"strings"
 	"time"
 
+	"github.com/its-a-feature/Mythic/eventing"
+
 	"github.com/its-a-feature/Mythic/database"
 	databaseStructs "github.com/its-a-feature/Mythic/database/structs"
 	"github.com/its-a-feature/Mythic/logging"
@@ -43,10 +44,6 @@ func processPtTaskCreateMessages(msg amqp.Delivery) {
 		go SendAllOperationsMessage(err.Error(), 0, "", database.MESSAGE_LEVEL_INFO, true)
 		return
 	}
-	_, err = database.DB.Exec(`UPDATE apitokens SET deleted=true AND active=false WHERE task_id=$1`, task.ID)
-	if err != nil {
-		logging.LogError(err, "Failed to update the apitokens to set to deleted")
-	}
 	//logging.LogInfo("got response back from create message", "resp", payloadMsg, "original", string(msg.Body))
 
 	var updateColumns []string
@@ -107,6 +104,7 @@ func processPtTaskCreateMessages(msg amqp.Delivery) {
 		} else {
 			task.Status = PT_TASK_FUNCTION_STATUS_COMPLETED
 		}
+		expireAPITokensForTask(task.ID)
 		task.Timestamp = time.Now().UTC()
 		updateColumns = append(updateColumns, "timestamp=:timestamp")
 		task.StatusTimestampSubmitted.Valid = true

mythic-docker/src/rabbitmq/recv_pt_task_opsec_post_response.go

@@ -43,10 +43,7 @@ func processPtTaskOPSECPostMessages(msg amqp.Delivery) {
 		go SendAllOperationsMessage(err.Error(), 0, "", database.MESSAGE_LEVEL_INFO, true)
 		return
 	}
-	_, err = database.DB.Exec(`UPDATE apitokens SET deleted=true AND active=false WHERE task_id=$1`, task.ID)
-	if err != nil {
-		logging.LogError(err, "Failed to update the apitokens to set to deleted")
-	}
+	expireAPITokensForTask(task.ID)
 	if payloadMsg.Success {
 		if !payloadMsg.OpsecPostBlocked || (payloadMsg.OpsecPostBlocked && payloadMsg.OpsecPostBypassed != nil && *payloadMsg.OpsecPostBypassed) {
 			task.Status = PT_TASK_FUNCTION_STATUS_SUBMITTED

mythic-docker/src/rabbitmq/recv_pt_task_opsec_pre_response.go

@@ -38,10 +38,7 @@ func processPtTaskOPSECPreMessages(msg amqp.Delivery) {
 			go SendAllOperationsMessage(err.Error(), 0, "", database.MESSAGE_LEVEL_INFO, true)
 			return
 		}
-		_, err = database.DB.Exec(`UPDATE apitokens SET deleted=true AND active=false WHERE task_id=$1`, task.ID)
-		if err != nil {
-			logging.LogError(err, "Failed to update the apitokens to set to deleted")
-		}
+		expireAPITokensForTask(task.ID)
 		if payloadMsg.Success {
 			shouldMoveToCreateTasking := false
 			if !payloadMsg.OpsecPreBlocked || (payloadMsg.OpsecPreBlocked && payloadMsg.OpsecPreBypassed != nil && *payloadMsg.OpsecPreBypassed) {

mythic-docker/src/rabbitmq/recv_pt_task_process_response.go

@@ -24,10 +24,7 @@ func processPtTaskProcessResponseMessages(msg amqp.Delivery) {
 		logging.LogError(err, "Failed to process PTTaskProcessResponseMessageResponse into struct")
 	} else {
 		// now process the create_tasking response body to update the task
-		_, err = database.DB.Exec(`UPDATE apitokens SET deleted=true AND active=false WHERE task_id=$1`, payloadMsg.TaskID)
-		if err != nil {
-			logging.LogError(err, "Failed to update the apitokens to set to deleted")
-		}
+		expireAPITokensForTask(payloadMsg.TaskID)
 		if !payloadMsg.Success {
 			go SendAllOperationsMessage(fmt.Sprintf("Failed to process response message for task %d:\n%s", payloadMsg.TaskID, payloadMsg.Error),
 				0, "", database.MESSAGE_LEVEL_INFO, true)

mythic-docker/src/rabbitmq/utils.go

@@ -852,10 +852,7 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) {
 	if !task.Completed {
 		return
 	}
-	_, err = database.DB.Exec(`UPDATE apitokens SET deleted=true AND active=false WHERE task_id=$1`, taskId)
-	if err != nil {
-		logging.LogError(err, "Failed to update the apitokens to set to deleted")
-	}
+	expireAPITokensForTask(taskId)
 	if task.ParentTaskID.Valid {
 		err = database.DB.Get(&parentTask, `SELECT 
 		task.id, task.status, task.completed, task.eventstepinstance_id, task.completed_callback_function, task.completed_callback_function_completed,

mythic-docker/src/rabbitmq/utils_api_tokens.go

@@ -0,0 +1,84 @@
+package rabbitmq
+
+import (
+	"time"
+
+	"github.com/its-a-feature/Mythic/authentication/mythicjwt"
+	"github.com/its-a-feature/Mythic/database"
+	"github.com/its-a-feature/Mythic/logging"
+)
+
+const (
+	shortLivedAPITokenTTL             = 5 * time.Minute
+	shortLivedAPITokenCleanupInterval = 1 * time.Minute
+)
+
+// expireAPIToken marks a single token unusable. The active/deleted guard keeps
+// repeated cleanup paths cheap and idempotent when a timer and sweeper overlap.
+func expireAPIToken(apitokenID int) {
+	_, err := database.DB.Exec(`UPDATE apitokens
+		SET active=false, deleted=true
+		WHERE id=$1 AND (active=true OR deleted=false)`, apitokenID)
+	if err != nil {
+		logging.LogError(err, "failed to mark apitoken as deleted", "apitoken_id", apitokenID)
+	}
+}
+
+// expireAPITokenAfterShortLivedTTL handles the best-case in-process expiry for
+// temporary tokens. A periodic sweeper also covers process restarts and failures
+// that happen after token creation but before this timer fires.
+func expireAPITokenAfterShortLivedTTL(apitokenID int) {
+	timer := time.NewTimer(shortLivedAPITokenTTL)
+	defer timer.Stop()
+	<-timer.C
+	expireAPIToken(apitokenID)
+}
+
+// expireAPITokensForTask retires all task-scoped tokens when the container-side
+// task function is done with them. Task tokens are lifecycle-driven, not TTL-driven.
+func expireAPITokensForTask(taskID int) {
+	_, err := database.DB.Exec(`UPDATE apitokens
+		SET active=false, deleted=true
+		WHERE task_id=$1 AND (active=true OR deleted=false)`, taskID)
+	if err != nil {
+		logging.LogError(err, "Failed to update the task apitokens to set to deleted", "task_id", taskID)
+	}
+}
+
+// expireAPITokensForPayload retires payload-scoped tokens as soon as the payload
+// build lifecycle finishes. The short-lived sweeper is only a fallback.
+func expireAPITokensForPayload(payloadID int) {
+	_, err := database.DB.Exec(`UPDATE apitokens
+		SET active=false, deleted=true
+		WHERE payload_id=$1 AND (active=true OR deleted=false)`, payloadID)
+	if err != nil {
+		logging.LogError(err, "Failed to update the payload apitokens to set to deleted", "payload_id", payloadID)
+	}
+}
+
+// startShortLivedAPITokenCleanup continuously cleans TTL-based tokens that lost
+// their original in-process timer because Mythic restarted or token creation failed mid-flow.
+func startShortLivedAPITokenCleanup() {
+	cleanupExpiredShortLivedAPITokens()
+	ticker := time.NewTicker(shortLivedAPITokenCleanupInterval)
+	defer ticker.Stop()
+	for range ticker.C {
+		cleanupExpiredShortLivedAPITokens()
+	}
+}
+
+func cleanupExpiredShortLivedAPITokens() {
+	cutoff := time.Now().UTC().Add(-shortLivedAPITokenTTL)
+	_, err := database.DB.Exec(`UPDATE apitokens
+		SET active=false, deleted=true
+		WHERE (active=true OR deleted=false)
+		AND creation_time < $1
+		AND token_type IN ($2, $3, $4)`,
+		cutoff,
+		mythicjwt.AUTH_METHOD_ON_START,
+		mythicjwt.AUTH_METHOD_CALLBACK,
+		mythicjwt.AUTH_METHOD_PAYLOAD)
+	if err != nil {
+		logging.LogError(err, "failed to clean expired short-lived apitokens")
+	}
+}