mirror of
https://github.com/its-a-feature/Mythic
synced 2026-06-08 14:55:38 +00:00
updating rabbitmq looping to add more retries
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -1 +1 @@
|
||||
3.3.1-rc47
|
||||
3.3.1-rc48
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user