117 lines
2.9 KiB
Go
117 lines
2.9 KiB
Go
package server
|
|
|
|
import "github.com/Adaptix-Framework/axc2"
|
|
|
|
type TaskTaskHandler struct{}
|
|
|
|
func (h *TaskTaskHandler) Create(tm *TaskManager, agent *Agent, taskData *adaptix.TaskData) {
|
|
if taskData.Sync {
|
|
tm.syncTaskCreate(taskData.AgentId, agent, taskData)
|
|
}
|
|
agent.HostedTasks.Push(*taskData)
|
|
}
|
|
|
|
func (h *TaskTaskHandler) Update(tm *TaskManager, agent *Agent, task *adaptix.TaskData, updateData *adaptix.TaskData) {
|
|
agent.RunningTasks.Delete(updateData.TaskId)
|
|
|
|
task.FinishDate = updateData.FinishDate
|
|
task.Completed = updateData.Completed
|
|
task.MessageType = updateData.MessageType
|
|
task.Message = updateData.Message
|
|
task.ClearText = updateData.ClearText
|
|
|
|
if task.HookId != "" && tm.ts.TsAxScriptIsServerHook(task.HookId) {
|
|
hookData := map[string]interface{}{
|
|
"agent": task.AgentId,
|
|
"task_id": task.TaskId,
|
|
"message": task.Message,
|
|
"text": task.ClearText,
|
|
"type": task.MessageType,
|
|
"completed": task.Completed,
|
|
}
|
|
result, _ := tm.ts.TsAxScriptExecPostHook(task.HookId, hookData)
|
|
if result != nil {
|
|
if msg, ok := result["message"].(string); ok {
|
|
task.Message = msg
|
|
}
|
|
if txt, ok := result["text"].(string); ok {
|
|
task.ClearText = txt
|
|
}
|
|
if mt, ok := result["type"].(int); ok {
|
|
task.MessageType = mt
|
|
}
|
|
}
|
|
|
|
if task.Sync {
|
|
if task.Completed {
|
|
tm.ts.TsAxScriptRemovePostHook(task.HookId)
|
|
task.HookId = ""
|
|
tm.completeTask(agent, task)
|
|
tm.executeServerHandler(task)
|
|
} else {
|
|
agent.RunningTasks.Put(task.TaskId, *task)
|
|
}
|
|
tm.syncTaskUpdate(task.AgentId, agent, task)
|
|
}
|
|
return
|
|
}
|
|
|
|
if task.HookId != "" && task.Client != "" && tm.ts.TsClientConnected(task.Client) {
|
|
agent.RunningTasks.Put(task.TaskId, *task)
|
|
|
|
packet := CreateSpAgentTaskHook(*task, 0)
|
|
tm.ts.TsSyncClient(task.Client, packet)
|
|
return
|
|
}
|
|
|
|
if task.Sync {
|
|
if task.Completed {
|
|
tm.completeTask(agent, task)
|
|
tm.executeServerHandler(task)
|
|
updateData.HandlerId = task.HandlerId
|
|
} else {
|
|
agent.RunningTasks.Put(task.TaskId, *task)
|
|
}
|
|
|
|
tm.syncTaskUpdate(task.AgentId, agent, task)
|
|
}
|
|
}
|
|
|
|
func (h *TaskTaskHandler) PostHook(tm *TaskManager, agent *Agent, task *adaptix.TaskData, hookData *adaptix.TaskData, jobIndex int) error {
|
|
_, ok := agent.RunningTasks.GetDelete(hookData.TaskId)
|
|
if !ok {
|
|
return nil
|
|
}
|
|
|
|
task.MessageType = hookData.MessageType
|
|
task.Message = hookData.Message
|
|
task.ClearText = hookData.ClearText
|
|
|
|
if task.Sync {
|
|
if task.Completed {
|
|
task.HookId = ""
|
|
tm.completeTask(agent, task)
|
|
} else {
|
|
agent.RunningTasks.Put(task.TaskId, *task)
|
|
}
|
|
|
|
tm.syncTaskUpdate(task.AgentId, agent, task)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (h *TaskTaskHandler) OnClientDisconnect(tm *TaskManager, agent *Agent, task *adaptix.TaskData, clientName string) {
|
|
if !task.Sync {
|
|
return
|
|
}
|
|
|
|
if task.Completed {
|
|
task.HookId = ""
|
|
agent.RunningTasks.Delete(task.TaskId)
|
|
tm.completeTask(agent, task)
|
|
}
|
|
|
|
tm.syncTaskUpdate(task.AgentId, agent, task)
|
|
}
|