iwa 0.0.33__py3-none-any.whl → 0.0.59__py3-none-any.whl

iwa/core/keys.py

@@ -210,14 +210,14 @@ class KeyStorage(BaseModel):
         if not self.get_address_by_tag("master"):
             logger.info("Master account not found. Creating new 'master' account...")
             try:
-                self.create_account("master")
+                self.generate_new_account("master")
             except Exception as e:
                 logger.error(f"Failed to create master account: {e}")
 
     @property
-    def master_account(self) -> EncryptedAccount:
+    def master_account(self) -> Optional[Union[EncryptedAccount, StoredSafeAccount]]:
         """Get the master account"""
-        master_account = self.get_account("master")
+        master_account = self.find_stored_account("master")
 
         if not master_account:
             return list(self.accounts.values())[0]
@@ -231,19 +231,34 @@ class KeyStorage(BaseModel):
             # Use backup directory relative to wallet path (supports tests with tmp_path)
             backup_dir = self._path.parent / "backup"
             backup_dir.mkdir(parents=True, exist_ok=True)
+            try:
+                os.chmod(backup_dir, 0o700)
+            except OSError as e:
+                logger.debug(f"Could not chmod backup dir (expected in some Docker setups): {e}")
             timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
             backup_path = backup_dir / f"wallet.json.{timestamp}.bkp"
             shutil.copy2(self._path, backup_path)
+            try:
+                os.chmod(backup_path, 0o600)
+            except OSError as e:
+                logger.debug(f"Could not chmod backup file: {e}")
             logger.debug(f"Backed up wallet to {backup_path}")
 
         # Ensure directory exists
         self._path.parent.mkdir(parents=True, exist_ok=True)
 
         with open(self._path, "w", encoding="utf-8") as f:
-            json.dump(self.model_dump(), f, indent=4)
+            # Use mode='json' to ensure all types (EthereumAddress) are correctly serialized
+            json.dump(self.model_dump(mode="json"), f, indent=4)
+            f.flush()
+            os.fsync(f.fileno())  # Force write to disk (critical for Docker volumes)
 
-        # Enforce read/write only for the owner
-        os.chmod(self._path, 0o600)
+        try:
+            os.chmod(self._path, 0o600)
+        except OSError as e:
+            logger.debug(f"Could not chmod wallet file: {e}")
+
+        logger.info(f"[KeyStorage] Wallet saved to {self._path} ({len(self.accounts)} accounts)")
 
     @staticmethod
     def _encrypt_mnemonic(mnemonic: str, password: str) -> dict:
@@ -328,21 +343,20 @@ class KeyStorage(BaseModel):
         encrypted_acct = EncryptedAccount.encrypt_private_key(
             private_key_hex, self._password, "master"
         )
-        self.accounts[encrypted_acct.address] = encrypted_acct
-        self.save()
-
+        self.register_account(encrypted_acct)
         return encrypted_acct, mnemonic_str
 
-    def create_account(self, tag: str) -> EncryptedAccount:
-        """Create account. Master is derived from mnemonic, others are random."""
+    def generate_new_account(self, tag: str) -> EncryptedAccount:
+        """Generate a brand new EOA account and register it with the given tag."""
+        # Note: register_account(tag) check is inside, but we handle 'master' logic here
         tags = [acct.tag for acct in self.accounts.values()]
         if not tags:
             tag = "master"  # First account is always master
-        if tag in tags:
-            raise ValueError(f"Tag '{tag}' already exists in wallet.")
 
         # Master account: derive from mnemonic
         if tag == "master":
+            if "master" in tags:
+                raise ValueError("Master account already exists in wallet.")
             encrypted_acct, mnemonic = self._create_master_from_mnemonic()
             self._pending_mnemonic = mnemonic  # Store temporarily for display
             return encrypted_acct
@@ -350,10 +364,28 @@ class KeyStorage(BaseModel):
         # Non-master: random key as before
         acct = Account.create()
         encrypted = EncryptedAccount.encrypt_private_key(acct.key.hex(), self._password, tag)
-        self.accounts[acct.address] = encrypted
-        self.save()
+        self.register_account(encrypted)
         return encrypted
 
