Add sanity checks to Derive client & fix manager address handling

derive_client/__init__.py

@@ -1,6 +1,7 @@
 """
 Init for the derive client
 """
+
 from .derive import DeriveClient
 
 DeriveClient

derive_client/cli.py

@@ -1,6 +1,7 @@
 """
 Cli module in order to allow interaction.
 """
+
 import os
 from textwrap import dedent
 

derive_client/clients/base_client.py

@@ -1,9 +1,11 @@
 """
 Base Client for the derive dex.
 """
+
 import json
 import random
 from decimal import Decimal
+from logging import Logger
 from time import sleep
 
 import eth_abi
@@ -18,6 +20,7 @@
 )
 from derive_action_signing.signed_action import SignedAction
 from derive_action_signing.utils import MAX_INT_32, get_action_nonce, sign_rest_auth_header, sign_ws_login, utc_now_ms
+from pydantic import validate_arguments
 from rich import print
 from web3 import Web3
 from websocket import WebSocketConnectionClosedException, create_connection
@@ -33,10 +36,14 @@
     Currency,
     Environment,
     InstrumentType,
+    MainnetCurrency,
+    ManagerAddress,
+    MarginType,
     OrderSide,
     OrderStatus,
     OrderType,
     RfqStatus,
+    SessionKey,
     SubaccountType,
     TimeInForce,
     TxResult,
@@ -64,35 +71,47 @@ def _create_signature_headers(self):
             session_key_or_wallet_private_key=self.signer._private_key,
         )
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def __init__(
         self,
-        wallet: str,
+        wallet: Address,
         private_key: str,
         env: Environment,
-        logger=None,
-        verbose=False,
-        subaccount_id=None,
-        referral_code=None,
+        logger: Logger | None = None,
+        verbose: bool = False,
+        subaccount_id: int | None = None,
+        referral_code: Address | None = None,
     ):
         self.verbose = verbose
         self.env = env
         self.config = CONFIGS[env]
         self.logger = logger or get_logger()
         self.web3_client = Web3()
         self.signer = self.web3_client.eth.account.from_key(private_key)
-        self.wallet = self.signer.address if not wallet else wallet
-        if subaccount_id is None:
-            subaccount_id = self._get_first_subaccount_id()
-        self.subaccount_id = int(subaccount_id)
+        self.wallet = wallet
+        self._verify_wallet(wallet)
+        self.subaccount_id = self._determine_subaccount_id(subaccount_id)
         self.referral_code = referral_code
 
-    def _get_first_subaccount_id(self) -> int:
-        self.logger.debug("No subaccount_id provided, fetching from API…")
+    def _verify_wallet(self, wallet: Address):
+        w3 = Web3(Web3.HTTPProvider(self.config.rpc_endpoint))
+        if not w3.is_connected():
+            raise ConnectionError(f"Failed to connect to RPC at {self.config.rpc_endpoint}")
+        if not w3.eth.get_code(wallet):
+            msg = f"{wallet} appears to be an EOA (no bytecode). Expected a smart-contract wallet on Derive."
+            raise ValueError(msg)
+        session_keys = self._get_session_keys(wallet)
+        if not any(self.signer.address == s.public_session_key for s in session_keys):
+            msg = f"{self.signer.address} is not among registered session keys for wallet {wallet}."
+            raise ValueError(msg)
+
+    def _determine_subaccount_id(self, subaccount_id: int | None) -> int:
         subaccounts = self.fetch_subaccounts()
-        self.logger.info(f"Subaccounts retrieved: {subaccounts!r}")
         if not (subaccount_ids := subaccounts.get("subaccount_ids", [])):
-            raise ValueError("No subaccounts found. Please create one on Derive first.")
-        subaccount_id = subaccount_ids[0]
+            raise ValueError(f"No subaccounts found for {self.wallet}. Please create one on Derive first.")
+        if subaccount_id is not None and subaccount_id not in subaccount_ids:
+            raise ValueError(f"Provided subaccount {subaccount_id} not among retrieved aubaccounts: {subaccounts!r}")
+        subaccount_id = subaccount_id or subaccount_ids[0]
         self.logger.info(f"Selected subaccount_id: {subaccount_id}")
         return subaccount_id
 
@@ -183,6 +202,15 @@ def fetch_instruments(
         }
         return self._send_request(url, json=payload, headers=PUBLIC_HEADERS)
 
