github.com/hungdoo/bot@v0.0.0-20240325145135-dd1f386f7b81/src/packages/command/balance/command.go (about)

     1  package balance
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"math/big"
     7  	"strings"
     8  
     9  	ethCommon "github.com/ethereum/go-ethereum/common"
    10  	"github.com/ethereum/go-ethereum/rpc"
    11  	"github.com/hungdoo/bot/src/common"
    12  	command "github.com/hungdoo/bot/src/packages/command/common"
    13  	"github.com/hungdoo/bot/src/packages/log"
    14  	"github.com/hungdoo/bot/src/packages/utils"
    15  	"github.com/shopspring/decimal"
    16  	"go.mongodb.org/mongo-driver/bson"
    17  )
    18  
    19  type Wallet struct {
    20  	Address    ethCommon.Address `bson:"address"`
    21  	PreBalance decimal.Decimal   `bson:"prebalance"`
    22  }
    23  
    24  func (w *Wallet) MarshalBSON() ([]byte, error) {
    25  	return bson.Marshal(&struct {
    26  		Address    string
    27  		PreBalance string
    28  	}{
    29  		Address:    w.Address.String(),
    30  		PreBalance: w.PreBalance.String(),
    31  	})
    32  }
    33  func (w *Wallet) UnmarshalBSON(data []byte) error {
    34  	var d bson.D
    35  	err := bson.Unmarshal(data, &d)
    36  	if err != nil {
    37  		return err
    38  	}
    39  	for _, v := range d {
    40  		if v.Key == "address" {
    41  			w.Address = ethCommon.HexToAddress(v.Value.(string))
    42  		} else if v.Key == "prebalance" {
    43  			w.PreBalance, err = decimal.NewFromString(v.Value.(string))
    44  			if err != nil {
    45  				return err
    46  			}
    47  		}
    48  	}
    49  	return nil
    50  }
    51  
    52  type BalanceCommand struct {
    53  	command.Command
    54  	Id      string   `bson:"_id,unique"`
    55  	Rpc     string   `bson:"rpc"`
    56  	Wallets []Wallet `bson:"wallets"`
    57  }
    58  
    59  func (c BalanceCommand) MarshalJSON() ([]byte, error) {
    60  	addresses := make([]string, len(c.Wallets))
    61  	for i, w := range c.Wallets {
    62  		addresses[i] = w.Address.String()
    63  	}
    64  	return json.Marshal(&struct {
    65  		Name string `json:"name"`
    66  		Type string `json:"type"`
    67  		// Data     []string `json:"data"`
    68  		IdleTime string   `json:"idletime"`
    69  		Rpc      string   `json:"rpc"`
    70  		Wallets  []Wallet `json:"wallets"`
    71  		Command  string   `json:"command"`
    72  	}{
    73  		Name:     c.Name,
    74  		Type:     c.Type.String(),
    75  		IdleTime: c.IdleTime.String(),
    76  		Rpc:      c.Rpc,
    77  		Wallets:  c.Wallets,
    78  		Command:  fmt.Sprintf("add balance %s %s %s", c.Name, c.Rpc, strings.Join(addresses, " ")),
    79  	})
    80  }
    81  
    82  func (c *BalanceCommand) Validate(data []string) error {
    83  	if len(data) < 2 {
    84  		return fmt.Errorf("invalid params: rpc, ...wallets")
    85  	}
    86  	return nil
    87  }
    88  
    89  func (c *BalanceCommand) SetData(newValue []string) (err error) {
    90  	if err = c.Validate(newValue); err != nil {
    91  		return err
    92  	}
    93  	// c.Data = newValue
    94  	c.Rpc = newValue[0]
    95  	for _, v := range newValue[1:] {
    96  		c.Wallets = append(
    97  			c.Wallets,
    98  			Wallet{
    99  				Address:    ethCommon.HexToAddress(v),
   100  				PreBalance: decimal.Zero,
   101  			},
   102  		)
   103  	}
   104  	return nil
   105  }
   106  
   107  func (c *BalanceCommand) Execute(mustReport bool, subcommand string) (string, *common.ErrorWithSeverity) {
   108  	rpcCli, err := rpc.Dial(c.Rpc)
   109  	if err != nil {
   110  		return "", common.NewErrorWithSeverity(common.Error, err.Error())
   111  	}
   112  
   113  	// Create a batch of RPC calls to get the balance of each address
   114  	batch := make([]rpc.BatchElem, len(c.Wallets))
   115  	for i, wallet := range c.Wallets {
   116  		batch[i] = rpc.BatchElem{
   117  			Method: "eth_getBalance",
   118  			Args: []interface{}{
   119  				wallet.Address,
   120  				"latest",
   121  			},
   122  			Result: new(string),
   123  		}
   124  	}
   125  
   126  	// Execute the batch call
   127  	err = rpcCli.BatchCall(batch)
   128  	if err != nil {
   129  		return "", common.NewErrorWithSeverity(common.Error, err.Error())
   130  	}
   131  
   132  	// Process the results
   133  	results := []string{}
   134  	var errors []string
   135  	for i, call := range batch {
   136  		wallet := c.Wallets[i]
   137  		if call.Error != nil {
   138  			errors = append(errors, fmt.Sprintf("address %s: err-%v", wallet.Address, call.Error))
   139  			continue
   140  		}
   141  
   142  		// Assert the type to string
   143  		resultString, ok := call.Result.(*string)
   144  		if !ok {
   145  			errors = append(errors, fmt.Sprintf("unexpected result type for address %s", wallet.Address))
   146  			continue
   147  		}
   148  
   149  		decimalValue, success := new(big.Int).SetString(*resultString, 0)
   150  		if !success {
   151  			errors = append(errors, fmt.Sprintf("error converting balance for address %s: %v", wallet.Address, resultString))
   152  			continue
   153  		}
   154  
   155  		balance := decimal.NewFromBigInt(decimalValue, 0)
   156  		log.GeneralLogger.Printf("%s: %.5f", wallet.Address, utils.DivDecimals(balance, 18).InexactFloat64())
   157  		if mustReport || !mustReport && !balance.Equal(wallet.PreBalance) {
   158  			c.Wallets[i].PreBalance = balance
   159  			results = append(results, fmt.Sprintf("%s\n=> [P:%.5f | V:%.5f]\n", wallet.Address,
   160  				utils.DivDecimals(wallet.PreBalance, 18).InexactFloat64(),
   161  				utils.DivDecimals(balance, 18).InexactFloat64()))
   162  		}
   163  	}
   164  
   165  	if len(errors) != 0 {
   166  		log.GeneralLogger.Printf("%s", strings.Join(errors, "\n"))
   167  		c.SetError(strings.Join(errors, "\n"))
   168  	}
   169  	return strings.Join(results, ""), nil
   170  }