#!/usr/bin/env python3
# kisstron: a simple to use TRON CLI wallet based on tronpy
#
# Supports:
# * BIP39 mnemonic wallet generation/restoration
# * converting private keys to BIP39 mnemonics and vice versa
# * basic TRX parameters display (balance, energy, bandwidth)
# * basic TRX operations (transfer, freeze for energy, unfreeze)
# * basic TRC20 token operations (balance display, transfer) for select tokens
#   (contract addresses are hardcoded, see KT_TRC20_TOKENS dictionary)
# * ability to select testnets (nile, shasta) if necessary
# * ability to connect to custom mainnet nodes if necessary
# * switching to TronGrid API by providing your own API key
# * offline mode with -o switch and bro[adcast] operation
#
# As of now, the main goal is to be able to operate on TRX and stablecoins.
# For mnemonic phrase generation, only English is currently supported.
#
# See README for more information, run with -h flag for usage instructions
#
# Created by Luxferre in 2024, released into public domain

import sys, json
import warnings # to suppress NotOpenSSLWarning on BSDs
warnings.filterwarnings("ignore", module="urllib3")
from tronpy import Tron
from tronpy.tron import Transaction
from tronpy.keys import PrivateKey
from tronpy.providers import HTTPProvider
from tronpy.exceptions import (BadAddress, AddressNotFound,
  BadKey, ValidationError, TransactionError)
from json.decoder import JSONDecodeError
import mnemonic, eth_utils

# some global parameters (hardcoded)
KT_TIMEOUT = 20.0 # network operation timeout
KT_FEE_LIMIT = 100000000 # transaction fee limit for mainnet (in Sun)
KT_TEST_FEE_LIMIT = 1000000000 # transaction fee limit for testnets (in Sun)
TRX_SCALE = 1000000.0 # TRX scale (million Suns in 1 TRX)
TG_API_KEY = None # TronGrid API Key (unused by default)
NODE_ADDR = 'http://18.196.99.16:8090' # public Tron node (Germany)

# TRC20 token definitions (also hardcoded)
# Make sure the mainnet tokens match the contract addresses
# specified on these pages:
# USDT: https://tron.network/usdt
# USDC: https://tron.network/usdc

KT_TRC20_TOKENS = {
  'mainnet': {
    'usdt': {
      'contract': 'TR7NHqjeKQxGTCi8q8ZY4pL8otSzgjLj6t',
      'ticker': 'USDT',
      'show_balance': True
    },
    'usdc': {
      'contract': 'TEkxiTehnzSmSe2XqrBj4w32RUN966rdz8',
      'ticker': 'USDC',
      'show_balance': True
    }
  },
  'nile': {
    'usdt': {
      'contract': 'TXLAQ63Xg1NAzckPwKHvzw7CSEmLMEqcdj',
      'ticker': 'USDT',
      'show_balance': True
    },
    'usdc': {
      'contract': 'TEMVynQpntMqkPxP6wXTW2K7e4sM3cRmWz',
      'ticker': 'USDC',
      'show_balance': True
    }
  },
  'shasta': {
    'usdt': {
      'contract': 'TG3XXyExBkPp9nzdajDZsozEu4BkaSJozs',
      'ticker': 'USDT',
      'show_balance': True
    },
    'usdc': {
      'contract': 'TSdZwNqpHofzP6BsBKGQUWdBeJphLmF6id',
      'ticker': 'USDC',
      'show_balance': True
    }
  }
}

# helper functions

# print operation result according to the below convention
def presult(resdict: dict):
  if 'success' not in resdict:
    resdict['success'] = False
  if 'message' not in resdict:
    resdict['message'] = ''
  if resdict['success'] == True:
    print ('Success:\n%s\n' % resdict['message'])
  else:
    print ('Error:\n%s\n' % resdict['message'])

# confirm transactions
def trconfirm(promptstr: str):
  choice = input(promptstr).lower()
  if choice != 'yes':
    print('Transaction not confirmed, aborting!')
    sys.exit(0)