+    def _get_session_keys(self, wallet: Address) -> list[SessionKey]:
+        url = f"{self.config.base_url}/private/session_keys"
+        payload = {"wallet": wallet}
+        session_keys = self._send_request(url, json=payload)
+        if not (public_session_keys := session_keys.get("public_session_keys")):
+            msg = f"No session keys registered for this wallet: {wallet}"
+            raise ValueError(msg)
+        return list(map(lambda kwargs: SessionKey(**kwargs), public_session_keys))
+
     def fetch_subaccounts(self):
         """
         Returns the subaccounts for a given wallet
@@ -191,7 +219,7 @@ def fetch_subaccounts(self):
         payload = {"wallet": self.wallet}
         return self._send_request(url, json=payload)
 
-    def fetch_subaccount(self, subaccount_id):
+    def fetch_subaccount(self, subaccount_id: int):
         """
         Returns information for a given subaccount
         """
@@ -209,7 +237,7 @@ def _internal_map_instrument(self, instrument_type, currency):
     def create_order(
         self,
         price,
-        amount,
+        amount: int,
         instrument_name: str,
         reduce_only=False,
         instrument_type: InstrumentType = InstrumentType.PERP,
@@ -255,6 +283,7 @@ def create_order(
             "referral_code": DEFAULT_REFERER if not self.referral_code else self.referral_code,
             **signed_action.to_json(),
         }
+        # breakpoint()
         response = self.submit_order(order)
         return response
 
@@ -504,7 +533,7 @@ def fetch_tickers(
 
     def create_subaccount(
         self,
-        amount=0,
+        amount: int = 0,
         subaccount_type: SubaccountType = SubaccountType.STANDARD,
         collateral_asset: CollateralAsset = CollateralAsset.USDC,
         underlying_currency: UnderlyingCurrency = UnderlyingCurrency.ETH,
@@ -629,7 +658,7 @@ def get_mmp_config(self, subaccount_id: int, currency: UnderlyingCurrency = None
 
     def set_mmp_config(
         self,
-        subaccount_id,
+        subaccount_id: int,
         currency: UnderlyingCurrency,
         mmp_frozen_time: int,
         mmp_interval: int,
@@ -775,7 +804,7 @@ def transfer_from_funding_to_subaccount(self, amount: int, asset_name: str, suba
             json=payload,
         )
 
-    def get_manager_for_subaccount(self, subaccount_id, asset_name):
+    def get_manager_for_subaccount(self, subaccount_id: int, asset_name):
         """
         Look up the manager for a subaccount
 
@@ -786,24 +815,25 @@ def get_manager_for_subaccount(self, subaccount_id, asset_name):
         deposit_currency = UnderlyingCurrency[asset_name]
         currency = self.fetch_currency(asset_name)
         underlying_address = currency['protocol_asset_addresses']['spot']
-        manager_addresses = currency['managers']
-
-        if len(manager_addresses) == 1:
-            manager_address = manager_addresses[0].get('address')
-        else:
-            to_account = self.fetch_subaccount(subaccount_id)
-            account_type = (
-                SubaccountType.STANDARD if to_account.get("margin_type") == "SM" else SubaccountType.PORTFOLIO
-            )
-            account_currency = UnderlyingCurrency[to_account.get("currency")]
-            index = (
-                0 if account_type is SubaccountType.STANDARD else 1 if account_currency is UnderlyingCurrency.ETH else 2
-            )
-            manager_address = manager_addresses[index].get('address')
-
-        if not manager_address or not underlying_address:
+        managers = list(map(lambda kwargs: ManagerAddress(**kwargs), currency['managers']))
+        manager_by_type = {}
+        for manager in managers:
+            manager_by_type.setdefault((manager.margin_type, manager.currency), []).append(manager)
+
+        to_account = self.fetch_subaccount(subaccount_id)
+        account_currency = MainnetCurrency[to_account.get("currency")]
+        margin_type = MarginType[to_account.get("margin_type")]
+
+        def get_unique_manager(margin_type, currency):
+            matches = manager_by_type.get((margin_type, currency), [])
+            if len(matches) != 1:
+                raise ValueError(f"Expected exactly one ManagerAddress for {(margin_type, currency)}, found {matches}")
+            return matches[0]
+
+        manager = get_unique_manager(margin_type, account_currency)
+        if not manager.address or not underlying_address:
             raise Exception(f"Unable to find manager address or underlying address for {asset_name}")
