arbitrage/web3_py_client.py

import os
import json
import logging
from decimal import Decimal, ROUND_DOWN

from web3 import Web3
from web3.middleware import ExtraDataToPOAMiddleware # For PoA networks like Goerli, Sepolia, BSC etc.
from eth_account import Account
from dotenv import load_dotenv


# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 加载环境变量
load_dotenv()

# 标准 IERC20 ABI (只包含常用函数)
IERC20_ABI = json.loads('''
[
    {
        "constant": true,
        "inputs": [],
        "name": "name",
        "outputs": [{"name": "", "type": "string"}],
        "payable": false,
        "stateMutability": "view",
        "type": "function"
    },
    {
        "constant": false,
        "inputs": [
            {"name": "_spender", "type": "address"},
            {"name": "_value", "type": "uint256"}
        ],
        "name": "approve",
        "outputs": [{"name": "", "type": "bool"}],
        "payable": false,
        "stateMutability": "nonpayable",
        "type": "function"
    },
    {
        "constant": true,
        "inputs": [],
        "name": "totalSupply",
        "outputs": [{"name": "", "type": "uint256"}],
        "payable": false,
        "stateMutability": "view",
        "type": "function"
    },
    {
        "constant": false,
        "inputs": [
            {"name": "_from", "type": "address"},
            {"name": "_to", "type": "address"},
            {"name": "_value", "type": "uint256"}
        ],
        "name": "transferFrom",
        "outputs": [{"name": "", "type": "bool"}],
        "payable": false,
        "stateMutability": "nonpayable",
        "type": "function"
    },
    {
        "constant": true,
        "inputs": [],
        "name": "decimals",
        "outputs": [{"name": "", "type": "uint8"}],
        "payable": false,
        "stateMutability": "view",
        "type": "function"
    },
    {
        "constant": true,
        "inputs": [{"name": "_owner", "type": "address"}],
        "name": "balanceOf",
        "outputs": [{"name": "balance", "type": "uint256"}],
        "payable": false,
        "stateMutability": "view",
        "type": "function"
    },
    {
        "constant": true,
        "inputs": [],
        "name": "symbol",
        "outputs": [{"name": "", "type": "string"}],
        "payable": false,
        "stateMutability": "view",
        "type": "function"
    },
    {
        "constant": false,
        "inputs": [
            {"name": "_to", "type": "address"},
            {"name": "_value", "type": "uint256"}
        ],
        "name": "transfer",
        "outputs": [{"name": "", "type": "bool"}],
        "payable": false,
        "stateMutability": "nonpayable",
        "type": "function"
    },
    {
        "constant": true,
        "inputs": [
            {"name": "_owner", "type": "address"},
            {"name": "_spender", "type": "address"}
        ],
        "name": "allowance",
        "outputs": [{"name": "", "type": "uint256"}],
        "payable": false,
        "stateMutability": "view",
        "type": "function"
    }
]
''')

