RedFlag/agent/internal/supplychain/binary_update.go

package supplychain

import (
	"context"
	"crypto/ed25519"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/Fimeg/RedFlag/agent/internal/capability"
	"github.com/Fimeg/RedFlag/agent/internal/constants"
)

const (
	AgentSelfPackageType   = "agent-self"
	HelperSelfPackageType  = "helper-self"
	DesktopSelfPackageType = "desktop-self"

	selfUpdateOperation = "upgrade"

	agentSelfStagingPath = "/var/lib/redflag/agent/pending-upgrade.bin"

	// helperSelfStagingPath is a cross-language contract: the Go agent stages the
	// new helper binary here, and the Rust helper reads it on self-upgrade.
	// The Rust counterpart is DEFAULT_HELPER_SELF_SOURCE in helper/src/main.rs.
	// Both must be updated together if this path changes.
	helperSelfStagingPath = "/var/lib/redflag/agent/pending-helper.bin"

	maxSelfUpdateBinarySize = 500 * 1024 * 1024
)

// Exit codes mirror helper/src/main.rs so capability receipts have one taxonomy
// across helper-executed and agent-local self-update paths.
const (
	policyExitOK            = 0
	policyExitBadToken      = 10
	policyExitTimeWindow    = 12
	policyExitKeyNotFound   = 14
	policyExitSignature     = 15
	policyExitArtifact      = 16
	policyExitReplay        = 17
	policyExitUnsupportedOp = 18
	policyExitExecFailed    = 19
	policyExitInternal      = 20
)

func (c *Consumer) processAgentSelfToken(ctx context.Context, token *capability.Token) (*PolicyResult, error) {
	if token.Operation != selfUpdateOperation {
		return nil, fmt.Errorf("agent-self operation %q not supported", token.Operation)
	}

	agentEntry, err := findClosureEntry(token, "redflag-agent")
	if err != nil {
		return nil, err
	}
	stagedAgent, err := c.stageClosureArtifact(agentEntry, agentSelfStagingPath)
	if err != nil {
		return nil, fmt.Errorf("stage agent-self artifact: %w", err)
	}
	defer os.Remove(stagedAgent)

	var helperArgs []string
	if helperEntry, err := findClosureEntry(token, "redflag-helper"); err == nil {
		stagedHelper, err := c.stageClosureArtifact(helperEntry, helperSelfStagingPath)
		if err != nil {
			return nil, fmt.Errorf("stage agent-self helper artifact: %w", err)
		}
		defer os.Remove(stagedHelper)
		helperArgs = []string{"--helper-file", stagedHelper}
	}

	return c.executor.Execute(ctx, token, helperArgs...)
}

func (c *Consumer) processHelperSelfToken(ctx context.Context, token *capability.Token) (*PolicyResult, error) {
	if token.Operation != selfUpdateOperation {
		return nil, fmt.Errorf("helper-self operation %q not supported", token.Operation)
	}

	entry, err := findClosureEntry(token, "redflag-helper")
	if err != nil {
		return nil, err
	}
	staged, err := c.stageClosureArtifact(entry, helperSelfStagingPath)
	if err != nil {
		return nil, fmt.Errorf("stage helper-self artifact: %w", err)
	}
	defer os.Remove(staged)

	return c.executor.Execute(ctx, token)
}

