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

package server

import (
	"AdaptixServer/core/eventing"
	"AdaptixServer/core/utils/logs"
	"fmt"

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

func (ts *Teamserver) TsGetPivotInfoByName(pivotName string) (string, string, string) {
	for value := range ts.pivots.Iterator() {
		pivot := value.Item.(*adaptix.PivotData)
		if pivot.PivotName == pivotName {
			return pivot.PivotId, pivot.ParentAgentId, pivot.ChildAgentId
		}
	}
	return "", "", ""
}

func (ts *Teamserver) TsGetPivotInfoById(pivotId string) (string, string, string) {
	for value := range ts.pivots.Iterator() {
		pivot := value.Item.(*adaptix.PivotData)
		if pivot.PivotId == pivotId {
			return pivot.PivotName, pivot.ParentAgentId, pivot.ChildAgentId
		}
	}
	return "", "", ""
}

func (ts *Teamserver) TsGetPivotByName(pivotName string) *adaptix.PivotData {
	for value := range ts.pivots.Iterator() {
		pivot := value.Item.(*adaptix.PivotData)
		if pivot.PivotName == pivotName {
			return pivot
		}
	}
	return nil
}

func (ts *Teamserver) TsGetPivotById(pivotId string) *adaptix.PivotData {
	for value := range ts.pivots.Iterator() {
		pivot := value.Item.(*adaptix.PivotData)
		if pivot.PivotId == pivotId {
			return pivot
		}
	}
	return nil
}

func (ts *Teamserver) TsPivotCreate(pivotId string, pAgentId string, chAgentId string, pivotName string, isRestore bool) error {
	// --- PRE HOOK ---
	preEvent := &eventing.EventDataPivotCreate{
		PivotId:       pivotId,
		ParentAgentId: pAgentId,
		ChildAgentId:  chAgentId,
		PivotName:     pivotName,
	}
	if !ts.EventManager.Emit(eventing.EventPivotCreate, eventing.HookPre, preEvent) {
		if preEvent.Error != nil {
			return preEvent.Error
		}
		return fmt.Errorf("operation cancelled by hook")
	}
	// ----------------

	pivotData := &adaptix.PivotData{
		PivotId:       pivotId,
		PivotName:     pivotName,
		ParentAgentId: pAgentId,
		ChildAgentId:  chAgentId,
	}

	if pivotData.PivotName == "" || ts.TsGetPivotByName(pivotData.PivotName) != nil {
		usedNames := make(map[string]bool)
		for value := range ts.pivots.Iterator() {
			pivot := value.Item.(*adaptix.PivotData)
			usedNames[pivot.PivotName] = true
		}

		ok := false
		for i := 0; i <= 9999; i++ {
			name := fmt.Sprintf("p%d", i)
			if !usedNames[name] {
				pivotData.PivotName = name
				ok = true
				break
			}
		}
		if !ok {
			return fmt.Errorf("the number of pivots has exceeded 9999")
		}
	}

	if parentAgent, err := ts.getAgent(pivotData.ParentAgentId); err == nil {
		parentAgent.PivotChilds.Put(pivotData)
	}

	if childAgent, err := ts.getAgent(pivotData.ChildAgentId); err == nil {
		childAgent.PivotParent = pivotData
	}

	//emptyMark := ""
	//_ = ts.TsAgentUpdateDataPartial(pivotData.ChildAgentId, struct {
	//	Mark *string `json:"mark"`
	//}{Mark: &emptyMark})

	ts.pivots.Put(pivotData)

	if !isRestore {
		err := ts.DBMS.DbPivotInsert(*pivotData)
		if err != nil {
			logs.Error("", err.Error())
		}
	}

	packet := CreateSpPivotCreate(*pivotData)
	ts.TsSyncAllClients(packet)

	// --- POST HOOK ---
	postEvent := &eventing.EventDataPivotCreate{
		PivotId:       pivotData.PivotId,
		ParentAgentId: pivotData.ParentAgentId,
		ChildAgentId:  pivotData.ChildAgentId,
		PivotName:     pivotData.PivotName,
	}
	ts.EventManager.EmitAsync(eventing.EventPivotCreate, postEvent)
	// -----------------

	return nil
}

func (ts *Teamserver) TsPivotDelete(pivotId string) error {
	// --- PRE HOOK ---
	preEvent := &eventing.EventDataPivotRemove{PivotId: pivotId}
	if !ts.EventManager.Emit(eventing.EventPivotRemove, eventing.HookPre, preEvent) {
		if preEvent.Error != nil {
			return preEvent.Error
		}
		return fmt.Errorf("operation cancelled by hook")
	}
	// ----------------

	pivotData := ts.TsGetPivotById(pivotId)
	if pivotData == nil {
		return fmt.Errorf("pivotId %s does not exist", pivotId)
	}

	if parentAgent, err := ts.getAgent(pivotData.ParentAgentId); err == nil {
		for i := uint(0); i < parentAgent.PivotChilds.Len(); i++ {
			valuePivot, ok := parentAgent.PivotChilds.Get(i)
			if ok {
				if pivot, ok := valuePivot.(*adaptix.PivotData); ok && pivot.PivotId == pivotId {
					parentAgent.PivotChilds.Delete(i)
					break
				}
			}
		}
	}

	if childAgent, err := ts.getAgent(pivotData.ChildAgentId); err == nil {
		childAgent.PivotParent = nil
	}

	unlinkMark := "Unlink"
	_ = ts.TsAgentUpdateDataPartial(pivotData.ChildAgentId, struct {
		Mark *string `json:"mark"`
	}{Mark: &unlinkMark})

	for i := uint(0); i < ts.pivots.Len(); i++ {
		valuePivot, ok := ts.pivots.Get(i)
		if ok {
			if valuePivot.(*adaptix.PivotData).PivotId == pivotId {
				ts.pivots.Delete(i)
				break
			}
		}
	}

	err := ts.DBMS.DbPivotDelete(pivotId)
	if err != nil {
		logs.Error("", err.Error())
	}

	packet := CreateSpPivotDelete(pivotId)
	ts.TsSyncAllClients(packet)

	// --- POST HOOK ---
	postEvent := &eventing.EventDataPivotRemove{PivotId: pivotId}
	ts.EventManager.EmitAsync(eventing.EventPivotRemove, postEvent)
	// -----------------

	return nil
}