class EthClient:
    def __init__(self, rpc_url: str = None, private_key: str = None):
        self.rpc_url = rpc_url or os.getenv("RPC_URL")
        _private_key = private_key or os.getenv("PRIVATE_KEY")

        if not self.rpc_url:
            raise ValueError("RPC_URL not provided or found in environment variables.")
        if not _private_key:
            raise ValueError("PRIVATE_KEY not provided or found in environment variables.")

        self.w3 = Web3(Web3.HTTPProvider(self.rpc_url))

        # 如果连接的是 PoA 网络 (如 Goerli, Sepolia, BSC, Polygon), 需要注入中间件
        # 对于主网，不需要此操作。可以根据 chain_id 动态判断，或者让用户明确。
        # 例如：if self.w3.eth.chain_id in [5, 11155111, 56, 137]: # Goerli, Sepolia, BSC, Polygon
        self.w3.middleware_onion.inject(ExtraDataToPOAMiddleware, layer=0)

        if not self.w3.is_connected():
            raise ConnectionError(f"Failed to connect to Ethereum node at {self.rpc_url}")

        self.account = Account.from_key(_private_key)
        self.address = self.account.address
        logging.info(f"EthClient initialized. Address: {self.address}, RPC: {self.rpc_url}, Connected: {self.w3.is_connected()}")

    def _get_nonce(self) -> int:
        """获取账户的下一个 nonce"""
        return self.w3.eth.get_transaction_count(self.address)

    def _estimate_gas(self, tx: dict) -> int:
        """估算交易的 gas limit"""
        return self.w3.eth.estimate_gas(tx)

    def _sign_and_send_transaction(self, tx: dict, gas_limit_multiplier: float = 1.2) -> str:
        """签署并发送交易，返回交易哈希"""
        try:
            # 填充 gas 和 nonce (如果未提供)
            if 'nonce' not in tx:
                logging.info('TODO nonce应该提前管理好')
                tx['nonce'] = self._get_nonce()
            
            '''
            在使用 web3.py 手动构建和签名交易时，需要使用支持 EIP-1559 交易类型的签名函数。
            通常，eth_account 库及其 sign_transaction 方法是支持的，只要您提供的交易字典包含了正确的 EIP-1559 字段 (chainId, nonce, to, value, gas, maxFeePerGas, maxPriorityFeePerGas 等)。
            您构建的交易字典可能混合了传统 Gas 字段 (gasPrice) 和 EIP-1559 字段 (maxFeePerGas, maxPriorityFeePerGas)。一个交易只能使用其中一种方式来指定 Gas 费用。
            解决方案： 确保您的交易字典中只有 EIP-1559 相关的 Gas 字段（maxFeePerGas, maxPriorityFeePerGas 和 gas），或者只有传统 Gas 字段（gasPrice 和 gas）。
            不要同时包含 gasPrice 和 maxFeePerGas/maxPriorityFeePerGas。
            '''
            # # 对于支持 EIP-1559 的网络，应该使用:
            # if 'maxPriorityFeePerGas' not in tx and 'maxFeePerGas' not in tx:
            #     latest_block = self.w3.eth.get_block('latest')
            #     tx['maxPriorityFeePerGas'] = '1000000000'
            #     tx['maxFeePerGas'] = int(latest_block['baseFeePerGas']) * 2 + int(tx['maxPriorityFeePerGas'])

            if 'chainId' not in tx:
                tx['chainId'] = self.w3.eth.chain_id

            # 增加成交成功率
            tx['gas'] = int(int(tx['gas']) * gas_limit_multiplier)

            signed_tx = self.w3.eth.account.sign_transaction(tx, self.account.key)
            tx_hash = self.w3.eth.send_raw_transaction(signed_tx.raw_transaction)
            return self.w3.to_hex(tx_hash)
        except Exception as e:
            logging.info(f"Error signing or sending transaction: {e}")
            # 可以进一步处理特定错误，例如 nonce 过低，余额不足等
            raise

    def wait_for_transaction_receipt(self, tx_hash: str, timeout: int = 120, poll_latency: int = 1):
        """等待交易被打包并返回收据"""
        try:
            receipt = self.w3.eth.wait_for_transaction_receipt(tx_hash, timeout=timeout, poll_latency=poll_latency)
            return receipt
        except Exception as e: # Web3.exceptions.TimeExhausted as e
            logging.info(f"Transaction {tx_hash} timed out after {timeout} seconds.")
            raise

    def send_eth(self, to_address: str, amount_ether: float, gas_limit: int = None, gas_price_gwei: float = None) -> str:
        """发送 ETH"""
        if not self.w3.is_address(to_address):
            raise ValueError(f"Invalid recipient address: {to_address}")

        value_wei = self.w3.to_wei(amount_ether, 'ether')
        tx = {
            'to': self.w3.to_checksum_address(to_address),
            'value': value_wei,
            'from': self.address # web3.py 会自动从签名账户中获取 from，但显式写上更清晰
        }
        if gas_limit:
            tx['gas'] = gas_limit
        if gas_price_gwei:
            tx['gasPrice'] = self.w3.to_wei(gas_price_gwei, 'gwei')

        logging.info(f"Preparing to send {amount_ether} ETH to {to_address}...")
        return self._sign_and_send_transaction(tx)

    def _get_erc20_contract(self, token_address: str):
        """获取 ERC20 合约实例"""
        if not self.w3.is_address(token_address):
            raise ValueError(f"Invalid token address: {token_address}")
        return self.w3.eth.contract(address=self.w3.to_checksum_address(token_address), abi=IERC20_ABI)

    def get_erc20_decimals(self, token_address: str) -> int:
        """获取 ERC20 代币的精度"""
        contract = self._get_erc20_contract(token_address)
        return contract.functions.decimals().call()

    def _to_token_units(self, token_address: str, amount_readable: float) -> int:
        """将可读的代币数量转换为最小单位 (例如 1.0 USDT -> 1000000 if decimals is 6)"""
        decimals = self.get_erc20_decimals(token_address)
        factor = Decimal(10) ** Decimal(decimals)
        return int(Decimal(str(amount_readable)) * factor)

    def _from_token_units(self, token_address: str, amount_units: int) -> Decimal:
        """将最小单位的代币数量转换为可读数量"""
        decimals = self.get_erc20_decimals(token_address)
        factor = Decimal(10) ** Decimal(decimals)
        return (Decimal(amount_units) / factor).quantize(Decimal('0.1') ** decimals, rounding=ROUND_DOWN)

    def transfer_erc20(self, token_address: str, to_address: str, amount_readable: float,
                         gas_limit: int = None, gas_price: float = None) -> str:
        """
        转移 ERC20 代币。
        :param token_address: 代币合约地址
        :param to_address: 接收者地址
        :param amount_readable: 要转移的代币数量 (例如 1.2345)
        :param gas_limit: 手动设置 gas limit (可选)
        :param gas_price: 手动设置 gas price (可选)
        :return: 交易哈希
        """
        if not self.w3.is_address(to_address):
            raise ValueError(f"Invalid recipient address: {to_address}")

        contract = self._get_erc20_contract(token_address)
        amount_in_smallest_units = self._to_token_units(token_address, amount_readable)

        tx_data = contract.functions.transfer(
            self.w3.to_checksum_address(to_address),
            amount_in_smallest_units
        ).build_transaction({
            'from': self.address,
            'nonce': self._get_nonce(),
            'gas': 200000, # 通常 transfer ERC20 消耗 50k-100k gas，可以估算或手动设置
            'gasPrice': self.w3.eth.gas_price # 或手动设置
        })

        if gas_limit:
            tx_data['gas'] = gas_limit
        if gas_price:
            tx_data['gasPrice'] = gas_price

        logging.info(f"Preparing to transfer {amount_readable} of token {token_address} to {to_address}...")
        return self._sign_and_send_transaction(tx_data, 1.2, 2)

    def approve_erc20(self, token_address: str, spender_address: str, amount_readable: float,
                        gas_limit: int = None, gas_price_gwei: float = None) -> str:
        """
        授权给 spender 地址一定数量的 ERC20 代币。
        :param token_address: 代币合约地址
        :param spender_address: 被授权者地址
        :param amount_readable: 要授权的代币数量 (例如 1.2345)
        :param gas_limit: 手动设置 gas limit (可选)
        :param gas_price_gwei: 手动设置 gas price Gwei (可选)
        :return: 交易哈希
        """
        if not self.w3.is_address(spender_address):
            raise ValueError(f"Invalid spender address: {spender_address}")

        contract = self._get_erc20_contract(token_address)
        amount_in_smallest_units = self._to_token_units(token_address, amount_readable)

        tx_data = contract.functions.approve(
            self.w3.to_checksum_address(spender_address),
            amount_in_smallest_units
        ).build_transaction({
            'from': self.address,
            'nonce': self._get_nonce(),
        })

        if gas_limit:
            tx_data['gas'] = gas_limit
        if gas_price_gwei:
            tx_data['gasPrice'] = self.w3.to_wei(gas_price_gwei, 'gwei')

        logging.info(f"Preparing to approve {amount_readable} of token {token_address} for spender {spender_address}...")
        return self._sign_and_send_transaction(tx_data)

    def get_erc20_balance(self, token_address: str, owner_address: str = None) -> Decimal:
        """获取指定地址的 ERC20 代币余额 (可读数量)"""
        target_address = owner_address or self.address
        if not self.w3.is_address(target_address):
            raise ValueError(f"Invalid owner address: {target_address}")

        contract = self._get_erc20_contract(token_address)
        balance_units = contract.functions.balanceOf(self.w3.to_checksum_address(target_address)).call()
        return self._from_token_units(token_address, balance_units)

    def get_erc20_allowance(self, token_address: str, spender_address: str, owner_address: str = None) -> Decimal:
        """获取 owner 授权给 spender 的 ERC20 代币数量 (可读数量)"""
        target_owner = owner_address or self.address
        if not self.w3.is_address(target_owner):
            raise ValueError(f"Invalid owner address: {target_owner}")
        if not self.w3.is_address(spender_address):
            raise ValueError(f"Invalid spender address: {spender_address}")

        contract = self._get_erc20_contract(token_address)
        allowance_units = contract.functions.allowance(
            self.w3.to_checksum_address(target_owner),
            self.w3.to_checksum_address(spender_address)
        ).call()
        return self._from_token_units(token_address, allowance_units)

    def get_erc20_total_supply(self, token_address: str) -> Decimal:
        """获取 ERC20 代币的总供应量 (可读数量)"""
        contract = self._get_erc20_contract(token_address)
        total_supply_units = contract.functions.totalSupply().call()
        return self._from_token_units(token_address, total_supply_units)

    def get_erc20_name(self, token_address: str) -> str:
        """获取 ERC20 代币的名称"""
        contract = self._get_erc20_contract(token_address)
        return contract.functions.name().call()

    def get_erc20_symbol(self, token_address: str) -> str:
        """获取 ERC20 代币的符号"""
        contract = self._get_erc20_contract(token_address)
        return contract.functions.symbol().call()

    def get_eth_balance(self, address: str = None) -> Decimal:
        """获取ETH余额 (单位 Ether)"""
        target_address = address or self.address
        if not self.w3.is_address(target_address):
            raise ValueError(f"Invalid address: {target_address}")
        balance_wei = self.w3.eth.get_balance(self.w3.to_checksum_address(target_address))
        return self.w3.from_wei(balance_wei, 'ether')