func (c *Consumer) processDesktopSelfToken(ctx context.Context, token *capability.Token) (*PolicyResult, error) {
	if err := validateDirectSelfToken(token); err != nil {
		return selfPolicyResult(token, "denied", "desktop_self_token_invalid", policyExitSignature, 0, err), nil
	}
	if token.Operation != selfUpdateOperation {
		return selfPolicyResult(
			token,
			"denied",
			"operation_not_allowed",
			policyExitUnsupportedOp,
			0,
			fmt.Errorf("operation=%s (desktop-self requires upgrade)", token.Operation),
		), nil
	}

	// Bound the in-process staging + install to match the executor timeout.
	// The agent-self and helper-self paths go through executor.Execute which
	// carries its own context timeout; desktop-self runs in-process and needs
	// an equivalent guard against hung downloads or stalled I/O.
	ctx, cancel := context.WithTimeout(ctx, executorTimeout)
	defer cancel()

	entry, err := findClosureEntry(token, "redflag-desktop")
	if err != nil {
		return selfPolicyResult(token, "denied", "desktop_self_no_entry", policyExitBadToken, 0, err), nil
	}

	// Check context before starting potentially long staging operation.
	if err := ctx.Err(); err != nil {
		return selfPolicyResult(token, "failed", "desktop_self_cancelled", policyExitExecFailed, 0, err), nil
	}

	stagingPath := filepath.Join(constants.GetAgentStateDir(), "desktop-self-upgrade.bin")
	staged, err := c.stageClosureArtifact(entry, stagingPath)
	if err != nil {
		return selfPolicyResult(token, "failed", "desktop_self_stage_failed", policyExitArtifact, 0, err), nil
	}
	defer os.Remove(staged)

	// Check replay before any side effects; record consumed only after a
	// successful install. The token is Ed25519-signed, hash-pinned, and
	// version-pinned, so re-executing the identical install after a transient
	// failure (disk full, backup fail) is idempotent — burning the token before
	// install would permanently prevent retry (ETHOS #3, #4).
	if err := replayCheckAgent(token.TokenID); err != nil {
		return selfPolicyResult(token, "denied", "token_already_consumed", policyExitReplay, 1, err), nil
	}

	target, err := installDesktopBinary(staged)
	if err != nil {
		log.Printf("[ERROR] [agent] [supplychain] desktop_self_install_failed token_id=%s error=%v", token.TokenID, err)
		return selfPolicyResult(token, "failed", "desktop_self_install_failed", policyExitExecFailed, 1, err), nil
	}

	// Install succeeded — mark token consumed so it cannot be replayed.
	if err := recordAgentTokenConsumed(token.TokenID); err != nil {
		// Log but do not abort: the binary is already in place. A duplicate
		// replay would be caught by hash + version checks, and the helper
		// would refuse an identical re-install anyway.
		log.Printf("[WARNING] [agent] [supplychain] desktop_self_replay_record_failed token_id=%s error=%v", token.TokenID, err)
	}

	if err := signalDesktopRestart(target); err != nil {
		log.Printf("[WARNING] [agent] [supplychain] desktop_restart_signal_failed token_id=%s error=%v", token.TokenID, err)
	}

	log.Printf("[SECURITY] [agent] [supplychain] desktop_self_upgraded token_id=%s path=%s", token.TokenID, target)
	return selfPolicyResult(token, "executed", "desktop_self_upgraded", policyExitOK, 1, nil), nil
}

func findClosureEntry(token *capability.Token, name string) (*capability.ClosureEntry, error) {
	for i := range token.Closure {
		if token.Closure[i].Name == name {
			return &token.Closure[i], nil
		}
	}
	return nil, fmt.Errorf("closure missing %s entry", name)
}

func (c *Consumer) stageClosureArtifact(entry *capability.ClosureEntry, dstPath string) (string, error) {
	if entry.ArtifactPath == "" {
		return "", fmt.Errorf("closure entry %s missing artifact_path", entry.Name)
	}
	if strings.TrimSpace(entry.SHA256) == "" {
		return "", fmt.Errorf("closure entry %s missing sha256", entry.Name)
	}

	dstDir := filepath.Dir(dstPath)
	if err := os.MkdirAll(dstDir, 0o700); err != nil {
		return "", fmt.Errorf("create staging dir: %w", err)
	}

	tmp, err := os.CreateTemp(dstDir, filepath.Base(dstPath)+".")
	if err != nil {
		return "", fmt.Errorf("create staging temp: %w", err)
	}
	tmpPath := tmp.Name()
	if err := tmp.Close(); err != nil {
		os.Remove(tmpPath)
		return "", fmt.Errorf("close staging temp: %w", err)
	}
	cleanup := true
	defer func() {
		if cleanup {
			os.Remove(tmpPath)
		}
	}()

	if err := c.fetchArtifactToFile(entry.ArtifactPath, tmpPath); err != nil {
		return "", err
	}
	actual, err := computeFileSHA256Hex(tmpPath)
	if err != nil {
		return "", err
	}
	if !strings.EqualFold(actual, entry.SHA256) {
		return "", fmt.Errorf("hash mismatch for %s: expected=%s actual=%s", entry.Name, strings.ToLower(entry.SHA256), actual)
	}
	if err := os.Chmod(tmpPath, 0o644); err != nil {
		return "", fmt.Errorf("chmod staged artifact: %w", err)
	}
	if err := os.Rename(tmpPath, dstPath); err != nil {
		return "", fmt.Errorf("commit staged artifact: %w", err)
	}
	cleanup = false
	return dstPath, nil
}