+    def register_account(self, account: Union[EncryptedAccount, StoredSafeAccount]):
+        """Register an account (EOA or Safe) in the storage with strict tag uniqueness checks."""
+        if not account.tag:
+            # Allow untagged accounts (rare but possible)
+            pass
+        else:
+            # Check for duplicate tags
+            for existing in self.accounts.values():
+                if existing.tag == account.tag and existing.address != account.address:
+                    raise ValueError(
+                        f"Tag '{account.tag}' is already used by address {existing.address}"
+                    )
+
+        self.accounts[account.address] = account
+        logger.info(
+            f"[KeyStorage] Registering account: tag='{account.tag}', address={account.address}"
+        )
+        self.save()
+
     def get_pending_mnemonic(self) -> Optional[str]:
         """Get and clear the pending mnemonic (for one-time display).
 
@@ -419,6 +451,24 @@ class KeyStorage(BaseModel):
         del self.accounts[account.address]
         self.save()
 
+    def rename_account(self, address_or_tag: str, new_tag: str):
+        """Rename an account's tag with uniqueness check."""
+        account = self.find_stored_account(address_or_tag)
+        if not account:
+            raise ValueError(f"Account '{address_or_tag}' not found.")
+
+        # Check if new tag is already used by a DIFFERENT account
+        for existing in self.accounts.values():
+            if existing.tag == new_tag and existing.address != account.address:
+                raise ValueError(f"Tag '{new_tag}' is already used by address {existing.address}")
+
+        old_tag = account.tag
+        account.tag = new_tag
+        logger.info(
+            f"[KeyStorage] Renaming account: '{old_tag}' -> '{new_tag}' (address={account.address})"
+        )
+        self.save()
+
     def _get_private_key(self, address: str) -> Optional[str]:
         """Get private key (Internal)"""
         account = self.accounts.get(EthereumAddress(address))

iwa/core/models.py

@@ -6,14 +6,26 @@ from typing import Dict, List, Optional, Type, TypeVar
 
 import tomli
 import tomli_w
-import yaml
 from pydantic import BaseModel, Field, PrivateAttr
 from pydantic_core import core_schema
+from ruamel.yaml import YAML
 
 from iwa.core.types import EthereumAddress  # noqa: F401 - re-exported for backwards compatibility
 from iwa.core.utils import singleton
 
 
+def _update_yaml_recursive(target: Dict, source: Dict) -> None:
+    """Recursively update a ruamel.yaml CommentedMap with data from a dict.
+
+    This preserves comments and structure in the target map.
+    """
+    for key, value in source.items():
+        if isinstance(value, dict) and key in target and isinstance(target[key], dict):
+            _update_yaml_recursive(target[key], value)
+        else:
+            target[key] = value
+
+
 class EncryptedData(BaseModel):
     """Encrypted data structure with explicit KDF parameters."""
 
@@ -67,6 +79,12 @@ class CoreConfig(BaseModel):
     )
     tenderly_olas_funds: float = Field(default=100000.0, description="OLAS amount for vNet funding")
 
+    # Safe Transaction Retry System
+    safe_tx_max_retries: int = Field(default=6, description="Maximum retries for Safe transactions")
+    safe_tx_gas_buffer: float = Field(
+        default=1.5, description="Gas buffer multiplier for Safe transactions"
+    )
+
 
 T = TypeVar("T", bound="StorableModel")
 
@@ -106,16 +124,31 @@ class StorableModel(BaseModel):
         self._path = path
 
     def save_yaml(self, path: Optional[Path] = None) -> None:
-        """Save to YAML file"""
+        """Save to YAML file preserving comments if file exists."""
         if path is None:
             if getattr(self, "_path", None) is None:
                 raise ValueError("Save path not specified and no previous path stored.")
             path = self._path
 
         path = path.with_suffix(".yaml")
+        ryaml = YAML()
+        ryaml.preserve_quotes = True
+        ryaml.indent(mapping=2, sequence=4, offset=2)
+
+        data = self.model_dump(mode="json")
+
+        if path.exists():
+            with path.open("r", encoding="utf-8") as f:
+                try:
+                    target = ryaml.load(f) or {}
+                    _update_yaml_recursive(target, data)
+                    data = target
+                except Exception:
+                    # Fallback to overwrite if load fails
+                    pass
 
         with path.open("w", encoding="utf-8") as f:
-            yaml.safe_dump(self.model_dump(), f, sort_keys=False, allow_unicode=True)
+            ryaml.dump(data, f)
         self._storage_format = "yaml"
         self._path = path
 
@@ -171,8 +204,9 @@ class StorableModel(BaseModel):
     def load_yaml(cls: Type[T], path: str | Path) -> T:
         """Load from YAML file"""
         path = Path(path)
+        ryaml = YAML()
         with path.open("r", encoding="utf-8") as f:
-            data = yaml.safe_load(f)
+            data = ryaml.load(f)
         obj = cls(**data)
         obj._storage_format = "yaml"
         obj._path = path
