From 72c319e926cb9018d3da9bb3b8fc7bc59fdd4cc2 Mon Sep 17 00:00:00 2001 From: Ryan Gilbert Date: Tue, 17 Dec 2024 11:40:10 -0500 Subject: [PATCH] chore: add network_id to WalletData --- CHANGELOG.md | 2 ++ README.md | 2 +- cdp/wallet.py | 3 ++- cdp/wallet_data.py | 27 +++++++++++++++++++++++---- tests/test_wallet.py | 19 +++++++++++++++++++ 5 files changed, 47 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b3e9c91..bac05be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## Unreleased +- Add `network_id` to `WalletData` so that it is saved with the seed data and surfaced via the export function + ### [0.12.1] - 2024-12-10 ### Added diff --git a/README.md b/README.md index 00883c6..3514725 100644 --- a/README.md +++ b/README.md @@ -189,7 +189,7 @@ list(address.trades()) The SDK creates wallets with [Developer-Managed (1-1)](https://docs.cdp.coinbase.com/mpc-wallet/docs/wallets#developer-managed-wallets) keys by default, which means you are responsible for securely storing the keys required to re-instantiate wallets. The below code walks you through how to export a wallet and store it in a secure location. ```python -# Export the data required to re-instantiate the wallet. The data contains the seed and the ID of the wallet. +# Export the data required to re-instantiate the wallet. The data contains the seed, the ID of the wallet, and the network ID. data = wallet.export_data() ``` diff --git a/cdp/wallet.py b/cdp/wallet.py index 46958e0..9405770 100644 --- a/cdp/wallet.py +++ b/cdp/wallet.py @@ -492,7 +492,7 @@ def export_data(self) -> WalletData: if self._master is None or self._seed is None: raise ValueError("Wallet does not have seed loaded") - return WalletData(self.id, self._seed) + return WalletData(self.id, self._seed, self.network_id) def save_seed(self, file_path: str, encrypt: bool | None = False) -> None: """Save the wallet seed to a file. @@ -530,6 +530,7 @@ def save_seed(self, file_path: str, encrypt: bool | None = False) -> None: "encrypted": encrypt, "auth_tag": auth_tag, "iv": iv, + "network_id": self.network_id, } with open(file_path, "w") as f: diff --git a/cdp/wallet_data.py b/cdp/wallet_data.py index dc4816e..65649da 100644 --- a/cdp/wallet_data.py +++ b/cdp/wallet_data.py @@ -1,16 +1,18 @@ class WalletData: """A class representing wallet data required to recreate a wallet.""" - def __init__(self, wallet_id: str, seed: str) -> None: + def __init__(self, wallet_id: str, seed: str, network_id: str | None = None) -> None: """Initialize the WalletData class. Args: wallet_id (str): The ID of the wallet. seed (str): The seed of the wallet. + network_id (str | None): The network ID of the wallet. Defaults to None. """ self._wallet_id = wallet_id self._seed = seed + self._network_id = network_id @property def wallet_id(self) -> str: @@ -32,6 +34,16 @@ def seed(self) -> str: """ return self._seed + @property + def network_id(self) -> str | None: + """Get the network ID of the wallet. + + Returns: + str: The network ID of the wallet. + + """ + return self._network_id + def to_dict(self) -> dict[str, str]: """Convert the wallet data to a dictionary. @@ -39,7 +51,10 @@ def to_dict(self) -> dict[str, str]: dict[str, str]: The dictionary representation of the wallet data. """ - return {"wallet_id": self.wallet_id, "seed": self.seed} + result = {"wallet_id": self.wallet_id, "seed": self.seed} + if self._network_id is not None: + result["network_id"] = self.network_id + return result @classmethod def from_dict(cls, data: dict[str, str]) -> "WalletData": @@ -52,7 +67,11 @@ def from_dict(cls, data: dict[str, str]) -> "WalletData": WalletData: The wallet data. """ - return cls(data["wallet_id"], data["seed"]) + return cls( + data["wallet_id"], + data["seed"], + data.get("network_id") + ) def __str__(self) -> str: """Return a string representation of the WalletData object. @@ -61,7 +80,7 @@ def __str__(self) -> str: str: A string representation of the wallet data. """ - return f"WalletData: (wallet_id: {self.wallet_id}, seed: {self.seed})" + return f"WalletData: (wallet_id: {self.wallet_id}, seed: {self.seed}, network_id: {self.network_id})" def __repr__(self) -> str: """Return a string representation of the WalletData object. diff --git a/tests/test_wallet.py b/tests/test_wallet.py index 751d56b..6ed56af 100644 --- a/tests/test_wallet.py +++ b/tests/test_wallet.py @@ -658,3 +658,22 @@ def test_wallet_quote_fund_no_default_address(wallet_factory): with pytest.raises(ValueError, match="Default address does not exist"): wallet.quote_fund(amount="1.0", asset_id="eth") + +@patch("cdp.Cdp.use_server_signer", False) +@patch("cdp.wallet.os") +@patch("cdp.wallet.Bip32Slip10Secp256k1") +def test_wallet_export_data(mock_bip32, mock_os, wallet_factory, master_key_factory): + """Test Wallet export_data method.""" + seed = b"\x00" * 64 + mock_urandom = Mock(return_value=seed) + mock_os.urandom = mock_urandom + mock_from_seed = Mock(return_value=master_key_factory(seed)) + mock_bip32.FromSeed = mock_from_seed + + wallet = wallet_factory() + + exported = wallet.export_data() + + assert exported.wallet_id == wallet.id + assert exported.seed == seed.hex() + assert exported.network_id == wallet.network_id