feat: address book for contract

zsystm · Nov 13, 2024 · 48ff902 · 48ff902
1 parent 7dc9a34
commit 48ff902
Show file tree

Hide file tree

Showing 4 changed files with 123 additions and 20 deletions.
diff --git a/cmd/solizard/embeds/address_book.json b/cmd/solizard/embeds/address_book.json
@@ -0,0 +1,5 @@
+{
+  "ERC20": {
+    "address": "0xdAC17F958D2ee523a2206206994597C13D831ec7"
+  }
+}
diff --git a/cmd/solizard/solizard.go b/cmd/solizard/solizard.go
@@ -9,12 +9,13 @@ import (
 	"time"
 
 	"github.com/ethereum/go-ethereum"
+	"github.com/ethereum/go-ethereum/accounts/abi"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/ethclient"
 
-	"github.com/zsystm/solizard/internal/abi"
+	internalabi "github.com/zsystm/solizard/internal/abi"
 	"github.com/zsystm/solizard/internal/config"
 	"github.com/zsystm/solizard/internal/ctx"
 	"github.com/zsystm/solizard/internal/prompt"
@@ -29,11 +30,13 @@ const SolizardDir = ".solizard"
 
 var (
 	// AbiDir is the directory where all abi files are stored
-	// default is $HOME/solizard/abis
-	AbiDir      = "abis"
-	ZeroAddr    = common.Address{}
-	ConfigExist = false
-	conf        *config.Config
+	// default is $HOME/.solizard/abis
+	AbiDir        = "abis"
+	ZeroAddr      = common.Address{}
+	ConfigExist   = false
+	AddrBookExist = false
+	conf          *config.Config
+	AddrBook      config.AddressBook
 )
 
 func init() {
@@ -76,6 +79,11 @@ func init() {
 					return err
 				}
 			}
+			if d.Type().IsRegular() && d.Name() == "address_book.json" {
+				if err := os.WriteFile(homeDir+"/"+SolizardDir+"/address_book.json", data, 0644); err != nil {
+					return err
+				}
+			}
 			return nil
 		}); err != nil {
 			fmt.Printf("failed to walk embedded files (reason: %v)\n", err)
@@ -89,17 +97,29 @@ func init() {
 		os.Exit(1)
 	}
 	ConfigPath := dir + "/" + "config.toml"
-	conf, err = config.ReadConfig(ConfigPath)
-	if err != nil {
+	if conf, err = config.ReadConfig(ConfigPath); err != nil {
 		fmt.Printf("failed to read config file (reason: %v)\n", err)
 		ConfigExist = false
+	} else {
+		ConfigExist = true
+	}
+
+	AddrBookPath := dir + "/" + "address_book.json"
+	if AddrBook, err = config.ReadAddressBook(AddrBookPath); err != nil {
+		fmt.Printf("failed to read address book (reason: %v)\n", err)
+		AddrBookExist = false
+	} else {
+		if err = AddrBook.Validate(); err != nil {
+			fmt.Printf("address book is invalid (reason: %v)\n", err)
+			panic(err)
+		}
+		AddrBookExist = true
 	}
-	ConfigExist = true
 }
 
 func Run() error {
-	fmt.Println(`🦎 Welcome to Solizard v1.4.1 🦎`)
-	mAbi, err := abi.LoadABIs(AbiDir)
+	fmt.Println(`🦎 Welcome to Solizard v1.5.0 🦎`)
+	mAbi, err := internalabi.LoadABIs(AbiDir)
 	if err != nil {
 		return err
 	}
@@ -111,10 +131,12 @@ func Run() error {
 		}
 	}
 
+	var selectedContractName string
+	var selectedAbi abi.ABI
 	// start the main loop
 	for {
 	STEP_SELECT_CONTRACT:
-		selectedAbi := prompt.MustSelectContractABI(mAbi)
+		selectedContractName, selectedAbi = prompt.MustSelectContractABI(mAbi)
 	INPUT_RPC_URL:
 		if sctx.EthClient() == nil {
 			rpcURL := prompt.MustInputRpcUrl()
@@ -126,14 +148,28 @@ func Run() error {
 			sctx.SetEthClient(client)
 		}
 	INPUT_CONTRACT_ADDRESS:
-		contractAddress := prompt.MustInputContractAddress()
+		// check if address book exists
+		useAddrBook := false
+		var contractAddress string
+		if AddrBookExist {
+			if value, exists := AddrBook[selectedContractName]; exists {
+				if yes := prompt.MustSelectAddressBookUsage(value.Address); yes {
+					contractAddress = value.Address
+					useAddrBook = true
+				}
+			}
+		}
+		if !useAddrBook {
+			contractAddress = prompt.MustInputContractAddress()
+		}
 		if err := validation.ValidateContractAddress(sctx, contractAddress); err != nil {
 			fmt.Printf("Invalid contract address (reason: %v)\n", err)
 			goto INPUT_CONTRACT_ADDRESS
 		}
+
 	SELECT_METHOD:
 		rw := prompt.MustSelectReadOrWrite()
-		if rw == abi.WriteMethod {
+		if rw == internalabi.WriteMethod {
 			// input private key
 			if sctx.PrivateKey() == nil {
 				pk := prompt.MustInputPrivateKey()
@@ -147,7 +183,7 @@ func Run() error {
 		}
 		methodName, method := prompt.MustSelectMethod(selectedAbi, rw)
 		input := prompt.MustCreateInputDataForMethod(method)
-		if rw == abi.ReadMethod {
+		if rw == internalabi.ReadMethod {
 			callMsg := ethereum.CallMsg{From: ZeroAddr, To: sctx.ContractAddress(), Data: input}
 			output, err := sctx.EthClient().CallContract(context.TODO(), callMsg, nil)
 			if err != nil {

diff --git a/internal/config/address_book.go b/internal/config/address_book.go
@@ -0,0 +1,44 @@
+package config
+
+import (
+	"encoding/json"
+	"fmt"
+	"os"
+
+	"github.com/ethereum/go-ethereum/common"
+)
+
+type AddressBook map[string]struct {
+	Address string `json:"address"`
+}
+
+// ReadAddressBook reads the address book from the given path
+func ReadAddressBook(path string) (AddressBook, error) {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return nil, err
+	}
+
+	var ab AddressBook
+	if err = json.Unmarshal(data, &ab); err != nil {
+		return nil, err
+	}
+	return ab, nil
+}
+
+func (ab AddressBook) Validate() error {
+	// check if the address book is empty
+	if len(ab) == 0 {
+		return fmt.Errorf("address book is empty")
+	}
+	for name, entry := range ab {
+		if entry.Address == "" {
+			return fmt.Errorf("address book entry %s has empty address", name)
+		}
+		if !common.IsHexAddress(entry.Address) {
+			return fmt.Errorf("address book entry %s has invalid address", name)
+		}
+	}
+
+	return nil
+}
diff --git a/internal/prompt/prompt.go b/internal/prompt/prompt.go
@@ -23,14 +23,17 @@ const DefaultPromptListSize = 10
 // which means the user wants to apply the config file
 func MustSelectApplyConfig() bool {
 	prompt := promptui.Prompt{
-		Label:     "found config file, but do you want to setup manually",
-		IsConfirm: true,
+		Label: "found config file at ~/.solizard/config.toml, apply? [Y/n]",
 	}
 	ret, _ := prompt.Run()
-	return strings.ToLower(ret) == "n"
+	if NoSelected(ret) {
+		return false
+	}
+	return true
 }
 
-func MustSelectContractABI(abis map[string]abi.ABI) abi.ABI {
+// MustSelectContractABI prompts the user to select a contract ABI and returns the selected contract name and ABI
+func MustSelectContractABI(abis map[string]abi.ABI) (string, abi.ABI) {
 	contractNames := make([]string, 0, len(abis))
 	for name := range abis {
 		contractNames = append(contractNames, name)
@@ -51,7 +54,7 @@ func MustSelectContractABI(abis map[string]abi.ABI) abi.ABI {
 	if err != nil {
 		panic(err)
 	}
-	return abis[selected]
+	return strings.TrimSuffix(selected, ".abi"), abis[selected]
 }
 
 func MustInputRpcUrl() string {
@@ -68,6 +71,17 @@ func MustInputRpcUrl() string {
 	return rpcURL
 }
 
+func MustSelectAddressBookUsage(contractAddr string) bool {
+	prompt := promptui.Prompt{
+		Label: fmt.Sprintf("Use %s as contract address? [Y/n]", contractAddr),
+	}
+	ret, _ := prompt.Run()
+	if NoSelected(ret) {
+		return false
+	}
+	return true
+}
+
 func MustInputContractAddress() string {
 	prompt := promptui.Prompt{
 		Label:    "Enter the contract address",
@@ -249,3 +263,7 @@ const SelectableListSize = 4
 func shouldSupportSearchMode(listLen int) bool {
 	return listLen > SelectableListSize
 }
+
+func NoSelected(s string) bool {
+	return strings.ToLower(s) == "n"
+}