if __name__ == "__main__":
    from ok_chain_client import swap
    import decimal
    import pprint

    client = EthClient()

    # CHAIN_ID = 1
    # IN_AMOUNT_TO_QUERY = decimal.Decimal('1')
    # IN_TOKEN_ADDRESS = '0xdAC17F958D2ee523a2206206994597C13D831ec7'     # USDT on Ethereum
    # IN_TOKEN_DECIMALS = decimal.Decimal(6)
    # OUT_TOKEN_ADDRESS = '0xf816507E690f5Aa4E29d164885EB5fa7a5627860'    # RATO on Ethereum
    USER_WALLET = '0xb1f33026Db86a86372493a3B124d7123e9045Bb4'
    # SLIPPAGE = 1
    # USER_EXCHANGE_WALLET = '0xc71835a042F4d870B0F4296cc89cAeb921a9f3DA'

    # rst = swap(CHAIN_ID, IN_AMOUNT_TO_QUERY * (10 ** IN_TOKEN_DECIMALS), IN_TOKEN_ADDRESS, OUT_TOKEN_ADDRESS, SLIPPAGE, USER_WALLET, USER_EXCHANGE_WALLET)
    # data = rst['data'][0]
    # tx = data['tx']
    try:
        # tx.pop('gasPrice', None)
        # tx.pop('value', None)
        # tx.pop('minReceiveAmount', None)
        # tx.pop('slippage', None)
        # tx.pop('maxSpendAmount', None)
        # tx.pop('signatureData', None)

        tx = {
            'from': USER_WALLET, 
            'to': USER_WALLET, 
            'gas': '40000',
            'value': 1,
            'maxPriorityFeePerGas': '1800000000'
        }
        latest_block = client.w3.eth.get_block('latest')
        tx['maxPriorityFeePerGas'] = int(tx['maxPriorityFeePerGas'])
        tx['maxFeePerGas'] = int(int(latest_block['baseFeePerGas']) * 2 + tx['maxPriorityFeePerGas'])

        pprint.pprint(tx)
        estimated_gas = client.w3.eth.estimate_gas(tx)
        estimated_wei = estimated_gas * (tx['maxPriorityFeePerGas'] + tx['maxFeePerGas'])
        estimated_eth = estimated_wei / (10 ** 18)
        logging.info(f"估算的燃气量: {estimated_gas}, eth消耗: {estimated_eth}")

        logging.info(f"餘額：{client.w3.eth.get_balance(USER_WALLET)}")

        # tx_hash = client._sign_and_send_transaction(tx)
        # receipt = client.wait_for_transaction_receipt(tx_hash)
        # logging.info(f"{tx_hash} 交易已确认! Status: {'Success' if receipt.status == 1 else 'Failed'}")
    except Exception as e:
        print(f"Gas 估算失败: {e}")
    
    # # --- 使用示例 ---
    # # 确保你的 .env 文件配置正确
    # # 并且你的账户中有足够的 ETH 来支付 Gas 费

    # # 替换为实际的ERC20代币地址和接收者地址 (例如USDT on Sepolia testnet)
    # # Sepolia USDT: 0xaA8E23Fb1079EA71e0a56F48S1a3ET28wpD1RLf
    # # Sepolia WETH: 0x7b79995e5f793A07Bc00c21412e50Ecn10yt2e_ (错误，应为 0x7b79995e5f793A07Bc00c21412e50Ecn10yt2eH)
    # # 更正: Sepolia WETH: 0xfFf9976782d46CC05630D1f6eBAb18b2324d6B14 (常用)
    # # 有些测试网可能没有标准的USDT，你可以找一个存在的ERC20代币进行测试，或者自己部署一个
    # # 为了演示，这里假设使用的是 Sepolia 测试网

    # # !!重要!!: 以下地址和代币地址仅为示例, 请替换为您测试网络上的真实地址和代币
    # # 如果您在主网操作，请务必小心，并使用小额资金测试。
    # TEST_RECIPIENT_ADDRESS = "0xb1f33026db86a86372493a3b124d7123e9045bb4" # 替换为你的测试接收地址
    # # Sepolia 上的一个示例 ERC20 token (你可以找一个你有的测试币)
    # # 用于测试的代币地址
    # TEST_ERC20_TOKEN_ADDRESS_SEPOLIA_LINK = "0xdAC17F958D2ee523a2206206994597C13D831ec7"

    # try:
    #     client = EthClient() # RPC_URL 和 PRIVATE_KEY 会从 .env 文件加载

    #     logging.info(f"\nMy ETH Balance: {client.get_eth_balance()} ETH")

    #     # 1. 发送 ETH (取消注释以测试, 确保接收地址正确且你有足够ETH)
    #     # logging.info(f"\nAttempting to send ETH...")
    #     # eth_tx_hash = client.send_eth(TEST_RECIPIENT_ADDRESS, 0.0001) # 发送 0.0001 ETH
    #     # logging.info(f"ETH transaction sent! Hash: {eth_tx_hash}")
    #     # receipt = client.wait_for_transaction_receipt(eth_tx_hash)
    #     # logging.info(f"ETH transaction confirmed! Status: {'Success' if receipt.status == 1 else 'Failed'}")

    #     # --- ERC20 操作示例 ---
    #     # 使用 Sepolia LINK 代币进行演示
    #     token_address = TEST_ERC20_TOKEN_ADDRESS_SEPOLIA_LINK
    #     if not client.w3.is_address(token_address): # 简单检查
    #          logging.info(f"Warning: {token_address} does not look like a valid address. Skipping ERC20 tests.")
    #     else:
    #         logging.info(f"\n--- ERC20 Token Operations for: {token_address} ---")
    #         token_name = client.get_erc20_name(token_address)
    #         token_symbol = client.get_erc20_symbol(token_address)
    #         token_decimals = client.get_erc20_decimals(token_address)
    #         logging.info(f"Token: {token_name} ({token_symbol}), Decimals: {token_decimals}")

    #         # ERC20余额查询以及基础功能测试（总供应量）
    #         my_token_balance = client.get_erc20_balance(token_address)
    #         logging.info(f"My {token_symbol} Balance: {my_token_balance} {token_symbol}")
    #         total_supply = client.get_erc20_total_supply(token_address)
    #         logging.info(f"Total Supply of {token_symbol}: {total_supply} {token_symbol}")

    #         # # 2. ERC20 转账 (取消注释以测试, 确保你有该代币且接收地址正确)
    #         # amount_to_transfer = 0.01 # 转移 0.01 个代币
    #         # if my_token_balance >= Decimal(str(amount_to_transfer)):
    #         #     logging.info(f"\nAttempting to transfer {amount_to_transfer} {token_symbol}...")
    #         #     erc20_tx_hash = client.transfer_erc20(token_address, TEST_RECIPIENT_ADDRESS, amount_to_transfer)
    #         #     logging.info(f"{token_symbol} transfer transaction sent! Block: {client.w3.eth.block_number} Hash: {erc20_tx_hash}")
    #         #     receipt = client.wait_for_transaction_receipt(erc20_tx_hash)
    #         #     logging.info(f"{token_symbol} transfer transaction confirmed! Block: {client.w3.eth.block_number} Status: {'Success' if receipt.status == 1 else 'Failed'}")
    #         #     logging.info(f"My new {token_symbol} Balance: {client.get_erc20_balance(token_address)} {token_symbol}")
    #         # else:
    #         #     logging.info(f"Insufficient {token_symbol} balance to transfer {amount_to_transfer} {token_symbol}.")

    #         # # 3. ERC20 Approve 和 Allowance (取消注释以测试)
    #         # spender_for_allowance = '0x156ACd2bc5fC336D59BAAE602a2BD9b5e20D6672' # 可以是任何你想授权的地址
    #         # amount_to_approve = 74547271788
    #         # token_address = "0xdAC17F958D2ee523a2206206994597C13D831ec7"
            
    #         # # current_allowance = client.get_erc20_allowance(token_address, spender_for_allowance)
    #         # # logging.info(f"\nCurrent allowance for {spender_for_allowance} to spend my {token_symbol}: {current_allowance} {token_symbol}")
            
    #         # # if my_token_balance >= Decimal(str(amount_to_approve)): # 确保有足够的代币去授权（虽然授权本身不消耗代币）
    #         # logging.info(f"\nAttempting to approve {amount_to_approve} {token_symbol} for spender {spender_for_allowance}...")
    #         # approve_tx_hash = client.approve_erc20(token_address, spender_for_allowance, amount_to_approve)
    #         # logging.info(f"{token_symbol} approve transaction sent! Hash: {approve_tx_hash}")
    #         # receipt = client.wait_for_transaction_receipt(approve_tx_hash)
    #         # logging.info(f"{token_symbol} approve transaction confirmed! Status: {'Success' if receipt.status == 1 else 'Failed'}")
        
    #         # new_allowance = client.get_erc20_allowance(token_address, spender_for_allowance)
    #         # logging.info(f"New allowance for {spender_for_allowance}: {new_allowance} {token_symbol}")
    #         # else:
    #         #      logging.info(f"Not enough balance to consider approving {amount_to_approve} {token_symbol} (though approval itself doesn't spend tokens).")

    # except ValueError as ve:
    #     logging.info(f"Configuration Error: {ve}")
    # except ConnectionError as ce:
    #     logging.info(f"Connection Error: {ce}")
    # except Exception as e:
    #     logging.info(f"An unexpected error occurred: {e}")
    #     import traceback
    #     traceback.logging.info_exc()