wallet-system/internal/data/wallet.go

package data

import (
	"context"
	"encoding/json"
	"errors"
	"github.com/fbsobreira/gotron-sdk/pkg/address"
	"github.com/fbsobreira/gotron-sdk/pkg/proto/api"
	"github.com/go-kratos/kratos/v2/log"
	"math/big"
	"wallet-system/internal/biz"
	"wallet-system/internal/pkg/logging/applogger"
	"wallet-system/internal/pkg/wallet/tron"
	"wallet-system/internal/pkg/wallet/tron/sign"
)

// walletRepo
// @Description:
type walletRepo struct {
	data *Data

	log *log.Helper
}

// NewWalletRepo
//
//	@Description:
//	@param data
//	@param logger
//	@return biz.WalletRepo
func NewWalletRepo(data *Data, logger log.Logger) biz.WalletRepo {
	return &walletRepo{
		data: data,
		log:  log.NewHelper(logger),
	}
}

/*
GenerateAddress Private key encryption management
1、KMS encryption
2、Leveldb storage
3、S3 storage
*/
func (r *walletRepo) GenerateAddress(ctx context.Context, wallet string, number uint32) ([]string, error) {
	var addressList []string
	for i := 1; i <= int(number); i++ {
		walletAddress, privateKey, err := tron.TronWalletAddress()
		if err != nil {
			applogger.Error("Error generating wallet address and private key:%v", err)
			return []string{}, err
		}

		encryptByte, err := r.data.kmsC.Encrypt(r.data.aws.Set.Key, privateKey)
		if err != nil {
			applogger.Error("GenerateAddress.Encrypt err:%v", err)
			return nil, err
		}
		applogger.Debug("AWS-KMS encrypted data:%v", encryptByte)

		if err = r.data.s3C.Storage(encryptByte, r.data.aws.Set.Bucket, walletAddress); err != nil {
			applogger.Error("GenerateAddress.storage err:%v", err)
			return nil, err
		}

		if err = r.data.levelC.WriteDB(ctx, walletAddress, encryptByte); err != nil {
			applogger.Error("Error writing wallet cache address:%v", err)
			return []string{}, err
		}

		addressList = append(addressList, walletAddress)
	}

	return addressList, nil
}

/*
SignatureTrc20Grpc 钱包签名转账上链
1、Determine signature type
2、Obtain wallet private key
3、Decrypting private keys
4、signature
5、Upper chain
6、Return status and transaction hash
*/
func (r *walletRepo) SignatureTrc20Grpc(ctx context.Context, wallet *biz.Wallet) (bool, string, error) {
	// 获取钱包私钥地址
	private, err := r.data.levelC.ReadPrivateKeyByAddress(ctx, wallet.From)
	if err != nil {
		return false, "", err
	}

	// 解密私钥
	pky, err := r.data.kmsC.Decrypt(private)
	if err != nil {
		return false, "", err
	}
	privateKey := string(pky)

	// 发送者地址
	from, err := address.Base58ToAddress(wallet.From)
	if err != nil {
		return false, "", err
	}

	// 接受者地址
	to, err := address.Base58ToAddress(wallet.To)
	if err != nil {
		return false, "", err
	}

	var contract address.Address     // 转账合约地址
	var tx *api.TransactionExtention // 创建交易信息

	switch wallet.TrcCheck {
	case tron.Trx:
		tx, err = r.data.grpcT.Transfer(from.String(), to.String(), tron.AmountInt(int(wallet.Amount)))
		if err != nil {
			applogger.Error("SignatureTrc20Grpc.Transfer err:%v", err)
			return false, "", err
		}

		return r.SignTransactionAndBroadcastTransaction(ctx, tx, privateKey)
	case tron.Trc20:
		contract, err = address.Base58ToAddress(wallet.Contract)
		if err != nil {
			return false, "", err
		}

		tx, err = r.data.grpcT.TransferTrc20(from.String(), to.String(), contract.String(), tron.AmountBigInt(int(wallet.Amount)), tron.FeeLimitInt(int(wallet.FeeLimit)))
		if err != nil {
			applogger.Error("SignatureTrc20Grpc.TransferTrc20 err:%v", err)
			return false, "", err
		}

		return r.SignTransactionAndBroadcastTransaction(ctx, tx, privateKey)
	default:
		return false, "", errors.New("Please enter the correct signature type.")
	}
}

/*
WalletApprove
1、Determine signature type
2、Obtain wallet private key
3、Decrypting private keys
4、signature
5、Upper chain
6、Return status and transaction hash
*/
func (r *walletRepo) WalletApprove(ctx context.Context, wallet *biz.Wallet) (bool, []byte, error) {
	tx, err := r.data.grpcT.TRC20Approve(wallet.From, wallet.To, wallet.Contract, big.NewInt(int64(wallet.Amount)), tron.FeeLimitInt(int(wallet.FeeLimit)))
	if err != nil {
		applogger.Error("TRC20Approve err:%v", err)
		return false, []byte{}, err
	}

	jms, err := json.Marshal(tx)
	if err != nil {
		return false, []byte{}, err
	}

	return false, jms, nil
}

// GetAllAddress
//
//	@Description:
//	@receiver r
//	@param ctx
//	@return []string
//	@return error
func (r *walletRepo) GetAllAddress(ctx context.Context) ([]string, error) {
	listAddress, err := r.data.levelC.ReadWalletAddrList(ctx)
	if err != nil {
		return []string{}, err
	}
	return listAddress, nil
}

// GetAllAddressAndPrivateKey
//
//	@Description:
//	@receiver r
//	@param ctx
//	@return map[string][]byte
//	@return error
func (r *walletRepo) GetAllAddressAndPrivateKey(ctx context.Context) (map[string][]byte, error) {
	listAddress, err := r.data.levelC.ReadWalletAddressPrivateKeyList(ctx)
	if err != nil {
		return map[string][]byte{}, err
	}

	return listAddress, nil
}

// GetPrivateKeyByAddress
//
//	@Description:
//	@receiver r
//	@param ctx
//	@param address
//	@return []byte
//	@return error
func (r *walletRepo) GetPrivateKeyByAddress(ctx context.Context, address string) (string, error) {
	private, err := r.data.levelC.ReadPrivateKeyByAddress(ctx, address)
	if err != nil {
		return "", err
	}

	pky, err := r.data.kmsC.Decrypt(private)
	if err != nil {
		return "", err
	}
	privateKey := string(pky)

	return privateKey, nil
}

// DeleteAddress
//
//	@Description:
//	@receiver r
//	@param ctx
//	@return error
func (r *walletRepo) DeleteAddress(ctx context.Context) error {
	if err := r.data.levelC.DeleteAllByKey(ctx); err != nil {
		return err
	}

	return nil
}

// SignTransactionAndBroadcastTransaction
//
//	@Description:
//	@receiver r
//	@param ctx
//	@param tx
//	@param privateKey
//	@return bool
//	@return string
//	@return error
func (r *walletRepo) SignTransactionAndBroadcastTransaction(ctx context.Context, tx *api.TransactionExtention, privateKey string) (bool, string, error) {
	signTx, err := sign.SignTransaction(tx.Transaction, privateKey)
	if err != nil {
		applogger.Error("SignTransactionAndBroadcastTransaction err:%v", err)
		return false, "", err
	}

	if err := r.data.grpcT.BroadcastTransaction(signTx); err != nil {
		return false, "", err
	}

	return true, tron.Encode(tx.GetTxid()), nil
}