Adaptix-Framework-AdaptixC2/AdaptixServer/core/server/ts_tasks.go

package server

import (
	"AdaptixServer/core/utils/krypt"
	"AdaptixServer/core/utils/logs"
	"AdaptixServer/core/utils/safe"
	"fmt"
	"github.com/Adaptix-Framework/axc2"
	"time"
)

func (ts *Teamserver) TsTaskRunningExists(agentId string, taskId string) bool {
	value, ok := ts.agents.Get(agentId)
	if !ok {
		logs.Error("", "TsTaskUpdate: agent %v not found", agentId)
		return false
	}
	agent, _ := value.(*Agent)

	return agent.RunningTasks.Contains(taskId)
}

func (ts *Teamserver) TsTaskCreate(agentId string, cmdline string, client string, taskData adaptix.TaskData) {
	value, ok := ts.agents.Get(agentId)
	if !ok {
		logs.Error("", "TsTaskCreate: agent %v not found", agentId)
		return
	}

	agent, _ := value.(*Agent)
	if agent.Active == false {
		return
	}

	if taskData.TaskId == "" {
		taskData.TaskId, _ = krypt.GenerateUID(8)
	}
	taskData.AgentId = agentId
	taskData.CommandLine = cmdline
	taskData.Client = client
	taskData.Computer = agent.Data.Computer
	taskData.StartDate = time.Now().Unix()
	if taskData.Completed {
		taskData.FinishDate = taskData.StartDate
	}

	taskData.User = agent.Data.Username
	if agent.Data.Impersonated != "" {
		taskData.User += fmt.Sprintf(" [%s]", agent.Data.Impersonated)
	}

	switch taskData.Type {

	case TYPE_TASK:
		if taskData.Sync {
			packet_task := CreateSpAgentTaskSync(taskData)
			ts.TsSyncAllClients(packet_task)

			packet_console := CreateSpAgentConsoleTaskSync(taskData)
			ts.TsSyncAllClients(packet_console)

			agent.OutConsole.Put(packet_console)
			_ = ts.DBMS.DbConsoleInsert(agentId, packet_console)
		}
		agent.TasksQueue.Put(taskData)

	case TYPE_BROWSER:
		agent.TasksQueue.Put(taskData)

	case TYPE_JOB:
		if taskData.Sync {
			packet_task := CreateSpAgentTaskSync(taskData)
			ts.TsSyncAllClients(packet_task)

			packet_console := CreateSpAgentConsoleTaskSync(taskData)
			ts.TsSyncAllClients(packet_console)

			agent.OutConsole.Put(packet_console)
			_ = ts.DBMS.DbConsoleInsert(agentId, packet_console)
		}
		agent.TasksQueue.Put(taskData)

	case TYPE_TUNNEL:
		if taskData.Sync {
			if taskData.Completed {
				agent.CompletedTasks.Put(taskData.TaskId, taskData)
			} else {
				agent.RunningTasks.Put(taskData.TaskId, taskData)
			}

			packet_task := CreateSpAgentTaskSync(taskData)
			ts.TsSyncAllClients(packet_task)

			packet_console := CreateSpAgentConsoleTaskSync(taskData)
			ts.TsSyncAllClients(packet_console)

			agent.OutConsole.Put(packet_console)
			_ = ts.DBMS.DbConsoleInsert(agentId, packet_console)

			if taskData.Completed {
				_ = ts.DBMS.DbTaskInsert(taskData)
			}
		}

	case TYPE_PROXY_DATA:
		logs.Debug("", "----TYPE_PROXY_DATA----")

	default:
		break
	}
}