@@ -223,10 +257,9 @@ class Config(StorableModel):
             return
 
         try:
-            import yaml
-
+            ryaml = YAML()
             with CONFIG_PATH.open("r", encoding="utf-8") as f:
-                data = yaml.safe_load(f) or {}
+                data = ryaml.load(f) or {}
 
             # Load core config
             if "core" in data:
@@ -270,25 +303,37 @@ class Config(StorableModel):
             self.save_config()
 
     def save_config(self) -> None:
-        """Persist current config to config.yaml."""
-        import yaml
-
+        """Persist current config to config.yaml preserving comments."""
         from iwa.core.constants import CONFIG_PATH
 
         data = {}
 
         if self.core:
-            data["core"] = self.core.model_dump()
+            data["core"] = self.core.model_dump(mode="json")
 
         data["plugins"] = {}
         for plugin_name, plugin_config in self.plugins.items():
             if isinstance(plugin_config, BaseModel):
-                data["plugins"][plugin_name] = plugin_config.model_dump()
+                data["plugins"][plugin_name] = plugin_config.model_dump(mode="json")
             elif isinstance(plugin_config, dict):
                 data["plugins"][plugin_name] = plugin_config
 
+        ryaml = YAML()
+        ryaml.preserve_quotes = True
+        ryaml.indent(mapping=2, sequence=4, offset=2)
+
+        if CONFIG_PATH.exists():
+            with CONFIG_PATH.open("r", encoding="utf-8") as f:
+                try:
+                    target = ryaml.load(f) or {}
+                    _update_yaml_recursive(target, data)
+                    data = target
+                except Exception:
+                    # Fallback to overwrite if load fails
+                    pass
+
         with CONFIG_PATH.open("w", encoding="utf-8") as f:
-            yaml.safe_dump(data, f, sort_keys=False, allow_unicode=True, default_flow_style=False)
+            ryaml.dump(data, f)
 
         self._path = CONFIG_PATH
         self._storage_format = "yaml"

iwa/core/pricing.py

@@ -4,9 +4,9 @@ import time
 from datetime import datetime, timedelta
 from typing import Dict, Optional
 
-import requests
 from loguru import logger
 
+from iwa.core.http import create_retry_session
 from iwa.core.secrets import secrets
 
 # Global cache shared across all PriceService instances
@@ -28,6 +28,11 @@ class PriceService:
             if self.secrets.coingecko_api_key
             else None
         )