# instantiate tronpy client
def tronclient(network='mainnet'):
  network = network.lower() # lowercase network name
  if network == 'mainnet': # use a public node or TronGrid API for mainnet
    if TG_API_KEY is not None: # TG_API_KEY must be set
      client = Tron(network=network, conf={'fee_limit': KT_FEE_LIMIT},
        provider=HTTPProvider(api_key=TG_API_KEY, timeout=KT_TIMEOUT))
    else: # NODE_ADDR must be set
      client = Tron(network=network, conf={'fee_limit': KT_FEE_LIMIT},
        provider=HTTPProvider(NODE_ADDR, timeout=KT_TIMEOUT))
  else: # any default API is fine for a testnet
    client = Tron(network=network,
      conf={'fee_limit': KT_TEST_FEE_LIMIT, 'timeout': KT_TIMEOUT})
  return client

# get TRC20 token balance of an address (in float)
def trc20balance(client: Tron, contractaddr: str, targetaddr: str):
  cntr = client.get_contract(contractaddr)
  precision = cntr.functions.decimals()
  return cntr.functions['balanceOf'](targetaddr) / (10 ** precision)

# operation definitions
# all operations return a dictionary with the following fields:
# success: True/False
# message: string
# (optional) result: dictionary

# create a new wallet
# input: network, passphrase, word count (12, 15, 18, 21, 24), derivation path
# output: private key, address, mnemonic
def newwallet(network: str, passphrase: str, wordcount: int, dpath: str):
  res = {'success': False, 'message': '', 'result': {}}
  if wordcount not in [12, 15, 18, 21, 24]:
    res['message'] = 'Invalid word count for mnemonic generation'
    return res
  client = tronclient(network)
  wallet, mnemo = client.generate_address_with_mnemonic(passphrase = passphrase,
    num_words = wordcount, language = 'english', account_path = dpath)
  wallet['mnemonic'] = mnemo
  res['success'] = True
  res['result'] = wallet
  res['message'] = 'Wallet created, save this info in a secure place!\n\n'
  res['message'] += 'Address: %s\n' % wallet['base58check_address']
  res['message'] += 'Address (hex): %s\n' % wallet['hex_address']
  res['message'] += 'Private key (hex): %s\n' % wallet['private_key']
  res['message'] += 'Public key (hex): %s\n' % wallet['public_key']
  res['message'] += 'Mnemonic: %s' % wallet['mnemonic']
  return res

# import a wallet
# input: network, mnemonic, passphrase, derivation path
# output: private key, address
def importwallet(network: str, mnemo: str, passphrase: str, dpath: str):
  res = {'success': False, 'message': '', 'result': {}}
  client = tronclient(network)
  try:
    wallet = client.generate_address_from_mnemonic(mnemonic = mnemo,
      passphrase = passphrase, account_path = dpath)
  except (mnemonic.mnemonic.ConfigurationError,
          eth_utils.exceptions.ValidationError):
    res['message'] = 'Invalid mnemonic phrase!'
    return res
  wallet['mnemonic'] = mnemo
  res['success'] = True
  res['result'] = wallet
  res['message'] = 'Wallet imported, save this info in a secure place!\n\n'
  res['message'] += 'Address: %s\n' % wallet['base58check_address']
  res['message'] += 'Address (hex): %s\n' % wallet['hex_address']
  res['message'] += 'Private key (hex): %s\n' % wallet['private_key']
  res['message'] += 'Public key (hex): %s\n' % wallet['public_key']
  res['message'] += 'Mnemonic: %s' % wallet['mnemonic']
  return res

