From 8c83ea6c87fa3d47369b336e60fe0dc491e0061d Mon Sep 17 00:00:00 2001
From: Dobromir Popov <d-popov@abv.bg>
Date: Wed, 2 Oct 2024 17:15:06 +0300
Subject: [PATCH] wip get ballances

---
 crypto/sol/app.py | 83 +++++++++++++++++++++++++++++++++--------------
 1 file changed, 59 insertions(+), 24 deletions(-)

diff --git a/crypto/sol/app.py b/crypto/sol/app.py
index cf6675b..3715762 100644
--- a/crypto/sol/app.py
+++ b/crypto/sol/app.py
@@ -13,6 +13,9 @@ import logging
 from solana.rpc.websocket_api import connect
 from solana.rpc.async_api import AsyncClient
 from solana.rpc.commitment import Confirmed
+from solana.rpc.types import TokenAccountOpts
+
+
 import os
 from dotenv import load_dotenv
 
@@ -71,13 +74,35 @@ async def get_token_balance(wallet_address, token_address):
         logging.error(f"Error getting balance for {token_address} in {wallet_address}: {str(e)}")
         return 0
 
+
 async def get_wallet_balances(wallet_address):
     balances = {}
     logging.info(f"Getting balances for wallet: {wallet_address}")
-    for token, address in TOKEN_ADDRESSES.items():
-        balance = await get_token_balance(wallet_address, address)
-        balances[token] = balance
-        logging.debug(f"Balance for {token}: {balance}")
+    
+    opts = TokenAccountOpts(
+        program_id=Pubkey.from_string("TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA")
+    )
+
+    # Get all token accounts for the wallet
+    response = await solana_client.get_token_accounts_by_owner(
+        Pubkey.from_string(wallet_address),
+        opts,
+        commitment=Confirmed
+    )
+    
+    if response.value:
+        for account in response.value:
+            mint = account.account.data.parsed.info.mint
+            balance_response = await solana_client.get_token_account_balance(account.pubkey)
+            if balance_response.value:
+                amount = float(balance_response.value.uiAmount)
+                balances[mint] = amount
+                logging.debug(f"Balance for {mint}: {amount}")
+    
+    sol_balance = await solana_client.get_balance(Pubkey.from_string(wallet_address))
+    if sol_balance.value:
+        balances['SOL'] = sol_balance.value / 1e9
+    
     return balances
 
 async def get_non_zero_token_balances(wallet_address):
@@ -153,43 +178,47 @@ async def follow_move(move):
         logging.warning(message)
         await send_telegram_message(message)
 
-async def on_logs(logs):
-    print(f"Received logs: {logs}")
-    if logs.value.err:
+async def on_logs(log):
+    print(f"Received log: {log}")
+    if 'err' in log and log['err']:
         return
     
-    tx = logs.value.signature
+    tx = log['signature']
     tx_result = await solana_client.get_transaction(tx)
     
-    if tx_result and tx_result.value:
-        for instruction in tx_result.value.transaction.message.instructions:
-            if instruction.program_id == Pubkey.from_string(TOKEN_ADDRESSES['SOL']):
+    if tx_result and 'result' in tx_result and tx_result['result']:
+        transaction = tx_result['result']['transaction']
+        message = transaction['message']
+        for instruction in message['instructions']:
+            if instruction['programId'] == TOKEN_ADDRESSES['SOL']:
                 # This is a token transfer
-                from_pubkey = instruction.accounts[0]
-                to_pubkey = instruction.accounts[1]
-                amount = instruction.data[8:]  # The amount is stored in the last 8 bytes
-                amount = int.from_bytes(amount, 'little') / 1e9  # Convert lamports to SOL
+                from_pubkey = instruction['accounts'][0]
+                to_pubkey = instruction['accounts'][1]
+                amount = int(instruction['data'], 16) / 1e9  # Convert lamports to SOL
                 
-                if from_pubkey == Pubkey.from_string(FOLLOWED_WALLET):
+                if from_pubkey == FOLLOWED_WALLET:
                     move = {
                         'token': 'SOL',
                         'amount': amount,
                         'to_token': 'Unknown'  # You might want to determine this based on the receiving address
                     }
-                    await follow_move(move)
-
+                    await follow_move(move)                    
 async def subscribe_to_wallet():
     uri = SOLANA_URL
     async with websockets.connect(uri) as websocket:
-        # Correct the `params` format to be an object
+        # Correct the `params` format to be an array
         request = {
             "jsonrpc": "2.0",
             "id": 1,
             "method": "logsSubscribe",
-            "params": {
-                "mentions": [YOUR_WALLET],
-                "commitment": "confirmed"
-            }
+            "params": [
+                {
+                    "mentions": [FOLLOWED_WALLET]  # Changed from YOUR_WALLET to FOLLOWED_WALLET
+                },
+                {
+                    "commitment": "confirmed"
+                }
+            ]
         }
 
         await websocket.send(json.dumps(request))
@@ -197,7 +226,13 @@ async def subscribe_to_wallet():
         # Listen for messages
         while True:
             response = await websocket.recv()
-            print(response)
+            response_data = json.loads(response)
+            if 'result' in response_data:
+                print(f"Subscription successful. Subscription id: {response_data['result']}")
+            elif 'params' in response_data:
+                await on_logs(response_data['params']['result'])
+            else:
+                print(f"Unexpected response: {response}")
             
 async def main():
     logging.basicConfig(level=logging.INFO)