kisstron/ops/ops.go

// Package ops implements the kisstron wallet operations (new, import, info,
// send, freeze, unfreeze, broadcast) mirroring the functions in kisstron.py.
package ops

import (
  "errors"
  "fmt"
  "math"
  "math/big"
  "strings"
  "time"

  "code.luxferre.top/luxferre/kisstron/cli"
  "code.luxferre.top/luxferre/kisstron/tron"
)

// Config carries the resolved command-line parameters shared by all operations.
type Config struct {
  Network    string
  Node       string
  APIKey     string
  FeeLimit   int64
  Passphrase string
  WordsCount int
  Derivation string
}

// client builds a TRON client for the operation's network, reproducing the
// tronclient() endpoint and fee-limit selection logic of kisstron.py.
func (c Config) client() *tron.Client {
  var endpoint string
  fee := int64(tron.FeeLimitTestnet)
  if c.Network == "mainnet" {
    fee = c.FeeLimit
    switch {
    case c.APIKey != "":
      endpoint = "https://api.trongrid.io"
    case c.Node != "":
      endpoint = c.Node
    default:
      endpoint = tron.DefaultNodeAddr
    }
  } else {
    endpoint = tron.DefaultEndpoint(c.Network)
  }
  return tron.NewClient(c.Network, endpoint, fee, c.APIKey)
}

// NewWallet creates a fresh wallet and returns its details.
func NewWallet(cfg Config) cli.Result {
  if !tron.ValidWordCount(cfg.WordsCount) {
    return cli.Result{Message: "Invalid word count for mnemonic generation"}
  }
  key, mnemonic, err := tron.GenerateWithMnemonic(cfg.Passphrase, cfg.WordsCount, cfg.Derivation)
  if err != nil {
    return cli.Result{Message: err.Error()}
  }
  return walletResult(key, mnemonic, "Wallet created, save this info in a secure place!")
}

// ImportWallet restores a wallet from a mnemonic.
func ImportWallet(cfg Config, mnemonic string) cli.Result {
  if !tron.ValidateMnemonic(mnemonic) {
    return cli.Result{Message: "Invalid mnemonic phrase!"}
  }
  key, err := tron.DeriveFromMnemonic(mnemonic, cfg.Passphrase, cfg.Derivation)
  if err != nil {
    return cli.Result{Message: "Invalid mnemonic phrase!"}
  }
  return walletResult(key, strings.Join(strings.Fields(mnemonic), " "), "Wallet imported, save this info in a secure place!")
}

func walletResult(key *tron.PrivateKey, mnemonic, banner string) cli.Result {
  info := tron.GenerateAddressInfo(key)
  msg := banner + "\n\n"
  msg += fmt.Sprintf("Address: %s\n", info["base58check_address"])
  msg += fmt.Sprintf("Address (hex): %s\n", info["hex_address"])
  msg += fmt.Sprintf("Private key (hex): %s\n", info["private_key"])
  msg += fmt.Sprintf("Public key (hex): %s\n", info["public_key"])
  msg += fmt.Sprintf("Mnemonic: %s", mnemonic)
  return cli.Result{Success: true, Message: msg, Result: info}
}

// frozenEntry is a single TRX freeze record.
type frozenEntry struct {
  Type   string
  Amount int64
}