# get account information by Tron address
# input: network, address
# output: info(dictionary)
def infobyaddr(network: str, address: str):
  res = {'success': False, 'message': '', 'result': {}}
  client = tronclient(network)
  try:
    rawinfo = client.get_account(address)
  except BadAddress:
    res['message'] = 'Invalid address! Set the correct one with -addr flag'
    return res
  except AddressNotFound:
    res['message'] = 'Address not found on the blockchain! Set the correct one with -addr flag'
    return res
  # get bandwidth and energy information
  rawresinfo = client.get_account_resource(address)
  # some values used in bandwidth calculation
  freenetlimit = 0
  freenetused = 0
  netlimit = 0
  netused = 0
  if 'freeNetLimit' in rawresinfo:
    freenetlimit = rawresinfo['freeNetLimit']
  if 'freeNetUsed' in rawresinfo:
    freenetlimit = rawresinfo['freeNetUsed']
  if 'NetLimit' in rawresinfo:
    netlimit = rawresinfo['NetLimit']
  if 'NetUsed' in rawresinfo:
    netused = rawresinfo['NetUsed']
  bandwidth = freenetlimit - freenetused + netlimit - netused
  # some values used in energy calculation
  energylimit = 0
  energyused = 0
  if 'EnergyLimit' in rawresinfo:
    energylimit = rawresinfo['EnergyLimit']
  if 'EnergyUsed' in rawresinfo:
    energyused = rawresinfo['EnergyUsed']
  # info dictionary received, parse it
  info = {
    'address': rawinfo['address'],     # wallet address
    'trxbalance': rawinfo['balance'],  # TRX wallet balance (in Sun)
    'frozen': [],                      # hold frozen TRX info here
    'trc20balance': {},                # hold TRC20 token balances here
    'bandwidth': bandwidth,            # bandwidth
    'energy': energylimit - energyused # energy
  }
  # iterate over TRX freeze records
  for el in rawinfo['frozenV2']:
    if 'type' in el and 'amount' in el and el['amount'] > 0:
      info['frozen'].append({'type': el['type'], 'amount': el['amount']})
  # iterate over hardcoded TRC20 tokens
  for tkn in KT_TRC20_TOKENS[network]:
    tknobj = KT_TRC20_TOKENS[network][tkn]
    if tknobj['show_balance'] == True:
      info['trc20balance'][tknobj['ticker']] = trc20balance(client,
        tknobj['contract'], address)
  res['success'] = True
  res['result'] = info
  res['message'] = 'Address: %s\nBalance:\n' % info['address']
  res['message'] += '%f TRX\n' % (info['trxbalance'] / TRX_SCALE)
  for ticker in info['trc20balance']:
    res['message'] += '%f %s\n' % (info['trc20balance'][ticker], ticker)
  res['message'] += 'Frozen assets:\n'
  for el in info['frozen']:
    res['message'] += '%s\t%f TRX\n' % (el['type'], el['amount'] / TRX_SCALE)
  res['message'] += 'Remaining bandwidth: %s\n' % str(info['bandwidth'])
  res['message'] += 'Remaining energy: %s' % str(info['energy'])
  return res

# get own account information by private key
# input: network, private key
# output: info(dictionary)
def infobypk(network: str, pk: str):
  res = {'success': False, 'message': '', 'result': {}}
  try:
    rpk = PrivateKey(bytes.fromhex(pk))
  except (BadKey, ValueError):
    res['message'] = 'Invalid private key!'
    return res
  client = tronclient(network)
  addr = client.generate_address(rpk)['base58check_address']
  return infobyaddr(network, addr)

# send TRX transaction (and wait for its completion)
# input: network, private key, target address, amount (in TRX)
# output: transaction receipt object
def sendtrx(network: str, pk: str, toaddr: str, amt: float,
            offline: bool = False):
  res = {'success': False, 'message': '', 'result': {}}
  if amt < 1:
    res['message'] = 'Nothing to transfer!'
    return res
  try:
    rpk = PrivateKey(bytes.fromhex(pk))
  except (BadKey, ValueError):
    res['message'] = 'Invalid private key!'
    return res
  client = tronclient(network)
  fromaddr = client.generate_address(rpk)['base58check_address']
  # build and sign the transaction
  txn = client.trx.transfer(fromaddr, toaddr,
          int(amt * TRX_SCALE)).build().sign(rpk)
  if offline == True: # generate transaction JSON
    res['success'] = True
    res['result'] = txn
    res['message'] = json.dumps(txn.to_json(), separators=(',', ':'))
  else: # send the transaction
    try:
      receipt = txn.broadcast().wait()
    except ValidationError as e:
      res['message'] = str(e)
      return res
    res['success'] = True
    res['message'] = 'Sent %f TRX to %s\nTransaction ID %s' % (amt, 
                        toaddr, receipt['id'])
    res['result'] = receipt
  return res