-        return manager_address, underlying_address, TOKEN_DECIMALS[deposit_currency]
+        return manager.address, underlying_address, TOKEN_DECIMALS[deposit_currency]
 
     def transfer_from_subaccount_to_funding(self, amount: int, asset_name: str, subaccount_id: int):
         """

derive_client/constants.py

@@ -34,6 +34,8 @@ def __getitem__(self, key):
 class EnvConfig(BaseModel, frozen=True):
     base_url: str
     ws_address: str
+    rpc_endpoint: str
+    block_explorer: str
     ACTION_TYPEHASH: str
     DOMAIN_SEPARATOR: str
     contracts: ContractAddresses
@@ -52,6 +54,8 @@ class EnvConfig(BaseModel, frozen=True):
     Environment.TEST: EnvConfig(
         base_url="https://api-demo.lyra.finance",
         ws_address="wss://api-demo.lyra.finance/ws",
+        rpc_endpoint="https://rpc-prod-testnet-0eakp60405.t.conduit.xyz",
+        block_explorer="https://explorer-prod-testnet-0eakp60405.t.conduit.xyz",
         ACTION_TYPEHASH="0x4d7a9f27c403ff9c0f19bce61d76d82f9aa29f8d6d4b0c5474607d9770d1af17",
         DOMAIN_SEPARATOR="0x9bcf4dc06df5d8bf23af818d5716491b995020f377d3b7b64c29ed14e3dd1105",
         contracts=ContractAddresses(
@@ -76,6 +80,8 @@ class EnvConfig(BaseModel, frozen=True):
     Environment.PROD: EnvConfig(
         base_url="https://api.lyra.finance",
         ws_address="wss://api.lyra.finance/ws",
+        rpc_endpoint="https://rpc.lyra.finance",
+        block_explorer="https://explorer.lyra.finance",
         ACTION_TYPEHASH="0x4d7a9f27c403ff9c0f19bce61d76d82f9aa29f8d6d4b0c5474607d9770d1af17",
         DOMAIN_SEPARATOR="0xd96e5f90797da7ec8dc4e276260c7f3f87fedf68775fbe1ef116e996fc60441b",
         contracts=ContractAddresses(

derive_client/data_types/__init__.py

@@ -7,6 +7,8 @@
     Currency,
     Environment,
     InstrumentType,
+    MainnetCurrency,
+    MarginType,
     OrderSide,
     OrderStatus,
     OrderType,
@@ -22,8 +24,10 @@
     CreateSubAccountData,
     CreateSubAccountDetails,
     DeriveAddresses,
+    ManagerAddress,
     MintableTokenData,
     NonMintableTokenData,
+    SessionKey,
     TxResult,
 )
 
@@ -45,9 +49,13 @@
     "ActionType",
     "RfqStatus",
     "Address",
+    "SessionKey",
     "MintableTokenData",
     "NonMintableTokenData",
     "DeriveAddresses",
     "CreateSubAccountDetails",
     "CreateSubAccountData",
+    "MainnetCurrency",
+    "MarginType",
+    "ManagerAddress",
 ]

derive_client/data_types/enums.py

@@ -38,6 +38,23 @@ class RPCEndPoints(Enum):
     DERIVE = LYRA = "https://rpc.lyra.finance"
 
 
+class SessionKeyScope(Enum):
+    ADMIN = "admin"
+    ACCOUNT = "account"
+    READ_ONLY = "read_only"
+
+
+class MainnetCurrency(Enum):
+    BTC = "BTC"
+    ETH = "ETH"
+
+
+class MarginType(Enum):
+    SM = "SM"
+    PM = "PM"
+    PM2 = "PM2"
+
+
 class InstrumentType(Enum):
     """Instrument types."""
 

derive_client/data_types/models.py

@@ -1,18 +1,33 @@
 """Models used in the bridge module."""
 
-
 from dataclasses import dataclass
 
 from derive_action_signing.module_data import ModuleData
 from derive_action_signing.utils import decimal_to_big_int
 from eth_abi.abi import encode
-from pydantic import BaseModel, ConfigDict
+from eth_utils import is_address, to_checksum_address
+from pydantic import BaseModel, ConfigDict, GetCoreSchemaHandler, GetJsonSchemaHandler
+from pydantic_core import core_schema
 from web3 import Web3
 from web3.datastructures import AttributeDict
 
-from .enums import ChainID, Currency, TxStatus
+from .enums import ChainID, Currency, MainnetCurrency, MarginType, SessionKeyScope, TxStatus
+
+
+class Address(str):
+    @classmethod
+    def __get_pydantic_core_schema__(cls, _source, _handler: GetCoreSchemaHandler) -> core_schema.CoreSchema:
+        return core_schema.no_info_before_validator_function(cls._validate, core_schema.str_schema())
+
+    @classmethod
+    def __get_pydantic_json_schema__(cls, _schema, _handler: GetJsonSchemaHandler) -> dict:
+        return {"type": "string", "format": "ethereum-address"}
 
-Address = str
+    @classmethod
+    def _validate(cls, v: str) -> str:
+        if not is_address(v):
+            raise ValueError(f"Invalid Ethereum address: {v}")
+        return to_checksum_address(v)
 
 
 @dataclass
@@ -69,6 +84,20 @@ class DeriveAddresses(BaseModel):
     chains: dict[ChainID, dict[Currency, MintableTokenData | NonMintableTokenData]]
 
 
+class SessionKey(BaseModel):
+    public_session_key: Address
+    expiry_sec: int
+    ip_whitelist: list
+    label: str
+    scope: SessionKeyScope
+
+
+class ManagerAddress(BaseModel):
+    address: Address
+    margin_type: MarginType
+    currency: MainnetCurrency | None
+
+
 @dataclass
 class TxResult:
     tx_hash: str

examples/fetch_instruments.py

@@ -1,6 +1,7 @@
 """
 Sample of fetching instruments from the derive client, and printing the result.
 """
+
 from rich import print
 
 from derive_client.data_types import Environment, InstrumentType, UnderlyingCurrency