fix: add olas balance

report.py

@@ -142,6 +142,11 @@ def generate_report():
                 usdc_balance_formatted = safe_info.get('usdc_balance_formatted', 'N/A')
                 _print_status("USDC Balance", usdc_balance_formatted)
 
+            # Check for OLAS balance on Staking chain
+            if chain_name.lower() == optimus_config.staking_chain:
+                olas_balance_formatted = safe_info.get('olas_balance_formatted', 'N/A')
+                _print_status("OLAS Balance", olas_balance_formatted)
+
             # Low balance check
             safe_threshold_wei = chain_config.get("chain_data", {}).get("user_params", {}).get("fund_requirements", {}).get("safe")
             if safe_threshold_wei:

wallet_info.py

@@ -29,14 +29,16 @@
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(message)s')
 
-USDC_ABI = [{
+TOKEN_ABI = [{
     "constant": True,
     "inputs": [{"name": "_owner", "type": "address"}],
     "name": "balanceOf",
     "outputs": [{"name": "balance", "type": "uint256"}],
     "type": "function"
 }]
 
+OLAS_ADDRESS = "0xcfD1D50ce23C46D3Cf6407487B2F8934e96DC8f9"
+
 def load_config():
     try:
         optimus_config = load_local_config()
@@ -69,13 +71,22 @@ def get_balance(web3, address):
 
 def get_usdc_balance(web3, address, chain_name):
     try:
-        usdc_contract = web3.eth.contract(address=USDC_ADDRESS, abi=USDC_ABI)
+        usdc_contract = web3.eth.contract(address=USDC_ADDRESS, abi=TOKEN_ABI)
         balance = usdc_contract.functions.balanceOf(address).call()
         return Decimal(balance) / Decimal(1e6)  # USDC has 6 decimal places
     except Exception as e:
         print(f"Error getting USDC balance for address {address}: {e}")
         return Decimal(0)
 
+def get_olas_balance(web3, address, chain_name):
+    try:
+        olas_address = OLAS_ADDRESS[chain_name]
+        olas_contract = web3.eth.contract(address=OLAS_ADDRESS, abi=TOKEN_ABI)
+        balance = olas_contract.functions.balanceOf(address).call()
+        return Decimal(balance) / Decimal(1e18)  # OLAS has 18 decimal places
+    except Exception as e:
+        print(f"Error getting OLAS balance for address {address}: {e}")
+        return Decimal(0)
 class DecimalEncoder(json.JSONEncoder):
     def default(self, o):
         if isinstance(o, Decimal):
@@ -148,6 +159,12 @@ def save_wallet_info():
                 safe_balances[chain_name]["usdc_balance"] = usdc_balance
                 safe_balances[chain_name]["usdc_balance_formatted"] = f"{usdc_balance:.2f} USDC"
 
+            # Get USDC balance for Principal Chain
+            if chain_name.lower() == optimus_config.staking_chain:
+                olas_balance = get_olas_balance(web3, safe_address, chain_name.lower())
+                safe_balances[chain_name]["olas_balance"] = usdc_balance
+                safe_balances[chain_name]["olas_balance_formatted"] = f"{olas_balance:.6f} OLAS"
+
         except Exception as e:
             print(f"An error occurred while processing chain ID {chain_id}: {e}")
             continue