# send TRC20 token transaction (and wait for its completion)
# input: network, private key, token ticker, target address, amount (in tokens)
# output: transaction receipt object
def sendtrc20(network: str, pk: str, ticker: str, toaddr: str,
              amt: float, offline: bool = False):
  res = {'success': False, 'message': '', 'result': {}}
  ticker = ticker.lower()
  if ticker not in KT_TRC20_TOKENS[network]:
    res['message'] = 'Token not supported by kisstron for %s!' % network
    return res
  if amt < 1:
    res['message'] = 'Nothing to transfer!'
    return res
  try:
    rpk = PrivateKey(bytes.fromhex(pk))
  except (BadKey, ValueError):
    res['message'] = 'Invalid private key!'
    return res
  client = tronclient(network)
  fromaddr = client.generate_address(rpk)['base58check_address']
  # instantiate the contract
  contractaddr = KT_TRC20_TOKENS[network][ticker]['contract']
  cntr = client.get_contract(contractaddr)
  precision = cntr.functions.decimals()
  scaler = 10 ** precision
  # build and sign the transaction
  txn = cntr.functions.transfer(toaddr,
          int(amt * scaler)).with_owner(fromaddr).build().sign(rpk)
  if offline == True: # generate transaction JSON
    res['success'] = True
    res['result'] = txn
    res['message'] = json.dumps(txn.to_json(), separators=(',', ':'))
  else: # send the transaction
    try:
      receipt = txn.broadcast().wait()
    except ValidationError as e:
      res['message'] = str(e)
      return res
    res['result'] = receipt
    if receipt['receipt']['result'] == 'OUT_OF_ENERGY':
      res['message'] = 'Not enough energy for transaction!'
      res['message'] += '\nTransaction ID %s' % receipt['id']
      return res
    res['success'] = True
    res['message'] = 'Sent %f %s to %s\nTransaction ID %s' % (amt, 
                        KT_TRC20_TOKENS[network][ticker]['ticker'],
                        toaddr, receipt['id'])
  return res

# freeze TRX for energy (and wait for its completion)
# input: network, private key, amount
# output: transaction receipt object
def freezetrx(network: str, pk: str, amt: float, offline: bool = False):
  res = {'success': False, 'message': '', 'result': {}}
  if amt < 1:
    res['message'] = 'Nothing to freeze!'
    return res
  try:
    rpk = PrivateKey(bytes.fromhex(pk))
  except (BadKey, ValueError):
    res['message'] = 'Invalid private key!'
    return res
  client = tronclient(network)
  myaddr = client.generate_address(rpk)['base58check_address']
  # build and sign the transaction
  txn = client.trx.freeze_balance(myaddr, int(amt * TRX_SCALE),
          "ENERGY").build().sign(rpk)
  if offline == True: # generate transaction JSON
    res['success'] = True
    res['result'] = txn
    res['message'] = json.dumps(txn.to_json(), separators=(',', ':'))
  else: # send the transaction
    try:
      receipt = txn.broadcast().wait()
    except ValidationError as e:
      res['message'] = str(e)
      return res
    res['success'] = True
    res['message'] = 'Frozen %f TRX at %s\nTransaction ID %s' % (amt, 
                        myaddr, receipt['id'])
    res['result'] = receipt
  return res

