From 970c65313e3251598bf831d69451a9f08216bcf5 Mon Sep 17 00:00:00 2001
From: Jacob Shufro <jacob@shuf.ro>
Date: Thu, 26 Sep 2024 19:02:34 -0400
Subject: [PATCH] Add cli utility for serializing NetworkState at a given slot

---
 go.mod                               |   2 +-
 shared/services/state/cli/.gitignore |   2 +
 shared/services/state/cli/cli.go     | 104 +++++++++++++++++++++++++++
 3 files changed, 107 insertions(+), 1 deletion(-)
 create mode 100644 shared/services/state/cli/.gitignore
 create mode 100644 shared/services/state/cli/cli.go

diff --git a/go.mod b/go.mod
index d59fbc313..5e5267825 100644
--- a/go.mod
+++ b/go.mod
@@ -178,4 +178,4 @@ replace github.com/wealdtech/go-merkletree v1.0.1-0.20190605192610-2bb163c2ea2a
 
 // replace github.com/wealdtech/go-eth2-types/v2 => github.com/rocket-pool/go-eth2-types/v2 v2.0.0-20230130220714-d88838162252
 
-// replace github.com/rocket-pool/rocketpool-go => ../rocketpool-go
+replace github.com/rocket-pool/rocketpool-go => ../rocketpool-go
diff --git a/shared/services/state/cli/.gitignore b/shared/services/state/cli/.gitignore
new file mode 100644
index 000000000..d369ba9e0
--- /dev/null
+++ b/shared/services/state/cli/.gitignore
@@ -0,0 +1,2 @@
+cli
+*.json
diff --git a/shared/services/state/cli/cli.go b/shared/services/state/cli/cli.go
new file mode 100644
index 000000000..a84a86c25
--- /dev/null
+++ b/shared/services/state/cli/cli.go
@@ -0,0 +1,104 @@
+package main
+
+import (
+	"encoding/json"
+	"flag"
+	"fmt"
+	"os"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/ethclient"
+	"github.com/rocket-pool/rocketpool-go/rocketpool"
+	"github.com/rocket-pool/smartnode/shared/services/beacon/client"
+	"github.com/rocket-pool/smartnode/shared/services/config"
+	"github.com/rocket-pool/smartnode/shared/services/state"
+	cfgtypes "github.com/rocket-pool/smartnode/shared/types/config"
+)
+
+// A basic CLI tool which can be used to serialize NetworkState objects to files
+// for future use.
+// Accepts arguments for a beacon node URL, an execution node URL, and a slot number
+// to get the state for.
+
+var bnFlag = flag.String("b", "http://localhost:5052", "The beacon node URL")
+var elFlag = flag.String("e", "http://localhost:8545", "The execution node URL")
+var slotFlag = flag.Uint64("slot", 0, "The slot number to get the state for")
+var networkFlag = flag.String("network", "mainnet", "The network to get the state for, i.e. 'mainnet' or 'holesky'")
+var prettyFlag = flag.Bool("p", false, "Pretty print the output")
+
+var validateFlag = flag.Bool("validate", false, "Validate the json from stdin can be unmarshalled into a NetworkState")
+
+func validate() {
+	decoder := json.NewDecoder(os.Stdin)
+	networkState := state.NetworkState{}
+	err := decoder.Decode(&networkState)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error decoding network state: %v\n", err)
+		os.Exit(1)
+	}
+	fmt.Fprintf(os.Stderr, "Network state validated\n")
+}
+
+func main() {
+	flag.Parse()
+
+	if *validateFlag {
+		validate()
+		return
+	}
+
+	sn := config.NewSmartnodeConfig(nil)
+	switch *networkFlag {
+	case "mainnet":
+		sn.Network.Value = cfgtypes.Network_Mainnet
+	case "holesky":
+		sn.Network.Value = cfgtypes.Network_Holesky
+	default:
+		fmt.Fprintf(os.Stderr, "Invalid network: %s\n", *networkFlag)
+		fmt.Fprintf(os.Stderr, "Valid networks are: mainnet, holesky\n")
+		os.Exit(1)
+	}
+
+	ec, err := ethclient.Dial(*elFlag)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error connecting to execution node: %v\n", err)
+		os.Exit(1)
+	}
+
+	contracts := sn.GetStateManagerContracts()
+	fmt.Fprintf(os.Stderr, "Contracts: %+v\n", contracts)
+
+	rocketStorage := sn.GetStorageAddress()
+
+	rp, err := rocketpool.NewRocketPool(ec, common.HexToAddress(rocketStorage))
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error creating rocketpool: %v\n", err)
+		os.Exit(1)
+	}
+	bc := client.NewStandardHttpClient(*bnFlag)
+	sm := state.NewNetworkStateManager(rp, contracts, bc, nil)
+
+	var networkState *state.NetworkState
+
+	if *slotFlag == 0 {
+		fmt.Fprintf(os.Stderr, "Slot number not provided, defaulting to head slot.\n")
+		networkState, err = sm.GetHeadState()
+	} else {
+		networkState, err = sm.GetStateForSlot(*slotFlag)
+	}
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error getting network state: %v\n", err)
+		os.Exit(1)
+	}
+
+	fmt.Fprintf(os.Stderr, "Network state fetched, outputting to stdout\n")
+	encoder := json.NewEncoder(os.Stdout)
+	if *prettyFlag {
+		encoder.SetIndent("", "  ")
+	}
+	err = encoder.Encode(networkState)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error encoding network state: %v\n", err)
+		os.Exit(1)
+	}
+}