AdaptixC2-Mod0/AdaptixServer/core/server/ts_agent_builder.go

package server

import (
	"AdaptixServer/core/eventing"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/rand/v2"
	"os/exec"
	"strings"
	"sync"

	"github.com/Adaptix-Framework/axc2"
	"github.com/gorilla/websocket"
)

func (ts *Teamserver) TsAgentBuildSyncOnce(agentName string, config string, listenersName []string) ([]byte, string, error) {

	conf := adaptix.BuildProfile{
		BuilderId:   "",
		AgentConfig: config,
	}

	for _, listener := range listenersName {
		listenerWM, listenerProfile, err := ts.TsListenerGetProfile(listener)
		if err != nil {
			return nil, "", err
		}

		transport := adaptix.TransportProfile{
			Watermark: listenerWM,
			Profile:   listenerProfile,
		}

		conf.ListenerProfiles = append(conf.ListenerProfiles, transport)
	}

	return ts.Extender.ExAgentGenerate(agentName, conf)
}

type BuildChannelData struct {
	AgentName     string   `json:"agent_name"`
	ListenersName []string `json:"listeners_name"`
}

func (ts *Teamserver) TsAgentBuildCreateChannel(buildData string, wsconn *websocket.Conn) error {
	var bd BuildChannelData
	if err := json.Unmarshal([]byte(buildData), &bd); err != nil {
		if wsconn != nil {
			_ = wsconn.WriteMessage(websocket.TextMessage, []byte("invalid build data"))
			wsconn.Close()
		}
		return errors.New("invalid build data")
	}

	if bd.AgentName == "" || len(bd.ListenersName) == 0 {
		if wsconn != nil {
			_ = wsconn.WriteMessage(websocket.TextMessage, []byte("invalid build data"))
			wsconn.Close()
		}
		return errors.New("invalid build data")
	}

	listenersName := bd.ListenersName

	_, wsMsg, err := wsconn.ReadMessage()
	if err != nil {
		_ = wsconn.WriteMessage(websocket.TextMessage, []byte("invalid agent config"))
		wsconn.Close()
		return errors.New("invalid agent config")
	}

	builder := &AgentBuilder{
		Id:            fmt.Sprintf("%08x", rand.Uint32()),
		Name:          bd.AgentName,
		ListenersName: listenersName,
		Config:        string(wsMsg),
		wsconn:        wsconn,
		mu:            sync.Mutex{},
		closed:        false,
	}
	ts.builders.Put(builder.Id, builder)

	var (
		fileContent []byte
		fileName    string
		conf        adaptix.BuildProfile
	)

	_ = ts.TsAgentBuildLog(builder.Id, adaptix.BUILD_LOG_INFO, "Building agent...")

	var postEvent *eventing.EventDataAgentGenerate
	// --- PRE HOOK ---
	preEvent := &eventing.EventDataAgentGenerate{
		AgentName:     builder.Name,
		ListenersName: builder.ListenersName,
		Config:        builder.Config,
	}
	if !ts.EventManager.Emit(eventing.EventAgentGenerate, eventing.HookPre, preEvent) {
		if preEvent.Error != nil {
			_ = ts.TsAgentBuildLog(builder.Id, adaptix.BUILD_LOG_ERROR, "Error: "+preEvent.Error.Error())
		} else {
			_ = ts.TsAgentBuildLog(builder.Id, adaptix.BUILD_LOG_ERROR, "Error: operation cancelled by hook")
		}
		goto RET
	}
	// ----------------

	conf = adaptix.BuildProfile{
		BuilderId:   builder.Id,
		AgentConfig: builder.Config,
	}

	for _, listener := range listenersName {
		listenerWM, listenerProfile, err := ts.TsListenerGetProfile(listener)
		if err != nil {
			_ = ts.TsAgentBuildLog(builder.Id, adaptix.BUILD_LOG_ERROR, fmt.Sprintf("Error: invalid '%s' listener profile", listener))
			goto RET
		}

		_ = ts.TsAgentBuildLog(builder.Id, adaptix.BUILD_LOG_INFO, fmt.Sprintf("Listener '%s' profile created", listener))

		transport := adaptix.TransportProfile{
			Watermark: listenerWM,
			Profile:   listenerProfile,
		}
		conf.ListenerProfiles = append(conf.ListenerProfiles, transport)
	}

	fileContent, fileName, err = ts.Extender.ExAgentGenerate(builder.Name, conf)
	if err != nil {
		_ = ts.TsAgentBuildLog(builder.Id, adaptix.BUILD_LOG_ERROR, "Error: agent builder failed")
		goto RET
	}
	_ = ts.TsAgentBuildLog(builder.Id, adaptix.BUILD_LOG_SUCCESS, "Agent built successfully")

	// --- POST HOOK ---
	postEvent = &eventing.EventDataAgentGenerate{
		AgentName:     builder.Name,
		ListenersName: builder.ListenersName,
		Config:        builder.Config,
		FileName:      fileName,
		FileContent:   fileContent,
	}
	ts.EventManager.EmitAsync(eventing.EventAgentGenerate, postEvent)
	// -----------------

	_ = ts.TsAgentBuildSendFile(builder.Id, fileName, fileContent)

RET:
	ts.TsAgentBuildClose(builder.Id)
	return err
}

