404 lines
15 KiB
Python
404 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import time
|
|
from dataclasses import dataclass
|
|
from decimal import Decimal, InvalidOperation
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from src.infrastructure.config import settings
|
|
from src.infrastructure.wallet_balances.token_registry import (
|
|
NATIVE_DECIMALS,
|
|
NATIVE_SYMBOLS,
|
|
get_coingecko_id,
|
|
get_evm_tokens,
|
|
get_sol_tokens,
|
|
get_trx_tokens,
|
|
)
|
|
|
|
TIMEOUT_SECONDS = 15
|
|
RETRY_ATTEMPTS = 2
|
|
PRICE_CACHE_SECONDS = 30
|
|
BLOCKSTREAM_URL = 'https://blockstream.info/api'
|
|
TRONGRID_URL = 'https://api.trongrid.io'
|
|
COINGECKO_URL = 'https://api.coingecko.com/api/v3/simple/price'
|
|
BALANCE_OF_SELECTOR = '70a08231'
|
|
BASE58_ALPHABET = '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz'
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class FormattedAmount:
|
|
raw: str
|
|
formatted: str
|
|
decimals: int
|
|
usd_price: Decimal | None = None
|
|
usd_value: Decimal | None = None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class WalletBalance:
|
|
wallet_id: str
|
|
chain: str
|
|
address: str
|
|
derivation_path: str
|
|
native_symbol: str
|
|
native: FormattedAmount
|
|
tokens: dict[str, FormattedAmount]
|
|
total_usd: Decimal | None
|
|
error: str | None = None
|
|
|
|
|
|
class WalletBalanceService:
|
|
def __init__(self) -> None:
|
|
self._prices_cache: dict[tuple[tuple[str,str],...],tuple[float,dict[str,Decimal | None]]] = {}
|
|
|
|
|
|
async def get_wallet_balances(self, wallets) -> list[WalletBalance]:
|
|
prices = await self._get_prices(wallets)
|
|
results = await asyncio.gather(
|
|
*(self._get_wallet_balance(wallet, prices) for wallet in wallets),
|
|
return_exceptions=True,
|
|
)
|
|
balances: list[WalletBalance] = []
|
|
for wallet, result in zip(wallets, results, strict=False):
|
|
if isinstance(result, Exception):
|
|
balances.append(self._failed_wallet(wallet, str(result)))
|
|
else:
|
|
balances.append(result)
|
|
return balances
|
|
|
|
async def _get_wallet_balance(self, wallet, prices: dict[str, Decimal | None]) -> WalletBalance:
|
|
chain = wallet.chain.upper()
|
|
if chain == 'BTC':
|
|
native_raw, token_raw = await self._btc_balance(wallet.address)
|
|
elif chain == 'ETH':
|
|
native_raw, token_raw = await self._evm_balance(
|
|
settings.ETH_RPC_URL,
|
|
wallet.address,
|
|
get_evm_tokens(chain),
|
|
)
|
|
elif chain == 'BSC':
|
|
native_raw, token_raw = await self._evm_balance(
|
|
settings.BSC_RPC_URL,
|
|
wallet.address,
|
|
get_evm_tokens(chain),
|
|
)
|
|
elif chain == 'TRX':
|
|
native_raw, token_raw = await self._trx_balance(wallet.address)
|
|
elif chain == 'SOL':
|
|
native_raw, token_raw = await self._sol_balance(wallet.address)
|
|
else:
|
|
raise ValueError(f'Unsupported chain: {wallet.chain}')
|
|
|
|
native = self._amount(native_raw, NATIVE_DECIMALS[chain], prices.get(f'{chain}:{NATIVE_SYMBOLS[chain]}'))
|
|
tokens = self._token_amounts(chain, token_raw, prices)
|
|
total_usd = self._total_usd(native, tokens)
|
|
return WalletBalance(
|
|
wallet_id=wallet.id,
|
|
chain=chain,
|
|
address=wallet.address,
|
|
derivation_path=wallet.derivation_path,
|
|
native_symbol=NATIVE_SYMBOLS[chain],
|
|
native=native,
|
|
tokens=tokens,
|
|
total_usd=total_usd,
|
|
)
|
|
|
|
async def _btc_balance(self, address: str) -> tuple[str, dict[str, str]]:
|
|
data = await self._get_json(f'{BLOCKSTREAM_URL}/address/{address}')
|
|
stats = data.get('chain_stats') or {}
|
|
funded = int(stats.get('funded_txo_sum') or 0)
|
|
spent = int(stats.get('spent_txo_sum') or 0)
|
|
return str(funded - spent), {}
|
|
|
|
async def _evm_balance(self, rpc_url: str, address: str, tokens: list[dict]) -> tuple[str, dict[str, str]]:
|
|
native_result = await self._rpc(rpc_url, 'eth_getBalance', [address, 'latest'])
|
|
native = self._hex_to_int_string(native_result)
|
|
token_items = await asyncio.gather(
|
|
*(self._evm_token_balance(rpc_url, address, token) for token in tokens),
|
|
return_exceptions=True,
|
|
)
|
|
result: dict[str, str] = {}
|
|
for token, item in zip(tokens, token_items, strict=False):
|
|
result[token['symbol']] = '0' if isinstance(item, Exception) else item
|
|
return native, result
|
|
|
|
async def _evm_token_balance(self, rpc_url: str, address: str, token: dict) -> str:
|
|
data = '0x' + BALANCE_OF_SELECTOR + address.lower().removeprefix('0x').rjust(64, '0')
|
|
result = await self._rpc(rpc_url, 'eth_call', [{'to': token['contract'], 'data': data}, 'latest'])
|
|
return self._hex_to_int_string(result)
|
|
|
|
async def _trx_balance(self, address: str) -> tuple[str, dict[str, str]]:
|
|
headers = self._tron_headers()
|
|
account = await self._get_json(f'{TRONGRID_URL}/v1/accounts/{address}', headers=headers)
|
|
trx = str(((account.get('data') or [{}])[0] or {}).get('balance') or 0)
|
|
address_hex = self._tron_address_to_hex(address).rjust(64, '0')
|
|
token_config = self._tokens_for_chain('TRX')
|
|
token_items = await asyncio.gather(
|
|
*(self._trx_token_balance(address, address_hex, token, headers) for token in token_config),
|
|
return_exceptions=True,
|
|
)
|
|
tokens: dict[str, str] = {}
|
|
for token, item in zip(token_config, token_items, strict=False):
|
|
tokens[token['symbol']] = '0' if isinstance(item, Exception) else item
|
|
return trx, tokens
|
|
|
|
async def _trx_token_balance(self, address: str, address_hex: str, token: dict, headers: dict[str, str]) -> str:
|
|
data = await self._post_json(
|
|
f'{TRONGRID_URL}/wallet/triggerconstantcontract',
|
|
{
|
|
'owner_address': address,
|
|
'contract_address': token['contract'],
|
|
'function_selector': 'balanceOf(address)',
|
|
'parameter': address_hex,
|
|
'visible': True,
|
|
},
|
|
headers={**headers, 'Content-Type': 'application/json'},
|
|
)
|
|
value = (data.get('constant_result') or [None])[0]
|
|
if not value or set(value) == {'0'}:
|
|
return '0'
|
|
return str(int(value, 16))
|
|
|
|
async def _sol_balance(self, address: str) -> tuple[str, dict[str, str]]:
|
|
native_data = await self._rpc(settings.SOL_RPC_URL, 'getBalance', [address])
|
|
native = str((native_data or {}).get('value') or 0)
|
|
tokens = {token['symbol']: '0' for token in get_sol_tokens()}
|
|
try:
|
|
accounts = await self._rpc(
|
|
settings.SOL_RPC_URL,
|
|
'getTokenAccountsByOwner',
|
|
[
|
|
address,
|
|
{'programId': 'TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA'},
|
|
{'encoding': 'jsonParsed'},
|
|
],
|
|
)
|
|
known = {token['mint']: token['symbol'] for token in get_sol_tokens()}
|
|
for item in accounts.get('value') or []:
|
|
info = (((item.get('account') or {}).get('data') or {}).get('parsed') or {}).get('info') or {}
|
|
mint = info.get('mint')
|
|
amount = ((info.get('tokenAmount') or {}).get('amount')) or '0'
|
|
if mint in known:
|
|
symbol = known[mint]
|
|
tokens[symbol] = str(int(tokens[symbol]) + int(amount))
|
|
except Exception:
|
|
pass
|
|
return native, tokens
|
|
|
|
async def _get_prices(self, wallets) -> dict[str, Decimal | None]:
|
|
pairs: dict[str, str] = {}
|
|
for wallet in wallets:
|
|
chain = wallet.chain.upper()
|
|
symbol = NATIVE_SYMBOLS.get(chain)
|
|
coin_id = get_coingecko_id(chain, symbol or chain)
|
|
if symbol and coin_id:
|
|
pairs[f'{chain}:{symbol}'] = coin_id
|
|
for token in self._tokens_for_chain(chain):
|
|
coin_id = get_coingecko_id(chain, token['symbol'])
|
|
if coin_id:
|
|
symbol = token['symbol']
|
|
pairs[f'{chain}:{symbol}'] = coin_id
|
|
if not pairs:
|
|
return {}
|
|
|
|
cache_key = tuple(sorted(pairs.items()))
|
|
cached = self._prices_cache.get(cache_key)
|
|
now = time.monotonic()
|
|
if cached is not None and now - cached[0] < PRICE_CACHE_SECONDS:
|
|
return dict(cached[1])
|
|
|
|
try:
|
|
response = await self._get_json(
|
|
COINGECKO_URL,
|
|
params={
|
|
'ids': ','.join(sorted(set(pairs.values()))),
|
|
'vs_currencies': 'usd',
|
|
},
|
|
headers=self._coingecko_headers(),
|
|
timeout=5,
|
|
)
|
|
except Exception:
|
|
return {key: None for key in pairs}
|
|
|
|
prices = {
|
|
key: self._decimal_or_none((response.get(coin_id) or {}).get('usd'))
|
|
for key, coin_id in pairs.items()
|
|
}
|
|
self._prices_cache[cache_key] = (now,prices)
|
|
return prices
|
|
|
|
def _token_amounts(
|
|
self,
|
|
chain: str,
|
|
raw: dict[str, str],
|
|
prices: dict[str, Decimal | None],
|
|
) -> dict[str, FormattedAmount]:
|
|
decimals = {token['symbol']: token['decimals'] for token in self._tokens_for_chain(chain)}
|
|
return {
|
|
symbol: self._amount(value, decimals.get(symbol, 0), prices.get(f'{chain}:{symbol}'))
|
|
for symbol, value in raw.items()
|
|
}
|
|
|
|
def _amount(self, raw: str, decimals: int, usd_price: Decimal | None) -> FormattedAmount:
|
|
formatted = self._format_units(raw, decimals)
|
|
usd_value = None
|
|
if usd_price is not None:
|
|
usd_value = self._round_usd(Decimal(formatted) * usd_price)
|
|
return FormattedAmount(
|
|
raw=raw,
|
|
formatted=formatted,
|
|
decimals=decimals,
|
|
usd_price=usd_price,
|
|
usd_value=usd_value,
|
|
)
|
|
|
|
def _total_usd(self, native: FormattedAmount, tokens: dict[str, FormattedAmount]) -> Decimal | None:
|
|
values = [native.usd_value, *(token.usd_value for token in tokens.values())]
|
|
present = [value for value in values if value is not None]
|
|
if not present:
|
|
return None
|
|
return self._round_usd(sum(present, Decimal('0')))
|
|
|
|
def _failed_wallet(self, wallet, error: str) -> WalletBalance:
|
|
chain = wallet.chain.upper()
|
|
native = FormattedAmount(
|
|
raw='0',
|
|
formatted='0',
|
|
decimals=NATIVE_DECIMALS.get(chain, 0),
|
|
)
|
|
return WalletBalance(
|
|
wallet_id=wallet.id,
|
|
chain=chain,
|
|
address=wallet.address,
|
|
derivation_path=wallet.derivation_path,
|
|
native_symbol=NATIVE_SYMBOLS.get(chain, chain),
|
|
native=native,
|
|
tokens={},
|
|
total_usd=None,
|
|
error=error,
|
|
)
|
|
|
|
async def _rpc(self, url: str, method: str, params: list[Any]) -> Any:
|
|
payload = {
|
|
'jsonrpc': '2.0',
|
|
'id': 1,
|
|
'method': method,
|
|
'params': params,
|
|
}
|
|
response = await self._post_json(url, payload)
|
|
if response.get('error'):
|
|
raise RuntimeError(str(response['error']))
|
|
return response.get('result')
|
|
|
|
async def _get_json(
|
|
self,
|
|
url: str,
|
|
*,
|
|
params: dict[str, str] | None = None,
|
|
headers: dict[str, str] | None = None,
|
|
timeout: int = TIMEOUT_SECONDS,
|
|
) -> Any:
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
for attempt in range(RETRY_ATTEMPTS + 1):
|
|
try:
|
|
response = await client.get(url, params=params, headers=headers)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
except (httpx.RequestError,httpx.HTTPStatusError,ValueError):
|
|
if attempt >= RETRY_ATTEMPTS:
|
|
raise
|
|
await asyncio.sleep(0.2 * (attempt + 1))
|
|
raise RuntimeError('HTTP GET failed')
|
|
|
|
async def _post_json(
|
|
self,
|
|
url: str,
|
|
payload: dict[str, Any],
|
|
*,
|
|
headers: dict[str, str] | None = None,
|
|
timeout: int = TIMEOUT_SECONDS,
|
|
) -> Any:
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
for attempt in range(RETRY_ATTEMPTS + 1):
|
|
try:
|
|
response = await client.post(url, json=payload, headers=headers)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
except (httpx.RequestError,httpx.HTTPStatusError,ValueError):
|
|
if attempt >= RETRY_ATTEMPTS:
|
|
raise
|
|
await asyncio.sleep(0.2 * (attempt + 1))
|
|
raise RuntimeError('HTTP POST failed')
|
|
|
|
def _tokens_for_chain(self, chain: str) -> list[dict]:
|
|
if chain in {'ETH', 'BSC'}:
|
|
return get_evm_tokens(chain)
|
|
if chain == 'TRX':
|
|
tokens = [dict(token) for token in get_trx_tokens()]
|
|
if settings.USDT_CONTRACT_ADDRESS:
|
|
for token in tokens:
|
|
if token['symbol'] == 'USDT':
|
|
token['contract'] = settings.USDT_CONTRACT_ADDRESS
|
|
return tokens
|
|
if chain == 'SOL':
|
|
return get_sol_tokens()
|
|
return []
|
|
|
|
def _format_units(self, raw: str, decimals: int) -> str:
|
|
if not raw or not raw.lstrip('-').isdigit():
|
|
return '0'
|
|
if decimals == 0:
|
|
return raw
|
|
negative = raw.startswith('-')
|
|
value = raw[1:] if negative else raw
|
|
padded = value.rjust(decimals + 1, '0')
|
|
whole = padded[:-decimals]
|
|
fraction = padded[-decimals:].rstrip('0')
|
|
result = f'{whole}.{fraction}' if fraction else whole
|
|
return f'-{result}' if negative else result
|
|
|
|
def _round_usd(self, value: Decimal) -> Decimal:
|
|
return value.quantize(Decimal('0.00000001'))
|
|
|
|
def _decimal_or_none(self, value: Any) -> Decimal | None:
|
|
try:
|
|
return Decimal(str(value))
|
|
except (InvalidOperation, TypeError, ValueError):
|
|
return None
|
|
|
|
def _hex_to_int_string(self, value: str | None) -> str:
|
|
if not value:
|
|
return '0'
|
|
return str(int(value, 16))
|
|
|
|
def _tron_headers(self) -> dict[str, str]:
|
|
headers = {'Accept': 'application/json'}
|
|
if settings.TRON_API_KEY:
|
|
headers['TRON-PRO-API-KEY'] = settings.TRON_API_KEY
|
|
return headers
|
|
|
|
def _coingecko_headers(self) -> dict[str, str]:
|
|
headers = {'Accept': 'application/json'}
|
|
if settings.COINGECKO_API_KEY:
|
|
headers['x-cg-demo-api-key'] = settings.COINGECKO_API_KEY
|
|
return headers
|
|
|
|
def _tron_address_to_hex(self, address: str) -> str:
|
|
num = 0
|
|
for char in address:
|
|
index = BASE58_ALPHABET.find(char)
|
|
if index == -1:
|
|
raise ValueError('Invalid TRON address')
|
|
num = num * 58 + index
|
|
raw = num.to_bytes(25, 'big')
|
|
payload = raw[:-4]
|
|
checksum = raw[-4:]
|
|
expected = hashlib.sha256(hashlib.sha256(payload).digest()).digest()[:4]
|
|
if checksum != expected or payload[0] != 0x41:
|
|
raise ValueError('Invalid TRON address')
|
|
return payload[1:].hex()
|