func (ts *Teamserver) TsTaskUpdate(agentId string, updateData adaptix.TaskData) {
	value, ok := ts.agents.Get(agentId)
	if !ok {
		logs.Error("", "TsTaskUpdate: agent %v not found", agentId)
		return
	}
	agent, _ := value.(*Agent)

	value, ok = agent.RunningTasks.Get(updateData.TaskId)
	if !ok {
		return
	}
	task, _ := value.(adaptix.TaskData)

	task.Data = []byte("")

	if task.Type == TYPE_JOB {
		updateData.AgentId = agentId

		if task.HookId != "" && task.Client != "" && ts.TsClientConnected(task.Client) {
			updateData.HookId = task.HookId

			hookJob := &HookJob{
				Job:       updateData,
				Processed: false,
				Sent:      false,
			}

			num := 0
			value2, ok := agent.RunningJobs.Get(task.TaskId)
			if ok {
				jobs := value2.(*safe.Slice)
				jobs.Put(hookJob)
				num = int(jobs.Len() - 1)
			} else {
				jobs := safe.NewSlice()
				jobs.Put(hookJob)
				agent.RunningJobs.Put(task.TaskId, jobs)
			}

			packet := CreateSpAgentTaskHook(updateData, num)
			ts.TsSyncClient(task.Client, packet)

		} else {
			if task.Sync {

				hookJob := &HookJob{
					Job:       updateData,
					Processed: true,
					Sent:      true,
				}

				value2, ok := agent.RunningJobs.Get(task.TaskId)

				if updateData.Completed {
					agent.RunningTasks.Delete(updateData.TaskId)

					if ok {
						jobs := value2.(*safe.Slice)
						jobs.Put(hookJob)
						jobs_array := jobs.CutArray()
						for _, job_value := range jobs_array {
							jobData := job_value.(*HookJob)
							if task.MessageType != CONSOLE_OUT_ERROR {
								task.MessageType = jobData.Job.MessageType
							}
							if task.Message == "" {
								task.Message = jobData.Job.Message
							}
							task.ClearText += jobData.Job.ClearText
						}

						agent.RunningJobs.Delete(task.TaskId)

					} else {
						task.MessageType = updateData.MessageType
						task.Message = updateData.Message
						task.ClearText = updateData.ClearText
					}

					task.FinishDate = updateData.FinishDate
					task.Completed = updateData.Completed

					agent.CompletedTasks.Put(task.TaskId, task)
					_ = ts.DBMS.DbTaskInsert(task)

				} else {
					if ok {
						jobs := value2.(*safe.Slice)
						jobs.Put(hookJob)
					} else {
						jobs := safe.NewSlice()
						jobs.Put(hookJob)
						agent.RunningJobs.Put(task.TaskId, jobs)
					}
				}

				packet_task_update := CreateSpAgentTaskUpdate(updateData)
				packet_console_update := CreateSpAgentConsoleTaskUpd(updateData)

				ts.TsSyncAllClients(packet_task_update)
				ts.TsSyncAllClients(packet_console_update)

				agent.OutConsole.Put(packet_console_update)
				_ = ts.DBMS.DbConsoleInsert(agentId, packet_console_update)
			}
		}

	} else if task.Type == TYPE_TUNNEL {
		agent.RunningTasks.Delete(updateData.TaskId)

		task.FinishDate = updateData.FinishDate
		task.Completed = updateData.Completed
		task.MessageType = updateData.MessageType

		var tmpTask = task
		tmpTask.Message = updateData.Message
		tmpTask.ClearText = updateData.ClearText

		if task.Message == "" {
			task.Message = updateData.Message
		}
		task.ClearText += updateData.ClearText

		if task.Sync {
			if task.Completed {
				agent.CompletedTasks.Put(task.TaskId, task)
				_ = ts.DBMS.DbTaskInsert(task)
			} else {
				agent.RunningTasks.Put(task.TaskId, task)
			}

			packet_task_update := CreateSpAgentTaskUpdate(tmpTask)
			packet_console_update := CreateSpAgentConsoleTaskUpd(tmpTask)

			ts.TsSyncAllClients(packet_task_update)
			ts.TsSyncAllClients(packet_console_update)

			agent.OutConsole.Put(packet_console_update)
			_ = ts.DBMS.DbConsoleInsert(agentId, packet_console_update)
		}

	} else if task.Type == TYPE_TASK || task.Type == TYPE_BROWSER {
		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 != "" && task.Client != "" && ts.TsClientConnected(task.Client) {

			agent.RunningTasks.Put(task.TaskId, task)

			packet := CreateSpAgentTaskHook(task, 0)
			ts.TsSyncClient(task.Client, packet)

		} else {
			if task.Sync {
				if task.Completed {
					agent.CompletedTasks.Put(task.TaskId, task)
					_ = ts.DBMS.DbTaskInsert(task)
				} else {
					agent.RunningTasks.Put(task.TaskId, task)
				}

				packet := CreateSpAgentTaskUpdate(task)
				ts.TsSyncAllClients(packet)

				packet2 := CreateSpAgentConsoleTaskUpd(task)
				ts.TsSyncAllClients(packet2)

				agent.OutConsole.Put(packet2)
				_ = ts.DBMS.DbConsoleInsert(agentId, packet2)
			}
		}
	}
}