// InfoByAddr displays account information for a TRON address.
func InfoByAddr(cfg Config, address string) cli.Result {
  client := cfg.client()
  rawInfo, err := client.GetAccount(address)
  if err != nil {
    switch {
    case errors.Is(err, tron.ErrBadAddress):
      return cli.Result{Message: "Invalid address! Set the correct one with -addr flag"}
    case errors.Is(err, tron.ErrAddressNotFound):
      return cli.Result{Message: "Address not found on the blockchain! Set the correct one with -addr flag"}
    default:
      return cli.Result{Message: err.Error()}
    }
  }
  rawRes, err := client.GetAccountResource(address)
  if err != nil {
    return cli.Result{Message: err.Error()}
  }

  // Bandwidth computation. NOTE: this faithfully reproduces the original
  // kisstron.py behaviour, including the latent bug where freeNetUsed is
  // assigned to the freeNetLimit variable (line 210 of kisstron.py), so wallet
  // output stays byte-identical to the Python implementation.
  var freeNetLimit, freeNetUsed, netLimit, netUsed int64
  if v, ok := rawRes["freeNetLimit"]; ok {
    freeNetLimit = tron.ValueInt64(v)
  }
  if v, ok := rawRes["freeNetUsed"]; ok {
    freeNetLimit = tron.ValueInt64(v)
  }
  if v, ok := rawRes["NetLimit"]; ok {
    netLimit = tron.ValueInt64(v)
  }
  if v, ok := rawRes["NetUsed"]; ok {
    netUsed = tron.ValueInt64(v)
  }
  bandwidth := freeNetLimit - freeNetUsed + netLimit - netUsed

  var energyLimit, energyUsed int64
  if v, ok := rawRes["EnergyLimit"]; ok {
    energyLimit = tron.ValueInt64(v)
  }
  if v, ok := rawRes["EnergyUsed"]; ok {
    energyUsed = tron.ValueInt64(v)
  }
  energy := energyLimit - energyUsed

  addressOut := tron.ValueString(rawInfo["address"])
  trxBalance := tron.ValueInt64(rawInfo["balance"])

  var frozen []frozenEntry
  if arr, ok := rawInfo["frozenV2"].([]any); ok {
    for _, e := range arr {
      m, ok := e.(map[string]any)
      if !ok {
        continue
      }
      t, hasT := m["type"]
      a, hasA := m["amount"]
      if !hasT || !hasA {
        continue
      }
      amount := tron.ValueInt64(a)
      if amount > 0 {
        frozen = append(frozen, frozenEntry{Type: tron.ValueString(t), Amount: amount})
      }
    }
  }

  type tokenBalance struct {
    ticker string
    value  float64
  }
  var balances []tokenBalance
  for _, token := range tron.TokensForNetwork(cfg.Network) {
    if !token.ShowBalance {
      continue
    }
    cntr, err := client.GetContract(token.Contract)
    if err != nil {
      return cli.Result{Message: err.Error()}
    }
    value, err := trc20Balance(cntr, address)
    if err != nil {
      return cli.Result{Message: err.Error()}
    }
    balances = append(balances, tokenBalance{ticker: token.Ticker, value: value})
  }

  msg := fmt.Sprintf("Address: %s\nBalance:\n", addressOut)
  msg += fmt.Sprintf("%f TRX\n", float64(trxBalance)/tron.TrxScale)
  for _, b := range balances {
    msg += fmt.Sprintf("%f %s\n", b.value, b.ticker)
  }
  msg += "Frozen assets:\n"
  for _, f := range frozen {
    msg += fmt.Sprintf("%s\t%f TRX\n", f.Type, float64(f.Amount)/tron.TrxScale)
  }
  msg += fmt.Sprintf("Remaining bandwidth: %d\n", bandwidth)
  msg += fmt.Sprintf("Remaining energy: %d", energy)
  return cli.Result{Success: true, Message: msg}
}

// trc20Balance returns the float token balance of addr for a contract.
func trc20Balance(cntr *tron.Contract, addr string) (float64, error) {
  precision, err := cntr.Decimals()
  if err != nil {
    return 0, err
  }
  bal, err := cntr.BalanceOf(addr)
  if err != nil {
    return 0, err
  }
  divisor := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(precision)), nil)
  f := new(big.Float).SetInt(bal)
  f.Quo(f, new(big.Float).SetInt(divisor))
  out, _ := f.Float64()
  return out, nil
}

// InfoByPK displays own account information derived from a private key.
func InfoByPK(cfg Config, pk string) cli.Result {
  key, err := tron.PrivateKeyFromHex(pk)
  if err != nil {
    return cli.Result{Message: "Invalid private key!"}
  }
  return InfoByAddr(cfg, key.Pub.ToBase58CheckAddress())
}

// SendTrx sends TRX (or produces a signed offline transaction).
func SendTrx(cfg Config, pk, toaddr string, amt float64, offline bool) cli.Result {
  if amt < 1 {
    return cli.Result{Message: "Nothing to transfer!"}
  }
  key, err := tron.PrivateKeyFromHex(pk)
  if err != nil {
    return cli.Result{Message: "Invalid private key!"}
  }
  builder, err := cfg.client().Transfer(key.Pub.ToBase58CheckAddress(), toaddr, int64(amt*tron.TrxScale))
  if err != nil {
    return cli.Result{Message: err.Error()}
  }
  tx, err := builder.Build()
  if err != nil {
    return cli.Result{Message: err.Error()}
  }
  if err := tx.Sign(key); err != nil {
    return cli.Result{Message: err.Error()}
  }
  if offline {
    return cli.Result{Success: true, Message: tx.JSON(), Result: tx}
  }
  return broadcastAndWait(tx, fmt.Sprintf("Sent %f TRX to %s", amt, toaddr))
}