func (ts *Teamserver) TsAgentBuildClose(builderId string) {
	if builderId == "" {
		return
	}

	value, ok := ts.builders.GetDelete(builderId)
	if !ok {
		return
	}
	builder, _ := value.(*AgentBuilder)

	builder.mu.Lock()
	defer builder.mu.Unlock()

	if !builder.closed {
		builder.closed = true
		if builder.wsconn != nil {
			_ = builder.wsconn.Close()
		}
	}
}

func (ts *Teamserver) TsAgentBuildExecute(builderId string, workingDir string, program string, args ...string) error {
	runner := exec.Command(program, args...)
	runner.Dir = workingDir

	if builderId == "" {
		var stderr strings.Builder
		runner.Stderr = &stderr

		err := runner.Run()
		if err != nil {
			return fmt.Errorf("%v: %s", err, stderr.String())
		}
		return nil
	}

	value, ok := ts.builders.Get(builderId)
	if !ok {
		return errors.New("builder not found")
	}
	builder, _ := value.(*AgentBuilder)

	builder.mu.Lock()
	if builder.closed {
		builder.mu.Unlock()
		return errors.New("channel closed")
	}
	builder.mu.Unlock()

	var (
		stdoutPipe io.ReadCloser
		stderrPipe io.ReadCloser
		err        error
	)

	if stdoutPipe, err = runner.StdoutPipe(); err != nil {
		return err
	}
	if stderrPipe, err = runner.StderrPipe(); err != nil {
		return err
	}
	if err = runner.Start(); err != nil {
		return err
	}

	var stderrBuf strings.Builder
	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		buf := make([]byte, 1024)
		for {
			n, err := stdoutPipe.Read(buf)
			if n > 0 {
				_ = ts.TsAgentBuildLog(builderId, adaptix.BUILD_LOG_NONE, string(buf[:n]))
			}
			if err != nil {
				break
			}
		}
	}()

	go func() {
		defer wg.Done()
		buf := make([]byte, 1024)
		for {
			n, err := stderrPipe.Read(buf)
			if n > 0 {
				stderrBuf.Write(buf[:n])
				_ = ts.TsAgentBuildLog(builderId, adaptix.BUILD_LOG_NONE, string(buf[:n]))
			}
			if err != nil {
				break
			}
		}
	}()

	wg.Wait()
	err = runner.Wait()
	if err != nil {
		return fmt.Errorf("%v: %s", err, stderrBuf.String())
	}

	return nil
}

type BuildLogMessage struct {
	Status  int    `json:"status"`
	Message string `json:"message"`
}

func (ts *Teamserver) TsAgentBuildLog(builderId string, status int, message string) error {
	if builderId == "" {
		return errors.New("builder ID is empty")
	}

	value, ok := ts.builders.Get(builderId)
	if !ok {
		return errors.New("builder not found")
	}
	builder, _ := value.(*AgentBuilder)

	builder.mu.Lock()
	defer builder.mu.Unlock()

	if builder.closed {
		return errors.New("channel closed")
	}

	logMsg := BuildLogMessage{Status: status, Message: message}
	jsonData, err := json.Marshal(logMsg)
	if err != nil {
		return err
	}
	return builder.wsconn.WriteMessage(websocket.TextMessage, jsonData)
}

type BuiltPayload struct {
	Status   int    `json:"status"`
	Filename string `json:"filename"`
	Content  []byte `json:"content"`
}

func (ts *Teamserver) TsAgentBuildSendFile(builderId string, filename string, content []byte) error {
	if builderId == "" {
		return errors.New("builder ID is empty")
	}

	value, ok := ts.builders.Get(builderId)
	if !ok {
		return errors.New("builder not found")
	}
	builder, _ := value.(*AgentBuilder)

	builder.mu.Lock()
	defer builder.mu.Unlock()

	if builder.closed {
		return errors.New("channel closed")
	}

	data := BuiltPayload{
		Status:   4,
		Filename: filename,
		Content:  content,
	}

	jsonData, err := json.Marshal(data)
	if err != nil {
		return err
	}

	return builder.wsconn.WriteMessage(websocket.TextMessage, jsonData)
}