func (ts *Teamserver) TsTaskPostHook(hookData adaptix.TaskData, jobIndex int) error {
	value, ok := ts.agents.Get(hookData.AgentId)
	if !ok {
		return fmt.Errorf("agent %v not found", hookData.AgentId)
	}
	agent, _ := value.(*Agent)

	value, ok = agent.RunningTasks.Get(hookData.TaskId)
	if !ok {
		return fmt.Errorf("task %v not found", hookData.TaskId)
	}
	task, _ := value.(adaptix.TaskData)

	if task.HookId == "" || task.HookId != hookData.HookId || task.Client != hookData.Client || !ts.TsClientConnected(task.Client) {
		return fmt.Errorf("Operation not available")
	}

	if task.Type == TYPE_JOB {

		if task.Sync {

			value2, ok := agent.RunningJobs.Get(task.TaskId)
			if !ok {
				return fmt.Errorf("job %v not found", task.TaskId)
			}
			jobs := value2.(*safe.Slice)

			jobValue, ok := jobs.Get(uint(jobIndex))
			if !ok {
				return fmt.Errorf("job %v not found", task.TaskId)
			}
			jobData := jobValue.(*HookJob)

			jobData.Job.MessageType = hookData.MessageType
			jobData.Job.Message = hookData.Message
			jobData.Job.ClearText = hookData.ClearText
			jobData.Processed = true

			completed := false
			sent := false

			jobs.DirectLock()
			defer jobs.DirectUnlock()

			slice := jobs.DirectSlice()
			for i := 0; i < len(slice); i++ {

				hookJob := slice[i].(*HookJob)
				if hookJob.Job.Completed {
					completed = true
				}
				hookJob.mu.Lock()
				if !hookJob.Sent {
					if hookJob.Processed {
						hookJob.Sent = true

						packet_task_update := CreateSpAgentTaskUpdate(hookJob.Job)
						packet_console_update := CreateSpAgentConsoleTaskUpd(hookJob.Job)

						ts.TsSyncAllClients(packet_task_update)
						ts.TsSyncAllClients(packet_console_update)

						agent.OutConsole.Put(packet_console_update)
						_ = ts.DBMS.DbConsoleInsert(task.AgentId, packet_console_update)
						sent = true
					} else {
						hookJob.mu.Unlock()
						break
					}
				}
				hookJob.mu.Unlock()
			}

			if completed && sent {
				agent.RunningTasks.Delete(task.TaskId)

				for i := 0; i < len(slice); i++ {
					hookJob := slice[i].(*HookJob)
					if task.MessageType != CONSOLE_OUT_ERROR {
						task.MessageType = hookJob.Job.MessageType
					}
					if task.Message == "" {
						task.Message = hookJob.Job.Message
					}
					task.ClearText += hookJob.Job.ClearText

					task.FinishDate = hookJob.Job.FinishDate
					task.Completed = hookJob.Job.Completed
				}

				agent.RunningJobs.Delete(task.TaskId)

				agent.CompletedTasks.Put(task.TaskId, task)
				_ = ts.DBMS.DbTaskInsert(task)
			}
		}

	} else if task.Type == TYPE_TUNNEL {

	} else if task.Type == TYPE_TASK || task.Type == TYPE_BROWSER {

		_, ok = agent.RunningTasks.GetDelete(hookData.TaskId)
		if !ok {
			return fmt.Errorf("task %v not found", hookData.TaskId)
		}

		task.MessageType = hookData.MessageType
		task.Message = hookData.Message
		task.ClearText = hookData.ClearText

		if task.Sync {
			if task.Completed {
				task.HookId = ""
				agent.CompletedTasks.Put(task.TaskId, task)
				_ = ts.DBMS.DbTaskInsert(task)
			} else {
				agent.RunningTasks.Put(task.TaskId, task)
			}

			packet := CreateSpAgentTaskUpdate(task)
			ts.TsSyncAllClients(packet)

			packet2 := CreateSpAgentConsoleTaskUpd(task)
			ts.TsSyncAllClients(packet2)

			agent.OutConsole.Put(packet2)
			_ = ts.DBMS.DbConsoleInsert(task.AgentId, packet2)
		}
	}
	return nil
}

