better support chat cancel

mythic-docker/src/rabbitmq/constants.go

@@ -128,6 +128,7 @@ const (
 	PT_TASK_PROCESS_RESPONSE = "pt_task_process_response"
 	PT_COMMAND_HELP_FUNCTION = "pt_command_help_function"
 	CHAT_REQUEST             = "chat_request"
+	CHAT_CANCEL              = "chat_cancel"
 )
 
 // Routes where container is consuming messages and responding back to Mythic

mythic-docker/src/rabbitmq/recv_chat_response.go

@@ -290,7 +290,7 @@ func setChatResponseStatus(request chatResponseRequest, status string, responseE
 	_, err := database.DB.Exec(`UPDATE chat_message
 		SET status=$3,
 			metadata = metadata || $4::jsonb
-		WHERE id=$1 AND operation_id=$2 AND deleted=false`, request.ResponseMessageID, request.OperationID, status, metadataJSON.String())
+		WHERE id=$1 AND operation_id=$2 AND deleted=false AND status <> 'cancelled'`, request.ResponseMessageID, request.OperationID, status, metadataJSON.String())
 	if err != nil {
 		return err
 	}

mythic-docker/src/rabbitmq/send_chat_request.go

@@ -31,6 +31,16 @@ type ChatContainerRequestMessage struct {
 	Secrets           map[string]interface{}        `json:"secrets" mapstructure:"secrets"`
 }
 
+type ChatContainerCancelRequestMessage struct {
+	ContainerName     string `json:"container_name" mapstructure:"container_name"`
+	OperationID       int    `json:"operation_id" mapstructure:"operation_id"`
+	ChannelID         int    `json:"channel_id" mapstructure:"channel_id"`
+	RequestID         int    `json:"request_id" mapstructure:"request_id"`
+	ResponseMessageID int    `json:"response_message_id" mapstructure:"response_message_id"`
+	Reason            string `json:"reason" mapstructure:"reason"`
+	CancelledBy       int    `json:"cancelled_by" mapstructure:"cancelled_by"`
+}
+
 func (r *rabbitMQConnection) SendChatContainerRequest(containerName string, chatMessage ChatContainerRequestMessage, authContext RabbitMQAuthContext) error {
 	chatMessage.ContainerName = containerName
 	headers, err := GenerateRabbitMQAuthTokenHeader(authContext)
@@ -52,3 +62,25 @@ func (r *rabbitMQConnection) SendChatContainerRequest(containerName string, chat
 	}
 	return nil
 }
+
+func (r *rabbitMQConnection) SendChatContainerCancelRequest(containerName string, cancelMessage ChatContainerCancelRequestMessage, authContext RabbitMQAuthContext) error {
+	cancelMessage.ContainerName = containerName
+	headers, err := GenerateRabbitMQAuthTokenHeader(authContext)
+	if err != nil {
+		logging.LogError(err, "Failed to generate auth context for chat cancellation")
+		return err
+	}
+	err = r.SendStructMessage(
+		MYTHIC_EXCHANGE,
+		GetChatContainerCancelRoutingKey(containerName),
+		"",
+		cancelMessage,
+		false,
+		headers,
+	)
+	if err != nil {
+		logging.LogError(err, "Failed to send chat cancellation", "container", containerName, "request_id", cancelMessage.RequestID)
+		return err
+	}
+	return nil
+}

mythic-docker/src/rabbitmq/utils_rabbitmq_routing.go

@@ -62,6 +62,9 @@ func GetPtCommandHelpRoutingKey(container string) string {
 func GetChatContainerRequestRoutingKey(container string) string {
 	return fmt.Sprintf("%s_%s", container, CHAT_REQUEST)
 }
+func GetChatContainerCancelRoutingKey(container string) string {
+	return fmt.Sprintf("%s_%s", container, CHAT_CANCEL)
+}
 
 // c2 rpc routing key functions
 func GetC2RPCOpsecChecksRoutingKey(container string) string {

mythic-docker/src/webserver/controllers/chat_update_webhook.go

@@ -550,8 +550,8 @@ func CancelChatRequestWebhook(c *gin.Context) {
 		chatRespondError(c, "only the request creator or an operation admin can cancel this request")
 		return
 	}
-	_, err = database.DB.Exec(`UPDATE chat_request
-		SET status='cancelled', cancelled_at=now(), error='Cancelled by operator'
+	result, err := database.DB.Exec(`UPDATE chat_request
+		SET status='cancelled', cancelled_at=now(), error='Operator issued cancel'
 		WHERE id=$1 AND operation_id=$2 AND status IN ('pending', 'streaming')`,
 		request.ID, operatorOperation.CurrentOperation.ID)
 	if err != nil {
@@ -559,10 +559,38 @@ func CancelChatRequestWebhook(c *gin.Context) {
 		chatRespondError(c, err.Error())
 		return
 	}
+	rowsAffected, err := result.RowsAffected()
+	if err != nil {
+		logging.LogError(err, "Failed to check cancelled chat request rows affected")
+		rowsAffected = 0
+	}
 	_, _ = database.DB.Exec(`UPDATE chat_message
 		SET status='cancelled'
 		WHERE id=$1 AND operation_id=$2 AND status IN ('pending', 'streaming')`,
 		request.ResponseMessageID, operatorOperation.CurrentOperation.ID)
+	if rowsAffected > 0 {
+		container := databaseStructs.ConsumingContainer{}
+		if err = database.DB.Get(&container, `SELECT *
+			FROM consuming_container
+			WHERE id=$1`, request.ChatContainerID); err != nil {
+			logging.LogError(err, "Failed to find chat container for cancellation", "request_id", request.ID)
+			chatRespondError(c, "cancelled in Mythic, but failed to find the chat container to notify")
+			return
+		}
+		authContext := authentication.RabbitMQAuthContextFromGin(c)
+		if err = rabbitmq.RabbitMQConnection.SendChatContainerCancelRequest(container.Name, rabbitmq.ChatContainerCancelRequestMessage{
+			OperationID:       operatorOperation.CurrentOperation.ID,
+			ChannelID:         channel.ID,
+			RequestID:         request.ID,
+			ResponseMessageID: request.ResponseMessageID,
+			Reason:            "Cancelled by operator",
+			CancelledBy:       operatorOperation.CurrentOperator.ID,
+		}, authContext); err != nil {
+			logging.LogError(err, "Failed to send chat cancellation", "request_id", request.ID, "container", container.Name)
+			chatRespondError(c, "cancelled in Mythic, but failed to notify the chat container")
+			return
+		}
+	}
 	c.JSON(http.StatusOK, ChatActionResponse{Status: "success", RequestID: request.ID, ResponseMessageID: request.ResponseMessageID, ChannelID: channel.ID})
 }