# unfreeze TRX for energy (and wait for its completion)
# input: network, private key, amount
# output: transaction receipt object
def unfreezetrx(network: str, pk: str, amt: float, offline: bool = False):
  res = {'success': False, 'message': '', 'result': {}}
  if amt < 1:
    res['message'] = 'Nothing to unfreeze!'
    return res
  try:
    rpk = PrivateKey(bytes.fromhex(pk))
  except (BadKey, ValueError):
    res['message'] = 'Invalid private key!'
    return res
  client = tronclient(network)
  myaddr = client.generate_address(rpk)['base58check_address']
  # build and sign the transaction
  txn = client.trx.unfreeze_balance(owner = myaddr,
          unfreeze_balance = int(amt * TRX_SCALE),
          resource = "ENERGY").build().sign(rpk)
  if offline == True: # generate transaction JSON
    res['success'] = True
    res['result'] = txn
    res['message'] = json.dumps(txn.to_json(), separators=(',', ':'))
  else: # send the transaction
    try:
      receipt = txn.broadcast().wait()
    except ValidationError as e:
      res['message'] = str(e)
      return res
    res['success'] = True
    res['message'] = 'Unfrozen %f TRX at %s\nTransaction ID %s' % (amt, 
                        myaddr, receipt['id'])
    res['result'] = receipt
  return res

# Broadcast transaction JSON saved in kisstron offline mode
# input: network, transaction JSON string
# output: transaction receipt object
def broadcastjson(network, tjson):
  res = {'success': False, 'message': '', 'result': {}}
  # import the transaction from JSON
  try:
    txn = Transaction.from_json(json.loads(tjson))
  except (TypeError, ValidationError, TransactionError,
          JSONDecodeError, KeyError):
    res['message'] = 'Malformed transaction JSON!'
    return res
  # instantiate a client for the transaction
  txn._client = tronclient(network)
  # send the transaction
  try:
    receipt = txn.broadcast().wait()
  except ValidationError as e:
    res['message'] = str(e)
    return res
  except TransactionError as e:
    res['message'] = str(e)
    return res
  res['result'] = receipt
  if ('receipt' in receipt and 'result' in receipt['receipt'] 
        and receipt['receipt']['result'] == 'OUT_OF_ENERGY'):
    res['message'] = 'Not enough energy for transaction!'
    res['message'] += '\nTransaction ID %s' % receipt['id']
    return res
  res['success'] = True
  res['message'] = 'Transaction sent\nTransaction ID %s' % receipt['id']
  return res