func (ts *Teamserver) TsTaskCancel(agentId string, taskId string) error {
	value, ok := ts.agents.Get(agentId)
	if !ok {
		return fmt.Errorf("agent %v not found", agentId)
	}
	agent, _ := value.(*Agent)

	var task adaptix.TaskData
	found := false
	for i := uint(0); i < agent.TasksQueue.Len(); i++ {
		if value, ok = agent.TasksQueue.Get(i); ok {
			task = value.(adaptix.TaskData)
			if task.TaskId == taskId {
				agent.TasksQueue.Delete(i)
				found = true
				break
			}
		}
	}

	if found {
		packet := CreateSpAgentTaskRemove(task)
		ts.TsSyncAllClients(packet)
		return nil
	}

	return nil
}

func (ts *Teamserver) TsTaskDelete(agentId string, taskId string) error {
	value, ok := ts.agents.Get(agentId)
	if !ok {
		return fmt.Errorf("agent %v not found", agentId)
	}
	agent, _ := value.(*Agent)

	var task adaptix.TaskData
	for i := uint(0); i < agent.TasksQueue.Len(); i++ {
		if value, ok = agent.TasksQueue.Get(i); ok {
			task = value.(adaptix.TaskData)
			if task.TaskId == taskId {
				return fmt.Errorf("task %v in process", taskId)
			}
		}
	}

	value, ok = agent.RunningTasks.Get(taskId)
	if ok {
		return fmt.Errorf("task %v in process", taskId)
	}
	value, ok = agent.CompletedTasks.GetDelete(taskId)
	if !ok {
		return fmt.Errorf("task %v not found", taskId)
	}

	task = value.(adaptix.TaskData)
	_ = ts.DBMS.DbTaskDelete(task.TaskId, "")

	packet := CreateSpAgentTaskRemove(task)
	ts.TsSyncAllClients(packet)
	return nil
}

///// Get Tasks

func (ts *Teamserver) TsTaskGetAvailableAll(agentId string, availableSize int) ([]adaptix.TaskData, error) {
	value, ok := ts.agents.Get(agentId)
	if !ok {
		return nil, fmt.Errorf("TsTaskQueueGetAvailable: agent %v not found", agentId)
	}
	agent, _ := value.(*Agent)

	var tasks []adaptix.TaskData
	tasksSize := 0

	/// TASKS QUEUE

	var sendTasks []string
	for i := uint(0); i < agent.TasksQueue.Len(); i++ {
		value, ok = agent.TasksQueue.Get(i)
		if ok {
			taskData := value.(adaptix.TaskData)
			if tasksSize+len(taskData.Data) < availableSize {
				tasks = append(tasks, taskData)
				if taskData.Sync || taskData.Type == TYPE_BROWSER {
					agent.RunningTasks.Put(taskData.TaskId, taskData)
				}
				agent.TasksQueue.Delete(i)
				i--
				sendTasks = append(sendTasks, taskData.TaskId)
				tasksSize += len(taskData.Data)
			} else {
				break
			}
		} else {
			break
		}
	}
	if len(sendTasks) > 0 {
		packet := CreateSpAgentTaskSend(sendTasks)
		ts.TsSyncAllClients(packet)
	}

	for i := uint(0); i < agent.TunnelConnectTasks.Len(); i++ {
		value, ok = agent.TunnelConnectTasks.Get(i)
		if ok {
			taskData := value.(adaptix.TaskData)
			if tasksSize+len(taskData.Data) < availableSize {
				tasks = append(tasks, taskData)
				agent.TunnelConnectTasks.Delete(i)
				i--
				tasksSize += len(taskData.Data)
			} else {
				break
			}
		} else {
			break
		}
	}

	/// TUNNELS QUEUE

	for i := uint(0); i < agent.TunnelQueue.Len(); i++ {
		value, ok = agent.TunnelQueue.Get(i)
		if ok {
			taskDataTunnel := value.(adaptix.TaskDataTunnel)
			if tasksSize+len(taskDataTunnel.Data.Data) < availableSize {
				tasks = append(tasks, taskDataTunnel.Data)
				agent.TunnelQueue.Delete(i)
				i--
				tasksSize += len(taskDataTunnel.Data.Data)
			} else {
				break
			}
		} else {
			break
		}
	}

	/// PIVOTS QUEUE

	for i := uint(0); i < agent.PivotChilds.Len(); i++ {
		value, ok = agent.PivotChilds.Get(i)
		if ok {
			pivotData := value.(*adaptix.PivotData)
			lostSize := availableSize - tasksSize
			if availableSize > 0 {
				data, err := ts.TsAgentGetHostedTasksAll(pivotData.ChildAgentId, lostSize)
				if err != nil {
					continue
				}
				pivotTaskData, err := ts.Extender.ExAgentPivotPackData(agent.Data.Name, pivotData.PivotId, data)
				if err != nil {
					continue
				}
				tasks = append(tasks, pivotTaskData)
				tasksSize += len(pivotTaskData.Data)
			}
		} else {
			break
		}
	}

	return tasks, nil
}