+        self.session = create_retry_session()
+
+    def close(self):
+        """Close the session."""
+        self.session.close()
 
     def get_token_price(self, token_id: str, vs_currency: str = "eur") -> Optional[float]:
         """Get token price in specified currency.
@@ -66,7 +71,8 @@ class PriceService:
                 if self.api_key:
                     headers["x-cg-demo-api-key"] = self.api_key
 
-                response = requests.get(url, params=params, headers=headers, timeout=10)
+                # Use session instead of direct requests
+                response = self.session.get(url, params=params, headers=headers, timeout=10)
 
                 if response.status_code == 401 and self.api_key:
                     logger.warning("CoinGecko API key invalid (401). Retrying without key...")
@@ -74,7 +80,7 @@ class PriceService:
                     headers.pop("x-cg-demo-api-key", None)
                     # Re-run with base URL
                     url = f"{self.BASE_URL}/simple/price"
-                    response = requests.get(url, params=params, headers=headers, timeout=10)
+                    response = self.session.get(url, params=params, headers=headers, timeout=10)
 
                 if response.status_code == 429:
                     logger.warning(
@@ -93,9 +99,7 @@ class PriceService:
                     return float(data[token_id][vs_currency])
 
                 # If we got response but price not found, it's likely a wrong ID
-                logger.debug(
-                    f"Price for {token_id} in {vs_currency} not found in response: {data}"
-                )
+                logger.debug(f"Price for {token_id} in {vs_currency} not found in response: {data}")
                 return None
 
             except Exception as e:

iwa/core/rpc_monitor.py

@@ -1,4 +1,5 @@
 """RPC Monitor for tracking API usage."""
+
 import threading
 from collections import defaultdict
 from typing import Dict

iwa/core/secrets.py

@@ -72,6 +72,33 @@ class Secrets(BaseSettings):
 
         return self
 
+    @model_validator(mode="after")
+    def strip_quotes_from_secrets(self) -> "Secrets":
+        """Strip leading/trailing quotes from SecretStr fields.
+
+        Docker env_file often preserves quotes (e.g. KEY="val" -> "val"),
+        which causes API authentication failures.
+        """
+        for field_name, field_value in self:
+            if isinstance(field_value, SecretStr):
+                raw_value = field_value.get_secret_value()
+                # Check for matching quotes at start and end
+                if len(raw_value) >= 2 and (
+                    (raw_value.startswith('"') and raw_value.endswith('"'))
+                    or (raw_value.startswith("'") and raw_value.endswith("'"))
+                ):
+                    clean_value = raw_value[1:-1]
+                    setattr(self, field_name, SecretStr(clean_value))
+            elif isinstance(field_value, str):
+                # Also strip quotes from plain string fields (like health_url)
+                if len(field_value) >= 2 and (
+                    (field_value.startswith('"') and field_value.endswith('"'))
+                    or (field_value.startswith("'") and field_value.endswith("'"))
+                ):
+                    clean_value = field_value[1:-1]
+                    setattr(self, field_name, clean_value)
+        return self
+
 
 # Global secrets instance
 secrets = Secrets()

iwa/core/services/account.py

@@ -20,7 +20,7 @@ class AccountService:
         self.key_storage = key_storage
 
     @property
-    def master_account(self) -> Optional[StoredSafeAccount]:
+    def master_account(self) -> Optional[Union["EncryptedAccount", StoredSafeAccount]]:
         """Get master account."""
         return self.key_storage.master_account
 

iwa/core/services/balance.py

@@ -1,9 +1,7 @@
 """Balance service module."""
 
-import time
 from typing import TYPE_CHECKING, Optional, Union
 
-from loguru import logger
 from web3.types import Wei
 
 from iwa.core.chain import ChainInterfaces
@@ -90,24 +88,3 @@ class BalanceService:
 
         contract = ERC20Contract(chain_name=chain_name, address=token_address)
         return contract.balance_of_wei(account.address)
-
-    def get_erc20_balance_with_retry(
-        self,
-        account_address: str,
-        token_address_or_name: str,
-        chain_name: str = "gnosis",
-        retries: int = 3,
-    ) -> Optional[float]:
-        """Fetch balance with retry logic."""
-        for attempt in range(retries):
-            try:
-                return self.get_erc20_balance_eth(
-                    account_address, token_address_or_name, chain_name
-                )
-            except Exception as e:
-                if attempt == retries - 1:
-                    logger.error(
-                        f"Failed to fetch balance for {token_address_or_name} after {retries} attempts: {e}"
-                    )
-                time.sleep(1)
-        return None

iwa/core/services/safe.py

@@ -1,6 +1,6 @@
 """Safe service module."""
 
-from typing import TYPE_CHECKING, List, Optional, Tuple
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 
 from loguru import logger
 from safe_eth.eth import EthereumClient
@@ -35,6 +35,7 @@ class SafeService:
         """Initialize SafeService."""
         self.key_storage = key_storage
         self.account_service = account_service
+        self._client_cache: Dict[str, EthereumClient] = {}
 
     def create_safe(
         self,
@@ -102,7 +103,12 @@ class SafeService:
 
         # Use ChainInterface which has proper RPC rotation and parsing
         chain_interface = ChainInterfaces().get(chain_name)
-        return EthereumClient(chain_interface.current_rpc)
+        rpc_url = chain_interface.current_rpc
+
+        if rpc_url not in self._client_cache:
+            self._client_cache[rpc_url] = EthereumClient(rpc_url)
+
+        return self._client_cache[rpc_url]
 
     def _deploy_safe_contract(
         self,
@@ -222,25 +228,25 @@ class SafeService:
         threshold: int,
         tag: Optional[str],
     ) -> StoredSafeAccount:
-        # Check if already exists
-        accounts = self.key_storage.accounts
-        if contract_address in accounts and isinstance(
-            accounts[contract_address], StoredSafeAccount
-        ):
-            safe_account = accounts[contract_address]
-            if chain_name not in safe_account.chains:
-                safe_account.chains.append(chain_name)
-        else:
-            safe_account = StoredSafeAccount(
-                tag=tag or f"Safe {contract_address[:6]}",
-                address=contract_address,
-                chains=[chain_name],
-                threshold=threshold,
-                signers=owner_addresses,
-            )
-            accounts[contract_address] = safe_account
+        # Check if already exists (by address)
+        existing = self.key_storage.find_stored_account(contract_address)
+        if existing and isinstance(existing, StoredSafeAccount):
+            if chain_name not in existing.chains:
+                existing.chains.append(chain_name)
+            self.key_storage.save()
+            return existing
+
+        # Create new Safe account object
+        safe_account = StoredSafeAccount(
+            tag=tag or f"Safe {contract_address[:6]}",
+            address=contract_address,
+            chains=[chain_name],
+            threshold=threshold,
+            signers=owner_addresses,
+        )
 
-        self.key_storage.save()
+        # Register via centralized method (enforces tag uniqueness)
+        self.key_storage.register_account(safe_account)
         return safe_account
 
     def redeploy_safes(self):
@@ -290,30 +296,51 @@ class SafeService:
 
         return signer_pkeys
 
-    def _sign_and_execute_safe_tx(self, safe_tx: SafeTx, signer_keys: List[str]) -> str:
+    def _sign_and_execute_safe_tx(
+        self,
+        safe_tx: SafeTx,
+        signer_keys: List[str],
+        chain_name: str,
+        safe_address: str,
+    ) -> str:
         """Sign and execute a SafeTx internally (INTERNAL USE ONLY).
 
         This method handles the signing and execution of a Safe transaction,
         keeping private keys internal to SafeService.
 
