LINUX-AETHERFORGE/cmd/agent/main.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	"forge-mesh/internal/api/types"
	"forge-mesh/internal/mining"
	"forge-mesh/internal/policy"

	"github.com/gorilla/websocket"
)

const agentVersion = "0.1.0-dev"

func main() {
	deckURL := flag.String("deck-url", "", "deck base URL (alias: -deck, -server)")
	deck := flag.String("deck", "", "deck base URL")
	server := flag.String("server", "", "deck base URL (deprecated alias)")
	secret := flag.String("secret", "", "fleet bearer secret")
	pubKey := flag.String("pubkey", "", "ed25519 public key hex for signed self-update")
	hostID := flag.String("host-id", "", "stable host id (optional, persisted under data dir)")
	stratumHost := flag.String("stratum-host", "127.0.0.1", "stratum proxy host")
	wallet := flag.String("wallet", "", "fallback wallet when no server profile")
	dataDir := flag.String("data-dir", defaultDataDir(), "agent state directory")
	flag.Parse()

	baseURL := coalesceEnv(
		*deckURL, *deck, *server,
		os.Getenv("FORGE_MESH_DECK_URL"),
		os.Getenv("FORGE_DECK_URL"),
		"http://127.0.0.1:8989",
	)
	baseURL = strings.TrimRight(baseURL, "/")

	*secret = coalesceEnv(*secret, os.Getenv("FORGE_MESH_FLEET_SECRET"), os.Getenv("FORGE_FLEET_SECRET"))
	if *secret == "" {
		log.Fatal("fleet secret required (-secret, FORGE_MESH_FLEET_SECRET, or FORGE_FLEET_SECRET)")
	}

	*pubKey = coalesceEnv(*pubKey, os.Getenv("FORGE_MESH_PUBKEY"), os.Getenv("FORGE_PUBKEY"))

	hostname, _ := os.Hostname()
	if err := os.MkdirAll(*dataDir, 0o755); err != nil {
		log.Fatalf("data dir: %v", err)
	}

	if *hostID == "" {
		*hostID = loadHostID(*dataDir)
	}
	if *hostID == "" {
		*hostID = registerHost(baseURL, *secret, hostname)
	}
	if *hostID != "" {
		saveHostID(*dataDir, *hostID)
	}

	if *wallet == "" {
		*wallet = os.Getenv("FORGE_WALLET")
	}

	profile := policy.DefaultMiningProfile(*wallet)
	chainCfg := mining.DefaultChainConfig()
	chainCfg.StratumHost = *stratumHost
	chainCfg.HostID = *hostID
	chain := mining.NewChain(profile, chainCfg)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		if err := chain.Run(ctx); err != nil && ctx.Err() == nil {
			log.Printf("mining chain stopped: %v", err)
		}
	}()

	agent := &Agent{
		ctx:       ctx,
		serverURL: baseURL,
		secret:    *secret,
		hostID:    *hostID,
		hostname:  hostname,
		pubKeyHex: *pubKey,
		chain:     chain,
	}

	go agent.run(ctx)
	go agent.selfUpdateLoop(ctx)

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	<-sig
	cancel()
	chain.Stop()
	wg.Wait()
}

type Agent struct {
	ctx       context.Context
	serverURL string
	secret    string
	hostID    string
	hostname  string
	pubKeyHex string
	chain     *mining.Chain

	mu   sync.Mutex
	conn *websocket.Conn
}

func (a *Agent) run(ctx context.Context) {
	backoff := time.Second
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		if err := a.connect(ctx); err != nil {
			log.Printf("agent ws: %v (retry in %s)", err, backoff)
			a.beaconFallback()
			select {
			case <-ctx.Done():
				return
			case <-time.After(backoff):
			}
			if backoff < 30*time.Second {
				backoff *= 2
			}
			continue
		}
		backoff = time.Second
	}
}

func (a *Agent) selfUpdateLoop(ctx context.Context) {
	if a.pubKeyHex == "" {
		return
	}
	ticker := time.NewTicker(30 * time.Minute)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			a.maybeSelfUpdate(a.pubKeyHex)
		}
	}
}