func (c *Consumer) fetchArtifactToFile(ref, dstPath string) error {
	if info, err := os.Stat(ref); err == nil {
		if info.IsDir() {
			return fmt.Errorf("artifact path is a directory: %s", ref)
		}
		return copyRegularFile(ref, dstPath)
	} else if !os.IsNotExist(err) {
		// Stat failed for a reason other than "not found" (e.g. permission denied).
		return fmt.Errorf("stat artifact %s: %w", ref, err)
	}

	// Local absolute path that doesn't exist — don't fall through to the
	// downloader; the file was expected on disk and is missing.
	if strings.HasPrefix(ref, "/") {
		return fmt.Errorf("artifact not staged at %s", ref)
	}

	if isDownloadRef(ref) {
		if c.downloader == nil {
			return fmt.Errorf("artifact downloader unavailable for %s", ref)
		}
		if _, err := c.downloader.DownloadAuthenticatedToFile(ref, dstPath, maxSelfUpdateBinarySize); err != nil {
			return fmt.Errorf("download artifact %s: %w", ref, err)
		}
		return nil
	}

	return copyRegularFile(ref, dstPath)
}

func isDownloadRef(ref string) bool {
	return strings.HasPrefix(ref, "http://") || strings.HasPrefix(ref, "https://")
}

func validateDirectSelfToken(token *capability.Token) error {
	if token.Version != capability.Version {
		return fmt.Errorf("unsupported token version=%d supported=%d", token.Version, capability.Version)
	}
	now := time.Now().UTC().Unix()
	if now < token.NotBefore {
		return fmt.Errorf("token not yet valid: now=%d not_before=%d", now, token.NotBefore)
	}
	if now > token.ExpiresAt {
		return fmt.Errorf("token expired: now=%d expires_at=%d", now, token.ExpiresAt)
	}

	pub, err := loadCapabilityPublicKey()
	if err != nil {
		return err
	}
	keyID := capability.KeyIDFor(ed25519.PublicKey(pub))
	if token.KeyID != keyID {
		return fmt.Errorf("key_id mismatch: token=%s local=%s", token.KeyID, keyID)
	}
	if err := token.Verify(ed25519.PublicKey(pub)); err != nil {
		return err
	}
	return nil
}

func loadCapabilityPublicKey() ([]byte, error) {
	keyPath := constants.GetServerPublicKeyPath()
	data, err := os.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("public key not found at %s: %w", keyPath, err)
	}
	if len(data) == ed25519.PublicKeySize {
		return data, nil
	}
	trimmed := strings.TrimSpace(string(data))
	decoded, decErr := hex.DecodeString(trimmed)
	if decErr == nil && len(decoded) == ed25519.PublicKeySize {
		return decoded, nil
	}
	return nil, fmt.Errorf("invalid public key size at %s: raw=%d decoded=%d", keyPath, len(data), len(decoded))
}

// replayCheckAgent returns an error if tokenID has already been consumed.
// It does NOT record the token; call recordAgentTokenConsumed after a
// successful install to close the replay window.
func replayCheckAgent(tokenID string) error {
	if _, err := safeTokenFilename(tokenID); err != nil {
		return err
	}
	if strings.ContainsAny(tokenID, "\r\n") {
		return fmt.Errorf("unsafe token_id: contains newline")
	}

	statePath := filepath.Join(constants.GetAgentStateDir(), "consumed-self-tokens")
	contents, err := os.ReadFile(statePath)
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("read replay state: %w", err)
	}
	if err == nil {
		for _, line := range strings.Split(string(contents), "\n") {
			if strings.TrimSpace(line) == tokenID {
				return fmt.Errorf("token_id=%s", tokenID)
			}
		}
	}
	return nil
}

