From bbfcb1936601955ea0cef084015877ba9fab19cc Mon Sep 17 00:00:00 2001 From: its-a-feature Date: Wed, 5 Mar 2025 13:56:57 -0600 Subject: [PATCH] updating rabbitmq looping to add more retries --- CHANGELOG.MD | 8 + VERSION | 2 +- mythic-docker/src/VERSION | 2 +- mythic-docker/src/rabbitmq/structs.go | 10 +- .../src/rabbitmq/util_create_task.go | 24 +- .../src/rabbitmq/utils_rabbitmq_routing.go | 210 +++++++++--------- 6 files changed, 136 insertions(+), 120 deletions(-) diff --git a/CHANGELOG.MD b/CHANGELOG.MD index b486a0bb..bad74199 100644 --- a/CHANGELOG.MD +++ b/CHANGELOG.MD @@ -4,6 +4,14 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [3.3.1-rc48] - 2025-03-05 + +### Changed + +- Updated the create tasking webhook to return sooner before sending to containers +- Updated some of the rabbitmq connections with loops to be a bit more robust to errors +- Updated some of the tasking status messages to try to be a bit more informative + ## [3.3.1-rc47] - 2025-02-26 ### Changed diff --git a/VERSION b/VERSION index 28fbc353..3e63ed37 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -3.3.1-rc47 \ No newline at end of file +3.3.1-rc48 \ No newline at end of file diff --git a/mythic-docker/src/VERSION b/mythic-docker/src/VERSION index 28fbc353..3e63ed37 100644 --- a/mythic-docker/src/VERSION +++ b/mythic-docker/src/VERSION @@ -1 +1 @@ -3.3.1-rc47 \ No newline at end of file +3.3.1-rc48 \ No newline at end of file diff --git a/mythic-docker/src/rabbitmq/structs.go b/mythic-docker/src/rabbitmq/structs.go index fd672ec2..77efd85b 100644 --- a/mythic-docker/src/rabbitmq/structs.go +++ b/mythic-docker/src/rabbitmq/structs.go @@ -152,12 +152,12 @@ type PT_TASK_FUNCTION_STATUS = string const ( PT_TASK_FUNCTION_STATUS_OPSEC_PRE PT_TASK_FUNCTION_STATUS = "OPSEC Pre Check Running..." - PT_TASK_FUNCTION_STATUS_OPSEC_PRE_ERROR = "Error: processing arguments - click cog to check stderr" + PT_TASK_FUNCTION_STATUS_OPSEC_PRE_ERROR = "Error: processing arguments" PT_TASK_FUNCTION_STATUS_OPSEC_PRE_BLOCKED = "OPSEC Pre Blocked" - PT_TASK_FUNCTION_STATUS_PREPROCESSING = "creating task..." - PT_TASK_FUNCTION_STATUS_PREPROCESSING_ERROR = "Error: creating task - click cog to check stderr" + PT_TASK_FUNCTION_STATUS_PREPROCESSING = "preparing task for agent..." + PT_TASK_FUNCTION_STATUS_PREPROCESSING_ERROR = "Error: creating task " PT_TASK_FUNCTION_STATUS_OPSEC_POST = "OPSEC Post Check Running..." - PT_TASK_FUNCTION_STATUS_OPSEC_POST_ERROR = "Error: opsec check - click cog to check stderr" + PT_TASK_FUNCTION_STATUS_OPSEC_POST_ERROR = "Error: opsec check failed to run" PT_TASK_FUNCTION_STATUS_OPSEC_POST_BLOCKED = "OPSEC Post Blocked" PT_TASK_FUNCTION_STATUS_SUBMITTED = "submitted" PT_TASK_FUNCTION_STATUS_PROCESSING = "agent processing" @@ -169,7 +169,7 @@ const ( PT_TASK_FUNCTION_STATUS_GROUP_COMPLETED_FUNCTION = "Group Completion Function Running..." PT_TASK_FUNCTION_STATUS_GROUP_COMPLETED_FUNCTION_ERROR = "Error: group completion function - click cog to check stderr" PT_TASK_FUNCTION_STATUS_COMPLETED = "success" - PT_TASK_FUNCTION_STATUS_PROCESSED = "processed, waiting for more messages..." + PT_TASK_FUNCTION_STATUS_PROCESSED = "processed, agent sending responses..." PT_TASK_FUNCTION_STATUS_INTERCEPTED = "intercepted for custom checks" PT_TASK_FUNCTION_STATUS_INTERCEPTED_ERROR = "Error: Task Interception Failed" ) diff --git a/mythic-docker/src/rabbitmq/util_create_task.go b/mythic-docker/src/rabbitmq/util_create_task.go index 7a1e8e8b..1422c3d8 100644 --- a/mythic-docker/src/rabbitmq/util_create_task.go +++ b/mythic-docker/src/rabbitmq/util_create_task.go @@ -516,7 +516,12 @@ func CreateTask(createTaskInput CreateTaskInput) CreateTaskResponse { TaskDisplayID: task.DisplayID, } } - return submitTaskToContainer(task.ID) + go submitTaskToContainer(task.ID) + return CreateTaskResponse{ + Status: "success", + TaskID: task.ID, + TaskDisplayID: task.DisplayID, + } } func associateUploadedFilesWithTask(task *databaseStructs.Task, files []string) { @@ -804,24 +809,15 @@ func handleHelpCommand(createTaskInput CreateTaskInput, callback databaseStructs } -func submitTaskToContainer(taskID int) CreateTaskResponse { - output := CreateTaskResponse{ - Status: "error", - Error: "not implemented", - } +func submitTaskToContainer(taskID int) { taskMessage := GetTaskConfigurationForContainer(taskID) - if err := RabbitMQConnection.SendPtTaskOPSECPre(taskMessage); err != nil { + err := RabbitMQConnection.SendPtTaskOPSECPre(taskMessage) + if err != nil { logging.LogError(err, "Failed to send task to payload type") - output.Error = err.Error() if _, err := database.DB.Exec(`UPDATE task SET status=$1 WHERE id=$2`, TASK_STATUS_CONTAINER_DOWN, taskID); err != nil { logging.LogError(err, "Failed to update task status") } - } else { - output.Status = "success" - output.Error = "" - output.TaskID = taskID - output.TaskDisplayID = taskMessage.Task.DisplayID } - return output + return } diff --git a/mythic-docker/src/rabbitmq/utils_rabbitmq_routing.go b/mythic-docker/src/rabbitmq/utils_rabbitmq_routing.go index ab958e19..496408b3 100644 --- a/mythic-docker/src/rabbitmq/utils_rabbitmq_routing.go +++ b/mythic-docker/src/rabbitmq/utils_rabbitmq_routing.go @@ -150,29 +150,28 @@ func (r *rabbitMQConnection) GetConnection() (*amqp.Connection, error) { defer r.mutex.Unlock() if r.conn != nil && !r.conn.IsClosed() { return r.conn, nil - } else { - for { - logging.LogInfo("Attempting to connect to rabbitmq") - conn, err := amqp.DialConfig(fmt.Sprintf("amqp://%s:%s@%s:%d/%s", - utils.MythicConfig.RabbitmqUser, - utils.MythicConfig.RabbitmqPassword, - utils.MythicConfig.RabbitmqHost, - utils.MythicConfig.RabbitmqPort, - utils.MythicConfig.RabbitmqVHost), - amqp.Config{ - Dial: func(network, addr string) (net.Conn, error) { - return net.DialTimeout(network, addr, 10*time.Second) - }, + } + for { + logging.LogInfo("Attempting to connect to rabbitmq") + conn, err := amqp.DialConfig(fmt.Sprintf("amqp://%s:%s@%s:%d/%s", + utils.MythicConfig.RabbitmqUser, + utils.MythicConfig.RabbitmqPassword, + utils.MythicConfig.RabbitmqHost, + utils.MythicConfig.RabbitmqPort, + utils.MythicConfig.RabbitmqVHost), + amqp.Config{ + Dial: func(network, addr string) (net.Conn, error) { + return net.DialTimeout(network, addr, 10*time.Second) }, - ) - if err != nil { - logging.LogError(err, "Failed to connect to rabbitmq") - time.Sleep(RETRY_CONNECT_DELAY) - continue - } - r.conn = conn - return conn, nil + }, + ) + if err != nil { + logging.LogError(err, "Failed to connect to rabbitmq") + time.Sleep(RETRY_CONNECT_DELAY) + continue } + r.conn = conn + return conn, nil } } func (r *rabbitMQConnection) SendStructMessage(exchange string, queue string, correlationId string, body interface{}, ignoreErrorMessage bool) error { @@ -255,7 +254,6 @@ func (r *rabbitMQConnection) SendMessage(exchange string, queue string, correlat } logging.LogError(err, "failed 3 times") return err - } func (r *rabbitMQConnection) SendRPCMessage(exchange string, queue string, body []byte, exclusiveQueue bool) ([]byte, error) { conn, err := r.GetConnection() @@ -358,82 +356,6 @@ func (r *rabbitMQConnection) ReceiveFromMythicDirectExchange(exchange string, qu // queue is where the messages get sent to (local name) // routingKey is the specific direct topic we're interested in for the exchange // handler processes the messages we get on our queue - for { - if conn, err := r.GetConnection(); err != nil { - logging.LogError(err, "Failed to connect to rabbitmq", "retry_wait_time", RETRY_CONNECT_DELAY) - time.Sleep(RETRY_CONNECT_DELAY) - continue - } else if ch, err := conn.Channel(); err != nil { - logging.LogError(err, "Failed to open rabbitmq channel", "retry_wait_time", RETRY_CONNECT_DELAY) - time.Sleep(RETRY_CONNECT_DELAY) - continue - } else if err = ch.ExchangeDeclare( - exchange, // exchange name - "direct", // type of exchange, ex: topic, fanout, direct, etc - true, // durable - true, // auto-deleted - false, // internal - false, // no-wait - nil, // arguments - ); err != nil { - logging.LogError(err, "Failed to declare exchange", "exchange", exchange, "exchange_type", "direct", "retry_wait_time", RETRY_CONNECT_DELAY) - time.Sleep(RETRY_CONNECT_DELAY) - continue - } else if q, err := ch.QueueDeclare( - queue, // name, queue - false, // durable - true, // delete when unused - exclusiveQueue, // exclusive - false, // no-wait - nil, // arguments - ); err != nil { - logging.LogError(err, "Failed to declare queue", "retry_wait_time", RETRY_CONNECT_DELAY) - ch.Close() - time.Sleep(RETRY_CONNECT_DELAY) - continue - } else if err = ch.QueueBind( - q.Name, // queue name - routingKey, // routing key - exchange, // exchange name - false, // nowait - nil, // arguments - ); err != nil { - logging.LogError(err, "Failed to bind to queue to receive messages", "retry_wait_time", RETRY_CONNECT_DELAY) - ch.Close() - time.Sleep(RETRY_CONNECT_DELAY) - continue - } else if msgs, err := ch.Consume( - q.Name, // queue name - "", // consumer - false, // auto-ack - false, // exclusive - false, // no local - false, // no wait - nil, // args - ); err != nil { - logging.LogError(err, "Failed to start consuming on queue", "queue", q.Name) - ch.Close() - } else { - forever := make(chan bool) - go func() { - for d := range msgs { - //logging.LogDebug("got direct message", "queue", q.Name, "msg", d.Body) - go handler(d) - if err = ch.Ack(d.DeliveryTag, false); err != nil { - logging.LogError(err, "Failed to Ack message") - } - } - forever <- true - }() - logging.LogInfo("Started listening for messages", "exchange", exchange, "queue", queue, "routingKey", routingKey) - <-forever - ch.Close() - logging.LogError(nil, "Stopped listening for messages", "exchange", exchange, "queue", queue, "routingKey", routingKey) - } - - } -} -func (r *rabbitMQConnection) ReceiveFromRPCQueue(exchange string, queue string, routingKey string, handler RPCQueueHandler, exclusiveQueue bool) { for { conn, err := r.GetConnection() if err != nil { @@ -484,8 +406,97 @@ func (r *rabbitMQConnection) ReceiveFromRPCQueue(exchange string, queue string, ) if err != nil { logging.LogError(err, "Failed to bind to queue to receive messages", "retry_wait_time", RETRY_CONNECT_DELAY) - time.Sleep(RETRY_CONNECT_DELAY) ch.Close() + time.Sleep(RETRY_CONNECT_DELAY) + continue + } + msgs, err := ch.Consume( + q.Name, // queue name + "", // consumer + false, // auto-ack + false, // exclusive + false, // no local + false, // no wait + nil, // args + ) + if err != nil { + logging.LogError(err, "Failed to start consuming on queue", "queue", q.Name) + ch.Close() + time.Sleep(RETRY_CONNECT_DELAY) + continue + } + forever := make(chan bool) + go func() { + for d := range msgs { + //logging.LogDebug("got direct message", "queue", q.Name, "msg", d.Body) + err = ch.Ack(d.DeliveryTag, false) + if err != nil { + logging.LogError(err, "Failed to Ack message") + } + go handler(d) + } + forever <- true + }() + logging.LogInfo("Started listening for messages", "exchange", exchange, "queue", queue, "routingKey", routingKey) + <-forever + ch.Close() + logging.LogError(nil, "Stopped listening for messages", "exchange", exchange, "queue", queue, "routingKey", routingKey) + } +} +func (r *rabbitMQConnection) ReceiveFromRPCQueue(exchange string, queue string, routingKey string, handler RPCQueueHandler, exclusiveQueue bool) { + for { + conn, err := r.GetConnection() + if err != nil { + logging.LogError(err, "Failed to connect to rabbitmq", "retry_wait_time", RETRY_CONNECT_DELAY) + time.Sleep(RETRY_CONNECT_DELAY) + continue + } + ch, err := conn.Channel() + if err != nil { + logging.LogError(err, "Failed to open rabbitmq channel", "retry_wait_time", RETRY_CONNECT_DELAY) + time.Sleep(RETRY_CONNECT_DELAY) + continue + } + err = ch.ExchangeDeclare( + exchange, // exchange name + "direct", // type of exchange, ex: topic, fanout, direct, etc + true, // durable + true, // auto-deleted + false, // internal + false, // no-wait + nil, // arguments + ) + if err != nil { + logging.LogError(err, "Failed to declare exchange", "exchange", exchange, "exchange_type", "direct", "retry_wait_time", RETRY_CONNECT_DELAY) + ch.Close() + time.Sleep(RETRY_CONNECT_DELAY) + continue + } + q, err := ch.QueueDeclare( + queue, // name, queue + false, // durable + true, // delete when unused + exclusiveQueue, // exclusive + false, // no-wait + nil, // arguments + ) + if err != nil { + logging.LogError(err, "Failed to declare queue", "retry_wait_time", RETRY_CONNECT_DELAY) + ch.Close() + time.Sleep(RETRY_CONNECT_DELAY) + continue + } + err = ch.QueueBind( + q.Name, // queue name + routingKey, // routing key + exchange, // exchange name + false, // nowait + nil, // arguments + ) + if err != nil { + logging.LogError(err, "Failed to bind to queue to receive messages", "retry_wait_time", RETRY_CONNECT_DELAY) + ch.Close() + time.Sleep(RETRY_CONNECT_DELAY) continue } msgs, err := ch.Consume( @@ -500,6 +511,7 @@ func (r *rabbitMQConnection) ReceiveFromRPCQueue(exchange string, queue string, if err != nil { logging.LogError(err, "Failed to start consuming messages on queue", "queue", q.Name) ch.Close() + time.Sleep(RETRY_CONNECT_DELAY) continue } forever := make(chan bool)