func (a *Agent) connect(ctx context.Context) error {
	wsURL := strings.Replace(a.serverURL, "http://", "ws://", 1)
	wsURL = strings.Replace(wsURL, "https://", "wss://", 1)
	wsURL += "/api/v1/ws/fleet?token=" + a.secret

	header := http.Header{}
	header.Set("Authorization", "Bearer "+a.secret)

	dialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second}
	conn, _, err := dialer.DialContext(ctx, wsURL, header)
	if err != nil {
		return err
	}

	a.mu.Lock()
	a.conn = conn
	a.mu.Unlock()
	defer func() {
		a.mu.Lock()
		a.conn = nil
		a.mu.Unlock()
		conn.Close()
	}()

	log.Printf("agent %s (%s) connected to %s", a.hostname, a.hostID, a.serverURL)

	if err := a.sendHeartbeat(); err != nil {
		return err
	}

	hbTicker := time.NewTicker(15 * time.Second)
	defer hbTicker.Stop()

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-hbTicker.C:
			if err := a.sendHeartbeat(); err != nil {
				return err
			}
		default:
		}

		_ = conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		_, data, err := conn.ReadMessage()
		if err != nil {
			return err
		}

		var msg types.WsMessage
		if err := json.Unmarshal(data, &msg); err != nil {
			continue
		}
		switch msg.Type {
		case "command":
			if msg.Command != nil {
				a.handleCommand(*msg.Command)
			}
		case "mining_profile":
			if profileRaw, ok := msg.Payload["profile"]; ok {
				b, _ := json.Marshal(profileRaw)
				var profile types.MiningProfile
				if json.Unmarshal(b, &profile) == nil {
					a.handleCommand(types.FleetCommand{Action: "mining_profile", Args: map[string]any{"profile": profileRaw}})
				}
			}
		}
	}
}

func (a *Agent) sendHeartbeat() error {
	st := a.chain.Status()
	hb := types.HeartbeatPayload{
		HostID:      a.hostID,
		Hostname:    a.hostname,
		Arch:        runtime.GOARCH,
		CurrentTier: st.CurrentTier,
		TierType:    st.TierType,
		TierState:   st.TierState,
	}
	hb.SetHashrateFields(st.HashrateHps)

	payload, _ := json.Marshal(map[string]any{
		"type":         "heartbeat",
		"host_id":      hb.HostID,
		"hostname":     hb.Hostname,
		"arch":         hb.Arch,
		"hashrate":     hb.Hashrate,
		"hashrate_hps": hb.HashrateHps,
		"current_tier": hb.CurrentTier,
		"tier_type":    hb.TierType,
		"tier_state":   hb.TierState,
	})

	a.mu.Lock()
	conn := a.conn
	a.mu.Unlock()
	if conn == nil {
		return nil
	}
	return conn.WriteMessage(websocket.TextMessage, payload)
}

func (a *Agent) handleCommand(cmd types.FleetCommand) {
	var cmdErr error
	switch cmd.Action {
	case "mining_profile":
		raw, ok := cmd.Args["profile"]
		if !ok {
			cmdErr = errCommand("missing profile")
			break
		}
		b, _ := json.Marshal(raw)
		var profile types.MiningProfile
		if err := json.Unmarshal(b, &profile); err != nil {
			cmdErr = err
			break
		}
		a.chain.UpdateProfile(profile)
		go func() {
			if err := a.chain.Run(a.ctx); err != nil && a.ctx.Err() == nil {
				log.Printf("mining chain restart: %v", err)
			}
		}()
		log.Printf("agent: mining profile updated wallet=%s tiers=%d", profile.WalletAddress, len(profile.Tiers))
	case "pause":
		a.chain.Stop()
	case "resume":
		go func() {
			if err := a.chain.Run(a.ctx); err != nil && a.ctx.Err() == nil {
				log.Printf("mining chain resume: %v", err)
			}
		}()
	case "reboot":
		cmdErr = exec.Command("systemctl", "reboot").Run()
	case "screenshot":
		cmdErr = a.takeScreenshot()
	default:
		log.Printf("unknown command: %s", cmd.Action)
		cmdErr = errCommand("unknown action")
	}
	a.sendCommandAck(cmd, cmdErr)
}