// recordAgentTokenConsumed appends tokenID to the consumed-self-tokens file.
// Call only after a successful install to avoid permanently burning the token
// on a transient failure.
func recordAgentTokenConsumed(tokenID string) error {
	statePath := filepath.Join(constants.GetAgentStateDir(), "consumed-self-tokens")
	if err := os.MkdirAll(filepath.Dir(statePath), 0o700); err != nil {
		return fmt.Errorf("create replay state dir: %w", err)
	}
	f, err := os.OpenFile(statePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
	if err != nil {
		return fmt.Errorf("open replay state: %w", err)
	}
	defer f.Close()
	if _, err := fmt.Fprintln(f, tokenID); err != nil {
		return fmt.Errorf("write replay state: %w", err)
	}
	return nil
}

func installDesktopBinary(staged string) (string, error) {
	target, err := desktopBinaryPath()
	if err != nil {
		return "", err
	}
	targetDir := filepath.Dir(target)
	if info, err := os.Stat(targetDir); err != nil {
		return "", fmt.Errorf("desktop install dir unavailable: %w", err)
	} else if !info.IsDir() {
		return "", fmt.Errorf("desktop install parent is not a directory: %s", targetDir)
	}

	if info, err := os.Stat(target); err == nil {
		if info.IsDir() {
			return "", fmt.Errorf("desktop target is a directory: %s", target)
		}
		// One generation back: overwrite stale .bak with current before replacing.
		os.Remove(target + ".bak")
		if err := copyRegularFile(target, target+".bak"); err != nil {
			return "", fmt.Errorf("backup desktop binary: %w", err)
		}
	} else if !os.IsNotExist(err) {
		return "", fmt.Errorf("stat desktop binary: %w", err)
	}

	if err := atomicReplaceFile(staged, target, 0o755); err != nil {
		return "", err
	}
	return target, nil
}

func desktopBinaryPath() (string, error) {
	if p := os.Getenv("REDFLAG_DESKTOP_BINARY"); p != "" {
		return p, nil
	}
	execPath, err := os.Executable()
	if err != nil {
		return "", fmt.Errorf("determine agent executable path: %w", err)
	}
	name := "redflag-desktop"
	if runtime.GOOS == "windows" {
		name += ".exe"
	}
	return filepath.Join(filepath.Dir(execPath), name), nil
}

func signalDesktopRestart(target string) error {
	if runtime.GOOS != "linux" {
		return nil
	}
	name := filepath.Base(target)
	cmd := exec.Command("pkill", "-x", name)
	out, err := cmd.CombinedOutput()
	if err == nil {
		log.Printf("[INFO] [agent] [supplychain] desktop_restart_signal_sent binary=%s", name)
		return nil
	}
	if exitErr, ok := err.(*exec.ExitError); ok && exitErr.ExitCode() == 1 {
		log.Printf("[INFO] [agent] [supplychain] desktop_restart_signal_skipped binary=%s reason=not_running", name)
		return nil
	}
	return fmt.Errorf("pkill -x %s failed: %w output=%s", name, err, strings.TrimSpace(string(out)))
}

func atomicReplaceFile(staged, target string, mode os.FileMode) error {
	tmp, err := os.CreateTemp(filepath.Dir(target), filepath.Base(target)+".new.")
	if err != nil {
		return fmt.Errorf("create install temp: %w", err)
	}
	tmpPath := tmp.Name()
	if err := tmp.Close(); err != nil {
		os.Remove(tmpPath)
		return fmt.Errorf("close install temp: %w", err)
	}
	cleanup := true
	defer func() {
		if cleanup {
			os.Remove(tmpPath)
		}
	}()

	if err := copyRegularFile(staged, tmpPath); err != nil {
		return err
	}
	if err := os.Chmod(tmpPath, mode); err != nil {
		return fmt.Errorf("chmod install temp: %w", err)
	}
	if runtime.GOOS == "windows" {
		_ = os.Remove(target)
	}
	if err := os.Rename(tmpPath, target); err != nil {
		return fmt.Errorf("replace binary: %w", err)
	}
	cleanup = false
	return nil
}

func copyRegularFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return fmt.Errorf("open source: %w", err)
	}
	defer in.Close()
	if info, err := in.Stat(); err != nil {
		return fmt.Errorf("stat source: %w", err)
	} else if info.IsDir() {
		return fmt.Errorf("source is a directory: %s", src)
	}

	out, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600)
	if err != nil {
		return fmt.Errorf("open destination: %w", err)
	}
	if _, err := io.Copy(out, in); err != nil {
		out.Close()
		return fmt.Errorf("copy file: %w", err)
	}
	if err := out.Close(); err != nil {
		return fmt.Errorf("close destination: %w", err)
	}
	return nil
}

func computeFileSHA256Hex(path string) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", fmt.Errorf("open for hash: %w", err)
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", fmt.Errorf("hash file: %w", err)
	}
	return hex.EncodeToString(h.Sum(nil)), nil
}

func selfPolicyResult(token *capability.Token, decision, reason string, exitCode, verifiedArtifacts int, resultErr error) *PolicyResult {
	var errText string
	if resultErr != nil {
		errText = resultErr.Error()
	}
	return &PolicyResult{
		TokenID:           token.TokenID,
		AgentID:           token.AgentID,
		PackageType:       token.PackageType,
		Operation:         token.Operation,
		Decision:          decision,
		Reason:            reason,
		Executed:          decision == "executed",
		VerifiedArtifacts: verifiedArtifacts,
		ExitCode:          exitCode,
		Error:             errText,
		Timestamp:         time.Now().UTC().Unix(),
	}
}