if __name__ == '__main__': # main app start
  from argparse import ArgumentParser
  parser = ArgumentParser(description='kisstron: a simple CLI wallet for Tron blockchain', epilog='(c) Luxferre 2024 --- No rights reserved <https://unlicense.org>')
  parser.add_argument('op', help='Operation: new, import, info, send, fre[eze], unf[reeze], bro[adcast]')
  parser.add_argument('coin', help='Coin to operate on: trx (default), usdt, usdc', default='trx')
  parser.add_argument('-net', '--network', default='mainnet', help='(op: all) Network: mainnet (default), nile, shasta')
  parser.add_argument('-node', '--tron-node', default=NODE_ADDR, help='(op: all) Custom Tron node HTTP address (default %s)' % NODE_ADDR)
  parser.add_argument('-tgkey', '--trongrid-api-key', default=None, help='(op: all) TronGrid API key (default None)')
  parser.add_argument('-fee', '--fee-limit', type=float, default=100.0, help='(op: send, fre, unf) Custom fee limit for mainnet transactions, in TRX (default 100)')
  parser.add_argument('-addr', '--address', default=None, help='(op: info, send) Tron address to send to or display information about (-pk overrides this to display own info)')
  parser.add_argument('-amt', '--amount', type=float, default=0, help='(op: send, fre, unf) Amount to send/freeze/unfreeze')
  parser.add_argument('-pk', '--private-key', default=None, help='(op: info, send, fre, unf) Private key for transactions or own information, in hex')
  parser.add_argument('-o', '--offline', action='store_true', help='(op: send, fre, unf) Build and sign the transaction in offline mode and output its JSON')
  parser.add_argument('-pp', '--passphrase', default='', help='(op: new, import) Mnemonic encryption passphrase (default empty)')
  parser.add_argument('-wc', '--words-count', type=int, default=12, help='(op: new) Word count in the mnemonic: 12 (default), 15, 18, 21, 24')
  parser.add_argument('-dp', '--derivation-path', default="m/44'/195'/0'/0/0", help="(op: new, import) BIP-39 derivation path for mnemonic (default is m/44'/195'/0'/0/0)")
  parser.add_argument('-nc', '--no-confirm', action='store_true', help='(op: send, fre, unf) Do not ask to type "yes" for transaction confirmation (dangerous)')

  args = parser.parse_args()

  # fix all relevant parameters (and lowercase where applicable)
  op = args.op.lower()            # operation
  coin = args.coin.lower()        # coin/token
  tronnet = args.network.lower()  # Tron network (mainnet, nile, shasta)
  targetaddr = args.address       # address parameter
  amount = float(args.amount)     # amount to send/freeze/unfreeze
  pk = args.private_key           # private key (in hex)
  passphrase = args.passphrase    # mnemonic encryption passphrase
  wcount = int(args.words_count)  # mnemonic word count (for new operation)
  dpath = args.derivation_path    # mnemonic derivation path (for new/import)
  trans_confirm = True            # ask for transaction confirmation ("yes") 
  NODE_ADDR = args.tron_node      # custom Tron node HTTP address
  TG_API_KEY = args.trongrid_api_key # optional TronGrid API key
  KT_FEE_LIMIT = int(args.fee_limit * TRX_SCALE) # custom fee limit
  offline_mode = False            # offline mode
  if args.offline == True:
    offline_mode = True
  if args.no_confirm == True:
    trans_confirm = False

  # check the constraints for every action and call appropriate functions

  opres = {} # operation result placeholder
  if op == 'new': # new wallet
    opres = newwallet(tronnet, passphrase, wcount, dpath)
  elif op == 'import': # import wallet
    mnemo = input('Your mnemonic phrase: ')
    opres = importwallet(tronnet, mnemo, passphrase, dpath)
  elif op == 'info': # info on own or foreign address
    if offline_mode == True:
      print('Info operation is not available in offline mode!')
      sys.exit(1)
    if pk == None: # public info, address must be set
      opres = infobyaddr(tronnet, targetaddr)  
    else: # own info, private key must be set
      opres = infobypk(tronnet, pk)
  elif op == 'send': # send transaction
    if trans_confirm == True:
      trconfirm('Enter "yes" to confirm sending funds: ')
    if coin == 'trx': # send TRX
      opres = sendtrx(tronnet, pk, targetaddr, amount, offline_mode)
    else: # send one of the predefined TRC20 tokens
      opres = sendtrc20(tronnet, pk, coin, targetaddr, amount, offline_mode)
  elif op == 'fre' or op == 'freeze': # freeze TRX for energy
    if trans_confirm == True:
      trconfirm('Enter "yes" to confirm freezing funds: ')
    opres = freezetrx(tronnet, pk, amount, offline_mode)
  elif op == 'unf' or op == 'unfreeze': # unfreeze TRX
    if trans_confirm == True:
      trconfirm('Enter "yes" to confirm unfreezing funds: ')
    opres = unfreezetrx(tronnet, pk, amount, offline_mode)
  elif op == 'bro' or op == 'broadcast': # broadcast signed JSON transaction
    if offline_mode == True:
      print('Broadcast operation is not available in offline mode!')
      sys.exit(1)
    transjson = input('Input transaction JSON: ')
    if trans_confirm == True:
      trconfirm('Enter "yes" to confirm sending the transaction: ')
    opres = broadcastjson(tronnet, transjson)
  else: # unknown operation
    print('Unknown operation!')
    print('Run the script with -h parameter to see available options')
    sys.exit(1)

  presult(opres) # print out the result