def test_find_program_address():
"""Test create associated_token_address."""
program_id = PublicKey("BPFLoader1111111111111111111111111111111111")
program_address, nonce = PublicKey.find_program_address([bytes()],
program_id)
assert program_address == PublicKey.create_program_address(
[bytes(), helpers.to_uint8_bytes(nonce)], program_id)
def get_associated_token_address(owner: PublicKey,
mint: PublicKey) -> PublicKey:
"""Derives the associated token address for the given wallet address and token mint."""
key, _ = PublicKey.find_program_address(
seeds=[bytes(owner),
bytes(TOKEN_PROGRAM_ID),
bytes(mint)],
program_id=ASSOCIATED_TOKEN_PROGRAM_ID)
return key
async def test_new_associated_account(test_token): # pylint: disable=redefined-outer-name
"""Test creating a new associated token account."""
new_acct = PublicKey(0)
token_account_pubkey = await test_token.create_associated_token_account(new_acct)
expected_token_account_key, _ = new_acct.find_program_address(
seeds=[bytes(new_acct), bytes(TOKEN_PROGRAM_ID), bytes(test_token.pubkey)],
program_id=ASSOCIATED_TOKEN_PROGRAM_ID,
)
assert token_account_pubkey == expected_token_account_key
def get_neon_erc20_account_address(self, neon_account_address: str):
neon_contract_address_bytes = bytes.fromhex(
self.neon_contract_address[2:])
neon_account_address_bytes = bytes.fromhex(neon_account_address[2:])
seeds = [
b"\1", b"ERC20Balance",
bytes(self.token.pubkey), neon_contract_address_bytes,
neon_account_address_bytes
]
return PublicKey.find_program_address(seeds, self.evm_loader_id)[0]
def ether2program(ether):
if isinstance(ether, str):
pass
elif isinstance(ether, EthereumAddress):
ether = str(ether)
else:
ether = ether.hex()
if ether[0:2] == '0x':
ether = ether[2:]
seed = [ACCOUNT_SEED_VERSION, bytes.fromhex(ether)]
(pda, nonce) = PublicKey.find_program_address(seed,
PublicKey(EVM_LOADER_ID))
return str(pda), nonce
def get_neon_account_address(self, neon_account_address: str) -> PublicKey:
neon_account_addressbytes = bytes.fromhex(neon_account_address[2:])
return PublicKey.find_program_address(
[b"\1", neon_account_addressbytes], self.evm_loader_id)[0]
def get_evm_loader_account_address(eth_address: str):
eth_addressbytes = bytes.fromhex(eth_address[2:])
return PublicKey.find_program_address([b"\1", eth_addressbytes],
EVM_LOADER_ID)
def collect(self,
api_endpoint,
pool_account,
collector,
skip_confirmation=True):
msg = ""
client = Client(api_endpoint)
msg += "Initialized client"
signers = [Account(self.private_key)]
pool = self.load_binary_option(api_endpoint, pool_account)
pool_account = PublicKey(pool_account)
collector_account = PublicKey(collector)
escrow_account = PublicKey(pool["escrow"])
escrow_mint_account = PublicKey(pool["escrow_mint"])
long_token_mint_account = PublicKey(pool["long_mint"])
short_token_mint_account = PublicKey(pool["short_mint"])
token_account = PublicKey(TOKEN_PROGRAM_ID)
escrow_authority_account = PublicKey.find_program_address(
[
bytes(long_token_mint_account),
bytes(short_token_mint_account),
bytes(token_account),
bytes(PublicKey(BINARY_OPTION_PROGRAM_ID))
],
PublicKey(BINARY_OPTION_PROGRAM_ID),
)[0]
# Transaction
tx = Transaction()
atas = []
for mint_account in (long_token_mint_account, short_token_mint_account,
escrow_mint_account):
token_pda_address = get_associated_token_address(
collector_account, mint_account)
associated_token_account_info = client.get_account_info(
token_pda_address)
account_info = associated_token_account_info['result']['value']
if account_info is not None:
account_state = ACCOUNT_LAYOUT.parse(
base64.b64decode(account_info['data'][0])).state
else:
account_state = 0
if account_state == 0:
msg += f" | Error Fetching PDA: {token_pda_address}"
raise Exception()
else:
msg += f" | Fetched PDA: {token_pda_address}"
atas.append(token_pda_address)
collect_ix = collect_instruction(
pool_account,
collector_account,
atas[0],
atas[1],
atas[2],
long_token_mint_account,
short_token_mint_account,
escrow_account,
escrow_authority_account,
token_account,
)
tx = tx.add(collect_ix)
try:
response = client.send_transaction(
tx,
*signers,
opts=types.TxOpts(skip_confirmation=skip_confirmation))
return json.dumps({
'status':
HTTPStatus.OK,
'msg':
msg + f" | Collect successful",
'tx':
response.get('result') if skip_confirmation else
response['result']['transaction']['signatures'],
})
except Exception as e:
msg += f" | ERROR: Encountered exception while attempting to send transaction: {e}"
print(msg)
raise (e)
def trade(self,
api_endpoint,
pool_account,
buyer_encrypted_private_key,
seller_encrypted_private_key,
size,
buyer_price,
seller_price,
skip_confirmation=True):
msg = ""
client = Client(api_endpoint)
msg += "Initialized client"
# Create account objects
buyer_private_key = list(
self.cipher.decrypt(buyer_encrypted_private_key))
seller_private_key = list(
self.cipher.decrypt(seller_encrypted_private_key))
assert (len(buyer_private_key) == 32)
assert (len(seller_private_key) == 32)
source_account = Account(self.private_key)
buyer = Account(buyer_private_key)
seller = Account(seller_private_key)
# Signers
signers = [buyer, seller, source_account]
pool = self.load_binary_option(api_endpoint, pool_account)
# List non-derived accounts
pool_account = PublicKey(pool_account)
escrow_account = PublicKey(pool["escrow"])
escrow_mint_account = PublicKey(pool["escrow_mint"])
long_token_mint_account = PublicKey(pool["long_mint"])
short_token_mint_account = PublicKey(pool["short_mint"])
buyer_account = buyer.public_key()
seller_account = seller.public_key()
token_account = PublicKey(TOKEN_PROGRAM_ID)
escrow_owner_account = PublicKey.find_program_address(
[
bytes(long_token_mint_account),
bytes(short_token_mint_account),
bytes(token_account),
bytes(PublicKey(BINARY_OPTION_PROGRAM_ID))
],
PublicKey(BINARY_OPTION_PROGRAM_ID),
)[0]
# Transaction
tx = Transaction()
atas = []
for acct in [buyer_account, seller_account]:
acct_atas = []
for mint_account in (long_token_mint_account,
short_token_mint_account,
escrow_mint_account):
token_pda_address = get_associated_token_address(
acct, mint_account)
associated_token_account_info = client.get_account_info(
token_pda_address)
account_info = associated_token_account_info['result']['value']
if account_info is not None:
account_state = ACCOUNT_LAYOUT.parse(
base64.b64decode(account_info['data'][0])).state
else:
account_state = 0
if account_state == 0:
msg += f" | Creating PDA: {token_pda_address}"
associated_token_account_ix = create_associated_token_account(
payer=source_account.public_key(),
owner=acct,
mint=mint_account,
)
tx = tx.add(associated_token_account_ix)
else:
msg += f" | Fetched PDA: {token_pda_address}"
acct_atas.append(token_pda_address)
atas.append(acct_atas)
trade_ix = trade_instruction(
pool_account,
escrow_account,
long_token_mint_account,
short_token_mint_account,
buyer_account,
seller_account,
atas[0][2],
atas[1][2],
atas[0][0],
atas[0][1],
atas[1][0],
atas[1][1],
escrow_owner_account,
token_account,
int(size),
int(buyer_price),
int(seller_price),
)
tx = tx.add(trade_ix)
# Send request
try:
response = client.send_transaction(
tx,
*signers,
opts=types.TxOpts(skip_confirmation=skip_confirmation))
return json.dumps({
'status':
HTTPStatus.OK,
'msg':
msg + f" | Trade successful",
'tx':
response.get('result') if skip_confirmation else
response['result']['transaction']['signatures'],
})
except Exception as e:
msg += f" | ERROR: Encountered exception while attempting to send transaction: {e}"
raise (e)
def refresh_user_ids(
redis: Redis,
db: SessionManager,
token_contract,
delegate_manager_contract,
staking_contract,
eth_web3,
waudio_token,
):
with db.scoped_session() as session:
lazy_refresh_user_ids = get_lazy_refresh_user_ids(
redis, session)[:MAX_LAZY_REFRESH_USER_IDS]
immediate_refresh_user_ids = get_immediate_refresh_user_ids(redis)
logger.info(
f"cache_user_balance.py | Starting refresh with {len(lazy_refresh_user_ids)} "
f"lazy refresh user_ids: {lazy_refresh_user_ids} and {len(immediate_refresh_user_ids)} "
f"immediate refresh user_ids: {immediate_refresh_user_ids}")
all_user_ids = lazy_refresh_user_ids + immediate_refresh_user_ids
user_ids = list(set(all_user_ids))
existing_user_balances: List[UserBalance] = ((
session.query(UserBalance)).filter(
UserBalance.user_id.in_(user_ids)).all())
all_user_balance = existing_user_balances
# Balances from current user lookup may
# not be present in the db, so make those
not_present_set = set(user_ids) - {
user.user_id
for user in existing_user_balances
}
new_balances: List[UserBalance] = [
UserBalance(
user_id=user_id,
balance="0",
associated_wallets_balance="0",
associated_sol_wallets_balance="0",
) for user_id in not_present_set
]
if new_balances:
session.add_all(new_balances)
all_user_balance = existing_user_balances + new_balances
logger.info(
f"cache_user_balance.py | adding new users: {not_present_set}")
user_balances = {user.user_id: user for user in all_user_balance}
# Grab the users & associated_wallets we need to refresh
user_associated_wallet_query: List[Tuple[int, str, str, str]] = (
session.query(
User.user_id,
User.wallet,
AssociatedWallet.wallet,
AssociatedWallet.chain,
).outerjoin(
AssociatedWallet,
and_(
AssociatedWallet.user_id == User.user_id,
AssociatedWallet.is_current == True,
AssociatedWallet.is_delete == False,
),
).filter(
User.user_id.in_(user_ids),
User.is_current == True,
).all())
user_bank_accounts_query: List[Tuple[int, str]] = (session.query(
User.user_id, UserBankAccount.bank_account).join(
UserBankAccount,
UserBankAccount.ethereum_address == User.wallet).filter(
User.user_id.in_(user_ids),
User.is_current == True,
).all())
user_id_bank_accounts = dict(user_bank_accounts_query)
# Combine query results for user bank, associated wallets,
# and primary owner wallet into a single metadata list
user_id_metadata: Dict[int, UserWalletMetadata] = {}
for user in user_associated_wallet_query:
user_id, user_wallet, associated_wallet, wallet_chain = user
if user_id not in user_id_metadata:
user_id_metadata[user_id] = {
"owner_wallet": user_wallet,
"associated_wallets": {
"eth": [],
"sol": []
},
"bank_account": None,
}
if user_id in user_id_bank_accounts:
user_id_metadata[user_id][
"bank_account"] = user_id_bank_accounts[user_id]
if associated_wallet:
if user_id in user_id_metadata:
user_id_metadata[user_id]["associated_wallets"][
wallet_chain # type: ignore
].append(associated_wallet)
logger.info(
f"cache_user_balance.py | fetching for {len(user_associated_wallet_query)} users: {user_ids}"
)
# mapping of user_id => balance change
needs_balance_change_update: Dict[int, Dict] = {}
# Fetch balances
# pylint: disable=too-many-nested-blocks
for user_id, wallets in user_id_metadata.items():
try:
owner_wallet = wallets["owner_wallet"]
owner_wallet = eth_web3.toChecksumAddress(owner_wallet)
owner_wallet_balance = token_contract.functions.balanceOf(
owner_wallet).call()
associated_balance = 0
waudio_balance: str = "0"
associated_sol_balance = 0
if "associated_wallets" in wallets:
for wallet in wallets["associated_wallets"]["eth"]:
wallet = eth_web3.toChecksumAddress(wallet)
balance = token_contract.functions.balanceOf(
wallet).call()
delegation_balance = (
delegate_manager_contract.functions.
getTotalDelegatorStake(wallet).call())
stake_balance = staking_contract.functions.totalStakedFor(
wallet).call()
associated_balance += (balance + delegation_balance +
stake_balance)
if waudio_token is not None:
for wallet in wallets["associated_wallets"]["sol"]:
try:
root_sol_account = PublicKey(wallet)
derived_account, _ = PublicKey.find_program_address(
[
bytes(root_sol_account),
bytes(SPL_TOKEN_ID_PK),
bytes(WAUDIO_MINT_PUBKEY
), # type: ignore
],
ASSOCIATED_TOKEN_PROGRAM_ID_PK,
)
bal_info = waudio_token.get_balance(
derived_account)
associated_waudio_balance: str = bal_info[
"result"]["value"]["amount"]
associated_sol_balance += int(
associated_waudio_balance)
except Exception as e:
logger.error(
" ".join([
"cache_user_balance.py | Error fetching associated ",
"wallet balance for user %s, wallet %s: %s",
]),
user_id,
wallet,
e,
)
if wallets["bank_account"] is not None:
if waudio_token is None:
logger.error(
"cache_user_balance.py | Missing Required SPL Confirguration"
)
else:
bal_info = waudio_token.get_balance(
PublicKey(wallets["bank_account"]))
waudio_balance = bal_info["result"]["value"]["amount"]
# update the balance on the user model
user_balance = user_balances[user_id]
# Convert Sol balances to wei
waudio_in_wei = to_wei(waudio_balance)
assoc_sol_balance_in_wei = to_wei(associated_sol_balance)
user_waudio_in_wei = to_wei(user_balance.waudio)
user_assoc_sol_balance_in_wei = to_wei(
user_balance.associated_sol_wallets_balance)
# Get values for user balance change
current_total_balance = (owner_wallet_balance +
associated_balance + waudio_in_wei +
assoc_sol_balance_in_wei)
prev_total_balance = (
int(user_balance.balance) +
int(user_balance.associated_wallets_balance) +
user_waudio_in_wei + user_assoc_sol_balance_in_wei)
# Write to user_balance_changes table
needs_balance_change_update[user_id] = {
"user_id": user_id,
"blocknumber": eth_web3.eth.block_number,
"current_balance": str(current_total_balance),
"previous_balance": str(prev_total_balance),
}
user_balance.balance = str(owner_wallet_balance)
user_balance.associated_wallets_balance = str(
associated_balance)
user_balance.waudio = waudio_balance
user_balance.associated_sol_wallets_balance = str(
associated_sol_balance)
except Exception as e:
logger.error(
f"cache_user_balance.py | Error fetching balance for user {user_id}: {(e)}"
)
# Outside the loop, batch update the UserBalanceChanges:
# Get existing user balances
user_balance_ids = list(needs_balance_change_update.keys())
existing_user_balance_changes: List[UserBalanceChange] = (
session.query(UserBalanceChange).filter(
UserBalanceChange.user_id.in_(user_balance_ids)).all())
# Find all the IDs that don't already exist in the DB
to_create_ids = set(user_balance_ids) - {
e.user_id
for e in existing_user_balance_changes
}
logger.info(
f"cache_user_balance.py | UserBalanceChanges needing update: {user_balance_ids},\
existing: {[e.user_id for e in existing_user_balance_changes]}, to create: {to_create_ids}"
)
# Create new entries for those IDs
balance_changes_to_add = [
UserBalanceChange(
user_id=user_id,
blocknumber=needs_balance_change_update[user_id]
["blocknumber"],
current_balance=needs_balance_change_update[user_id]
["current_balance"],
previous_balance=needs_balance_change_update[user_id]
["previous_balance"],
) for user_id in to_create_ids
]
session.add_all(balance_changes_to_add)
# Lastly, update all the existing entries
for change in existing_user_balance_changes:
new_values = needs_balance_change_update[change.user_id]
change.blocknumber = new_values["blocknumber"]
change.current_balance = new_values["current_balance"]
change.previous_balance = new_values["previous_balance"]
# Commit the new balances
session.commit()
# Remove the fetched balances from Redis set
logger.info(
f"cache_user_balance.py | Got balances for {len(user_associated_wallet_query)} users, removing from Redis."
)
if lazy_refresh_user_ids:
redis.srem(LAZY_REFRESH_REDIS_PREFIX, *lazy_refresh_user_ids)
if immediate_refresh_user_ids:
redis.srem(IMMEDIATE_REFRESH_REDIS_PREFIX,
*immediate_refresh_user_ids)
def get_base_address(mint, program_id):
return PublicKey.find_program_address([bytes(mint)], program_id)