func (ts *Teamserver) TsTaskGetAvailableTasks(agentId string, availableSize int) ([]adaptix.TaskData, int, error) {
	value, ok := ts.agents.Get(agentId)
	if !ok {
		return nil, 0, fmt.Errorf("TsTaskQueueGetAvailable: agent %v not found", agentId)
	}
	agent, _ := value.(*Agent)

	var tasks []adaptix.TaskData
	tasksSize := 0

	/// TASKS QUEUE

	var sendTasks []string
	for i := uint(0); i < agent.TasksQueue.Len(); i++ {
		value, ok = agent.TasksQueue.Get(i)
		if ok {
			taskData := value.(adaptix.TaskData)
			if tasksSize+len(taskData.Data) < availableSize {
				tasks = append(tasks, taskData)
				if taskData.Sync || taskData.Type == TYPE_BROWSER {
					agent.RunningTasks.Put(taskData.TaskId, taskData)
				}
				agent.TasksQueue.Delete(i)
				i--
				sendTasks = append(sendTasks, taskData.TaskId)
				tasksSize += len(taskData.Data)
			} else {
				break
			}
		} else {
			break
		}
	}
	if len(sendTasks) > 0 {
		packet := CreateSpAgentTaskSend(sendTasks)
		ts.TsSyncAllClients(packet)
	}

	for i := uint(0); i < agent.TunnelConnectTasks.Len(); i++ {
		value, ok = agent.TunnelConnectTasks.Get(i)
		if ok {
			taskData := value.(adaptix.TaskData)
			if tasksSize+len(taskData.Data) < availableSize {
				tasks = append(tasks, taskData)
				agent.TunnelConnectTasks.Delete(i)
				i--
				//sendTasks = append(sendTasks, taskData.TaskId)
				tasksSize += len(taskData.Data)
			} else {
				break
			}
		} else {
			break
		}
	}

	return tasks, tasksSize, nil
}

/// Get Pivot Tasks

func (ts *Teamserver) TsTasksPivotExists(agentId string, first bool) bool {
	value, ok := ts.agents.Get(agentId)
	if !ok {
		return false
	}
	agent := value.(*Agent)

	if !first {
		if agent.TasksQueue.Len() > 0 || agent.TunnelQueue.Len() > 0 {
			return true
		}
	}

	for i := uint(0); i < agent.PivotChilds.Len(); i++ {
		value, ok = agent.PivotChilds.Get(i)
		if ok {
			pivotData := value.(*adaptix.PivotData)
			if ts.TsTasksPivotExists(pivotData.ChildAgentId, false) {
				return true
			}
		}
	}
	return false
}

func (ts *Teamserver) TsTaskGetAvailablePivotAll(agentId string, availableSize int) ([]adaptix.TaskData, error) {
	value, ok := ts.agents.Get(agentId)
	if !ok {
		return nil, fmt.Errorf("TsTaskQueueGetAvailable: agent %v not found", agentId)
	}
	agent, _ := value.(*Agent)

	var tasks []adaptix.TaskData
	tasksSize := 0

	/// PIVOTS QUEUE

	for i := uint(0); i < agent.PivotChilds.Len(); i++ {
		value, ok = agent.PivotChilds.Get(i)
		if ok {
			pivotData := value.(*adaptix.PivotData)
			lostSize := availableSize - tasksSize
			if availableSize > 0 {
				data, err := ts.TsAgentGetHostedTasksAll(pivotData.ChildAgentId, lostSize)
				if err != nil {
					continue
				}
				pivotTaskData, err := ts.Extender.ExAgentPivotPackData(agent.Data.Name, pivotData.PivotId, data)
				if err != nil {
					continue
				}
				tasks = append(tasks, pivotTaskData)
				tasksSize += len(pivotTaskData.Data)
			}
		} else {
			break
		}
	}

	return tasks, nil
}