// SendTrc20 sends a supported TRC20 token.
func SendTrc20(cfg Config, pk, ticker, toaddr string, amt float64, offline bool) cli.Result {
  ticker = strings.ToLower(ticker)
  token, ok := tron.TokenByKey(cfg.Network, ticker)
  if !ok {
    return cli.Result{Message: fmt.Sprintf("Token not supported by kisstron for %s!", cfg.Network)}
  }
  if amt < 1 {
    return cli.Result{Message: "Nothing to transfer!"}
  }
  key, err := tron.PrivateKeyFromHex(pk)
  if err != nil {
    return cli.Result{Message: "Invalid private key!"}
  }
  client := cfg.client()
  cntr, err := client.GetContract(token.Contract)
  if err != nil {
    return cli.Result{Message: err.Error()}
  }
  precision, err := cntr.Decimals()
  if err != nil {
    return cli.Result{Message: err.Error()}
  }
  amount := new(big.Int).SetInt64(int64(amt * math.Pow(10, float64(precision))))
  builder, err := cntr.Transfer(key.Pub.ToBase58CheckAddress(), toaddr, amount)
  if err != nil {
    return cli.Result{Message: err.Error()}
  }
  tx, err := builder.Build()
  if err != nil {
    return cli.Result{Message: err.Error()}
  }
  if err := tx.Sign(key); err != nil {
    return cli.Result{Message: err.Error()}
  }
  if offline {
    return cli.Result{Success: true, Message: tx.JSON(), Result: tx}
  }
  return broadcastAndWait(tx, fmt.Sprintf("Sent %f %s to %s", amt, token.Ticker, toaddr))
}

// FreezeTrx freezes TRX for energy.
func FreezeTrx(cfg Config, pk string, amt float64, offline bool) cli.Result {
  if amt < 1 {
    return cli.Result{Message: "Nothing to freeze!"}
  }
  key, err := tron.PrivateKeyFromHex(pk)
  if err != nil {
    return cli.Result{Message: "Invalid private key!"}
  }
  myaddr := key.Pub.ToBase58CheckAddress()
  builder, err := cfg.client().FreezeBalance(myaddr, int64(amt*tron.TrxScale), "ENERGY")
  if err != nil {
    return cli.Result{Message: err.Error()}
  }
  tx, err := builder.Build()
  if err != nil {
    return cli.Result{Message: err.Error()}
  }
  if err := tx.Sign(key); err != nil {
    return cli.Result{Message: err.Error()}
  }
  if offline {
    return cli.Result{Success: true, Message: tx.JSON(), Result: tx}
  }
  return broadcastAndWait(tx, fmt.Sprintf("Frozen %f TRX at %s", amt, myaddr))
}

// UnfreezeTrx unfreezes TRX.
func UnfreezeTrx(cfg Config, pk string, amt float64, offline bool) cli.Result {
  if amt < 1 {
    return cli.Result{Message: "Nothing to unfreeze!"}
  }
  key, err := tron.PrivateKeyFromHex(pk)
  if err != nil {
    return cli.Result{Message: "Invalid private key!"}
  }
  myaddr := key.Pub.ToBase58CheckAddress()
  builder, err := cfg.client().UnfreezeBalance(myaddr, int64(amt*tron.TrxScale), "ENERGY")
  if err != nil {
    return cli.Result{Message: err.Error()}
  }
  tx, err := builder.Build()
  if err != nil {
    return cli.Result{Message: err.Error()}
  }
  if err := tx.Sign(key); err != nil {
    return cli.Result{Message: err.Error()}
  }
  if offline {
    return cli.Result{Success: true, Message: tx.JSON(), Result: tx}
  }
  return broadcastAndWait(tx, fmt.Sprintf("Unfrozen %f TRX at %s", amt, myaddr))
}

// BroadcastJSON broadcasts a transaction previously produced in offline mode.
func BroadcastJSON(cfg Config, tjson string) cli.Result {
  tx, err := cfg.client().ParseTransaction(tjson)
  if err != nil {
    return cli.Result{Message: "Malformed transaction JSON!"}
  }
  if _, err := tx.Broadcast(); err != nil {
    return cli.Result{Message: err.Error()}
  }
  receipt, err := tx.Wait(30*time.Second, 1600*time.Millisecond)
  if err != nil {
    return cli.Result{Message: err.Error()}
  }
  if isOutOfEnergy(receipt) {
    return cli.Result{Message: "Not enough energy for transaction!\nTransaction ID " + receiptID(receipt), Result: receipt}
  }
  return cli.Result{Success: true, Message: "Transaction sent\nTransaction ID " + receiptID(receipt), Result: receipt}
}

// broadcastAndWait broadcasts tx, waits for the receipt and formats a success
// message with the given header line.
func broadcastAndWait(tx *tron.Transaction, header string) cli.Result {
  if _, err := tx.Broadcast(); err != nil {
    return cli.Result{Message: err.Error()}
  }
  receipt, err := tx.Wait(30*time.Second, 1600*time.Millisecond)
  if err != nil {
    return cli.Result{Message: err.Error()}
  }
  if isOutOfEnergy(receipt) {
    return cli.Result{Message: "Not enough energy for transaction!\nTransaction ID " + receiptID(receipt), Result: receipt}
  }
  return cli.Result{Success: true, Message: header + "\nTransaction ID " + receiptID(receipt), Result: receipt}
}

func receiptID(receipt map[string]any) string {
  return tron.ValueString(receipt["id"])
}

func isOutOfEnergy(receipt map[string]any) bool {
  r, ok := receipt["receipt"].(map[string]any)
  if !ok {
    return false
  }
  return tron.ValueString(r["result"]) == "OUT_OF_ENERGY"
}