+        Uses SafeTransactionExecutor for retry logic and gas handling.
+
         SECURITY: Keys are overwritten with zeros and cleared after use.
         """
+        from iwa.core.chain import ChainInterfaces
+        from iwa.core.services.safe_executor import SafeTransactionExecutor
+
         try:
-            # Sign with all available signers
+            # Sign with all available signers (local operation)
             for pk in signer_keys:
-                safe_tx.sign(pk)
+                if pk:
+                    safe_tx.sign(pk)
 
-            # Verify the transaction will succeed
-            safe_tx.call()
+            chain_interface = ChainInterfaces().get(chain_name)
+            executor = SafeTransactionExecutor(chain_interface)
 
-            # Execute using the first signer
-            safe_tx.execute(signer_keys[0])
+            success, tx_hash_or_error, receipt = executor.execute_with_retry(
+                safe_address=safe_address,
+                safe_tx=safe_tx,
+                signer_keys=signer_keys,
+                operation_name=f"safe_tx_{safe_address[:10]}",
+            )
+
+            if success:
+                return tx_hash_or_error
+            else:
+                raise ValueError(f"Safe transaction failed: {tx_hash_or_error}")
 
-            return f"0x{safe_tx.tx_hash.hex()}"
         finally:
             # SECURITY: Overwrite keys with zeros before clearing (best effort)
             for i in range(len(signer_keys)):
-                signer_keys[i] = "0" * len(signer_keys[i]) if signer_keys[i] else ""
+                if signer_keys[i]:
+                    signer_keys[i] = "0" * len(signer_keys[i])
             signer_keys.clear()
 
     def execute_safe_transaction(
@@ -358,17 +385,16 @@ class SafeService:
 
         # Get signer keys, execute, and immediately clear
         signer_keys = self._get_signer_keys(safe_account)
-        try:
-            tx_hash = self._sign_and_execute_safe_tx(safe_tx, signer_keys)
-            logger.info(f"Safe transaction executed. Tx Hash: {tx_hash}")
-            return tx_hash
-        finally:
-            # Clear keys from memory (best effort)
-            for i in range(len(signer_keys)):
-                signer_keys[i] = None
-            del signer_keys
+        tx_hash = self._sign_and_execute_safe_tx(
+            safe_tx=safe_tx,
+            signer_keys=signer_keys,
+            chain_name=chain_name,
+            safe_address=safe_account.address,
+        )
+        logger.info(f"Safe transaction executed. Tx Hash: {tx_hash}")
+        return tx_hash
 
-    def get_sign_and_execute_callback(self, safe_address_or_tag: str):
+    def get_sign_and_execute_callback(self, safe_address_or_tag: str, chain_name: str):
         """Get a callback function that signs and executes a SafeTx.
 
         This method returns a callback that can be passed to SafeMultisig.send_tx().
@@ -376,6 +402,7 @@ class SafeService:
 
         Args:
             safe_address_or_tag: The Safe account address or tag
+            chain_name: The chain name for context
 
         Returns:
             A callable that takes a SafeTx and returns the transaction hash
@@ -387,11 +414,11 @@ class SafeService:
 
         def _sign_and_execute(safe_tx: SafeTx) -> str:
             signer_keys = self._get_signer_keys(safe_account)
-            try:
-                return self._sign_and_execute_safe_tx(safe_tx, signer_keys)
-            finally:
-                for i in range(len(signer_keys)):
-                    signer_keys[i] = None
-                del signer_keys
+            return self._sign_and_execute_safe_tx(
+                safe_tx=safe_tx,
+                signer_keys=signer_keys,
+                chain_name=chain_name,
+                safe_address=safe_account.address,
+            )
 
         return _sign_and_execute