func (a *Agent) sendCommandAck(cmd types.FleetCommand, err error) {
	ack := map[string]any{
		"type":    "command_ack",
		"command": cmd.Action,
		"ok":      err == nil,
	}
	if err != nil {
		ack["error"] = err.Error()
	}
	payload, _ := json.Marshal(ack)

	a.mu.Lock()
	conn := a.conn
	a.mu.Unlock()
	if conn == nil {
		return
	}
	_ = conn.WriteMessage(websocket.TextMessage, payload)
}

func (a *Agent) beaconFallback() {
	st := a.chain.Status()
	hb := types.HeartbeatPayload{
		HostID:      a.hostID,
		Hostname:    a.hostname,
		Arch:        runtime.GOARCH,
		CurrentTier: st.CurrentTier,
		TierType:    st.TierType,
		TierState:   st.TierState,
	}
	hb.SetHashrateFields(st.HashrateHps)

	body, _ := json.Marshal(hb)
	for _, path := range []string{"/api/v1/fleet/beacon", "/api/v1/beacon"} {
		req, err := http.NewRequest(http.MethodPost, a.serverURL+path, bytes.NewReader(body))
		if err != nil {
			continue
		}
		req.Header.Set("Authorization", "Bearer "+a.secret)
		req.Header.Set("Content-Type", "application/json")

		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			continue
		}
		if resp.StatusCode != http.StatusOK {
			resp.Body.Close()
			continue
		}

		var br types.BeaconResponse
		if err := json.NewDecoder(resp.Body).Decode(&br); err != nil {
			resp.Body.Close()
			continue
		}
		resp.Body.Close()
		for _, cmd := range br.Commands {
			a.handleCommand(cmd)
		}
		return
	}
}

func (a *Agent) takeScreenshot() error {
	for _, tool := range []string{"grim", "scrot", "import"} {
		if _, err := exec.LookPath(tool); err != nil {
			continue
		}
		out := "/tmp/forge-mesh-screenshot.png"
		var cmd *exec.Cmd
		switch tool {
		case "grim":
			cmd = exec.Command("grim", out)
		case "scrot":
			cmd = exec.Command("scrot", out)
		default:
			cmd = exec.Command("import", "-window", "root", out)
		}
		if err := cmd.Run(); err == nil {
			log.Printf("screenshot saved to %s", out)
			return nil
		}
	}
	log.Println("screenshot: no tool available (grim/scrot/import)")
	return nil
}

func registerHost(baseURL, secret, hostname string) string {
	body, _ := json.Marshal(map[string]string{"hostname": hostname})
	req, err := http.NewRequest(http.MethodPost, baseURL+"/api/v1/fleet/register", bytes.NewReader(body))
	if err != nil {
		return ""
	}
	req.Header.Set("Authorization", "Bearer "+secret)
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil || resp.StatusCode != http.StatusOK {
		return ""
	}
	defer resp.Body.Close()
	var out struct {
		HostID string `json:"host_id"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return ""
	}
	return out.HostID
}

func defaultDataDir() string {
	if runtime.GOOS == "linux" {
		if _, err := os.Stat("/opt/forge-mesh"); err == nil {
			return "/opt/forge-mesh"
		}
	}
	home, _ := os.UserHomeDir()
	return filepath.Join(home, ".forge-mesh")
}

func hostIDPath(dataDir string) string {
	return filepath.Join(dataDir, "host.id")
}

func loadHostID(dataDir string) string {
	data, err := os.ReadFile(hostIDPath(dataDir))
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(data))
}

func saveHostID(dataDir, id string) {
	_ = os.WriteFile(hostIDPath(dataDir), []byte(id), 0o644)
}

func coalesceEnv(values ...string) string {
	for _, v := range values {
		if strings.TrimSpace(v) != "" {
			return strings.TrimSpace(v)
		}
	}
	return ""
}

type commandError string

func (e commandError) Error() string { return string(e) }

func errCommand(msg string) error { return commandError(msg) }