[CHIA-307] (long-lived) Vault

chia/_tests/clvm/test_p2_singletons.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+import pytest
+from chia_rs import G2Element
+
+from chia._tests.util.spend_sim import CostLogger, sim_and_client
+from chia.types.blockchain_format.program import Program
+from chia.types.blockchain_format.sized_bytes import bytes32
+from chia.types.coin_spend import make_spend
+from chia.types.mempool_inclusion_status import MempoolInclusionStatus
+from chia.types.spend_bundle import SpendBundle
+from chia.util.errors import Err
+from chia.util.ints import uint64
+from chia.wallet.conditions import CreateCoin
+from chia.wallet.puzzles import p2_singleton_via_delegated_puzzle_safe as dp_safe
+from chia.wallet.util.curry_and_treehash import shatree_atom
+
+ACS = Program.to(1)
+ACS_HASH = ACS.get_tree_hash()
+
+MOCK_SINGLETON_MOD = Program.to([2, 5, 7])  # (a 5 11) - (mod (_ PUZZLE . solution) (a PUZZLE solution))
+MOCK_SINGLETON_MOD_HASH = MOCK_SINGLETON_MOD.get_tree_hash()
+MOCK_SINGLETON_LAUNCHER_ID = bytes32([0] * 32)
+MOCK_SINGLETON_LAUNCHER_HASH = bytes32([1] * 32)
+MOCK_SINGLETON = MOCK_SINGLETON_MOD.curry(
+    Program.to((MOCK_SINGLETON_MOD_HASH, (MOCK_SINGLETON_LAUNCHER_ID, MOCK_SINGLETON_LAUNCHER_HASH))),
+    ACS,
+)
+MOCK_SINGLETON_HASH = MOCK_SINGLETON.get_tree_hash()
+dp_safe.PRE_HASHED_HASHES[MOCK_SINGLETON_MOD_HASH] = shatree_atom(MOCK_SINGLETON_MOD_HASH)
+dp_safe.PRE_HASHED_HASHES[MOCK_SINGLETON_LAUNCHER_HASH] = shatree_atom(MOCK_SINGLETON_LAUNCHER_HASH)
+
+
+@pytest.mark.anyio
+async def test_dp_safe_lifecycle(cost_logger: CostLogger) -> None:
+    P2_SINGLETON = dp_safe.construct(MOCK_SINGLETON_LAUNCHER_ID, MOCK_SINGLETON_MOD_HASH, MOCK_SINGLETON_LAUNCHER_HASH)
+    P2_SINGLETON_HASH = dp_safe.construct_hash(
+        MOCK_SINGLETON_LAUNCHER_ID, MOCK_SINGLETON_MOD_HASH, MOCK_SINGLETON_LAUNCHER_HASH
+    )
+    assert dp_safe.match(P2_SINGLETON) is not None
+    assert dp_safe.match(ACS) is None
+    assert dp_safe.match(MOCK_SINGLETON) is None
+
+    async with sim_and_client() as (sim, sim_client):
+        await sim.farm_block(P2_SINGLETON_HASH)
+        await sim.farm_block(MOCK_SINGLETON_HASH)
+        p2_singleton = (await sim_client.get_coin_records_by_puzzle_hash(P2_SINGLETON_HASH, include_spent_coins=False))[
+            0
+        ].coin
+        singleton = (await sim_client.get_coin_records_by_puzzle_hash(MOCK_SINGLETON_HASH, include_spent_coins=False))[
+            0
+        ].coin
+
+        dp = Program.to(1)
+        dp_hash = dp.get_tree_hash()
+        bundle = cost_logger.add_cost(
+            "p2_singleton_w_mock_singleton",
+            SpendBundle(
+                [
+                    make_spend(
+                        p2_singleton,
+                        P2_SINGLETON,
+                        dp_safe.solve(
+                            ACS_HASH,
+                            dp,
+                            Program.to([CreateCoin(bytes32([0] * 32), uint64(0)).to_program()]),
+                            p2_singleton.name(),
+                        ),
+                    ),
+                    make_spend(
+                        singleton,
+                        MOCK_SINGLETON,
+                        Program.to([dp_safe.required_announcement(dp_hash, p2_singleton.name()).to_program()]),
+                    ),
+                ],
+                G2Element(),
+            ),
+        )
+        result = await sim_client.push_tx(bundle)
+        assert result == (MempoolInclusionStatus.SUCCESS, None)
+        checkpoint = sim.block_height
+        await sim.farm_block()
+
+        assert len(await sim_client.get_coin_records_by_puzzle_hash(bytes32([0] * 32), include_spent_coins=False)) == 1
+
+        await sim.rewind(checkpoint)
+
+        result = await sim_client.push_tx(
+            SpendBundle(
+                [
+                    make_spend(
+                        p2_singleton,
+                        P2_SINGLETON,
+                        dp_safe.solve(
+                            ACS_HASH,
+                            dp,
+                            Program.to([CreateCoin(bytes32([0] * 32), uint64(0)).to_program()]),
+                            bytes32([0] * 32),
+                        ),
+                    ),
+                    make_spend(
+                        singleton,
+                        MOCK_SINGLETON,
+                        Program.to([dp_safe.required_announcement(dp_hash, p2_singleton.name()).to_program()]),
+                    ),
+                ],
+                G2Element(),
+            )
+        )
+        assert result == (MempoolInclusionStatus.FAILED, Err.ASSERT_MY_COIN_ID_FAILED)
+
+        result = await sim_client.push_tx(
+            SpendBundle(
+                [
+                    make_spend(
+                        p2_singleton,
+                        P2_SINGLETON,
+                        dp_safe.solve(
+                            ACS_HASH,
+                            dp,
+                            Program.to([CreateCoin(bytes32([0] * 32), uint64(0)).to_program()]),
+                            p2_singleton.name(),
+                        ),
+                    ),
+                    make_spend(
+                        singleton,
+                        MOCK_SINGLETON,
+                        Program.to([]),
+                    ),
+                ],
+                G2Element(),
+            )
+        )
+        assert result == (MempoolInclusionStatus.FAILED, Err.ASSERT_ANNOUNCE_CONSUMED_FAILED)
+
+        result = await sim_client.push_tx(
+            SpendBundle(
+                [
+                    make_spend(
+                        p2_singleton,
+                        P2_SINGLETON,
+                        dp_safe.solve(
+                            ACS_HASH,
+                            dp,
+                            Program.to([CreateCoin(bytes32([0] * 32), uint64(0)).to_program()]),
+                            p2_singleton.name(),
+                        ),
+                    ),
+                    make_spend(
+                        singleton,
+                        MOCK_SINGLETON,
+                        Program.to([dp_safe.required_announcement(dp_hash, bytes32([0] * 32)).to_program()]),
+                    ),
+                ],
+                G2Element(),
+            )
+        )
+        assert result == (MempoolInclusionStatus.FAILED, Err.ASSERT_ANNOUNCE_CONSUMED_FAILED)

chia/_tests/cmds/wallet/test_vault.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from chia_rs import G1Element
+
+from chia._tests.cmds.cmd_test_utils import TestRpcClients, TestWalletRpcClient, run_cli_command_and_assert
+from chia._tests.cmds.wallet.test_consts import FINGERPRINT_ARG, STD_TX, STD_UTX, WALLET_ID_ARG
+from chia.rpc.wallet_request_types import VaultCreate, VaultCreateResponse, VaultRecovery, VaultRecoveryResponse
+from chia.util.ints import uint64
+from chia.wallet.conditions import ConditionValidTimes
+from chia.wallet.util.tx_config import TXConfig
+
+test_condition_valid_times: ConditionValidTimes = ConditionValidTimes(min_time=uint64(100), max_time=uint64(150))
+
+
+def test_vault_create(capsys: object, get_test_cli_clients: tuple[TestRpcClients, Path]) -> None:
+    test_rpc_clients, root_dir = get_test_cli_clients
+
+    # set RPC clients
+    class CreateVaultRpcClient(TestWalletRpcClient):
+        async def vault_create(
+            self,
+            args: VaultCreate,
+            tx_config: TXConfig,
+            timelock_info: ConditionValidTimes,
+        ) -> VaultCreateResponse:
+            return VaultCreateResponse([STD_UTX], [STD_TX])
+
+    inst_rpc_client = CreateVaultRpcClient()  # pylint: disable=no-value-for-parameter
+    test_rpc_clients.wallet_rpc_client = inst_rpc_client
+    pk = bytes(G1Element()).hex()
+    recovery_pk = bytes(G1Element()).hex()
+    timelock = "100"
+    hidden_puzzle_index = "10"
+    fee = "0.1"
+    command_args = [
+        "vault",
+        "create",
+        FINGERPRINT_ARG,
+        "-pk",
+        pk,
+        "-rk",
+        recovery_pk,
+        "-rt",
+        timelock,
+        "-i",
+        hidden_puzzle_index,
+        "-m",
+        fee,
+        "--valid-at",
+        "100",
+        "--expires-at",
+        "150",
+    ]
+    assert_list = ["Successfully created a Vault wallet"]
+    run_cli_command_and_assert(capsys, root_dir, command_args, assert_list)
+
+
+def test_vault_recovery(capsys: object, get_test_cli_clients: tuple[TestRpcClients, Path], tmp_path: Path) -> None:
+    test_rpc_clients, root_dir = get_test_cli_clients
+
+    # set RPC clients
+    class CreateVaultRpcClient(TestWalletRpcClient):
+        async def vault_recovery(
+            self,
+            args: VaultRecovery,
+            tx_config: TXConfig,
+            timelock_info: ConditionValidTimes,
+        ) -> VaultRecoveryResponse:
+            return VaultRecoveryResponse([STD_UTX, STD_UTX], [STD_TX, STD_TX], STD_TX.name, STD_TX.name)
+
+    inst_rpc_client = CreateVaultRpcClient()  # pylint: disable=no-value-for-parameter
+    test_rpc_clients.wallet_rpc_client = inst_rpc_client
+    pk = bytes(G1Element()).hex()
+    recovery_pk = bytes(G1Element()).hex()
+    timelock = "100"
+    hidden_puzzle_index = "10"
+    command_args = [
+        "vault",
+        "recover",
+        "-pk",
+        pk,
+        "-rk",
+        recovery_pk,
+        "-rt",
+        timelock,
+        "-i",
+        hidden_puzzle_index,
+        "-ri",
+        str(tmp_path / "recovery_init.json"),
+        "-rf",
+        str(tmp_path / "recovery_finish.json"),
+        "--valid-at",
+        "100",
+        "--expires-at",
+        "150",
+    ]
+    assert_list = [
+        "Writing transactions to file ",
+        "recovery_init.json",
+        "recovery_finish.json",
+    ]
+    run_cli_command_and_assert(capsys, root_dir, command_args + [FINGERPRINT_ARG, WALLET_ID_ARG], assert_list)

chia/_tests/core/daemon/test_daemon.py

@@ -10,7 +10,7 @@
 import pytest
 from aiohttp import WSMessage
 from aiohttp.web_ws import WebSocketResponse
-from chia_rs import G1Element
+from chia_rs import G1Element, PrivateKey
 from pytest_mock import MockerFixture
 
 from chia._tests.util.misc import Marks, datacases
@@ -32,7 +32,7 @@
 from chia.simulator.setup_services import setup_full_node
 from chia.util.config import load_config
 from chia.util.json_util import dict_to_json_str
-from chia.util.keychain import Keychain, KeyData, supports_os_passphrase_storage
+from chia.util.keychain import Keychain, KeyData, KeyTypes, supports_os_passphrase_storage
 from chia.util.keyring_wrapper import DEFAULT_PASSPHRASE_IF_NO_MASTER_PASSPHRASE, KeyringWrapper
 from chia.util.ws_message import create_payload, create_payload_dict
 from chia.wallet.derive_keys import master_sk_to_farmer_sk, master_sk_to_pool_sk
@@ -205,6 +205,8 @@ async def get_keys_for_plotting(self, request: dict[str, Any]) -> dict[str, Any]
     "hammer stable page grunt venture purse canyon discover "
     "egg vivid spare immune awake code announce message"
 )
+assert isinstance(test_key_data.private_key, PrivateKey)
+assert isinstance(test_key_data_2.private_key, PrivateKey)
 
 success_response_data = {
     "success": True,
@@ -237,6 +239,7 @@ async def get_keys_for_plotting(self, request: dict[str, Any]) -> dict[str, Any]
 def add_private_key_response_data(fingerprint: int) -> dict[str, object]:
     return {
         "success": True,
+        "key_type": KeyTypes.G1_ELEMENT.value,
         "fingerprint": fingerprint,
     }
 

chia/_tests/core/daemon/test_keychain_proxy.py

@@ -6,6 +6,7 @@
 from typing import Any
 
 import pytest
+from chia_rs import G1Element
 
 from chia.daemon.keychain_proxy import KeychainProxy, connect_to_keychain_and_validate
 from chia.simulator.block_tools import BlockTools
@@ -48,21 +49,22 @@ async def test_add_private_key(keychain_proxy: KeychainProxy) -> None:
 @pytest.mark.anyio
 async def test_add_public_key(keychain_proxy: KeychainProxy) -> None:
     keychain = keychain_proxy
-    await keychain.add_key(bytes(TEST_KEY_3.public_key).hex(), TEST_KEY_3.label, private=False)
+    assert isinstance(TEST_KEY_3.observation_root, G1Element)
+    await keychain.add_key(TEST_KEY_3.public_key.hex(), TEST_KEY_3.label, private=False)
     with pytest.raises(Exception, match="already exists"):
-        await keychain.add_key(bytes(TEST_KEY_3.public_key).hex(), "", private=False)
+        await keychain.add_key(TEST_KEY_3.public_key.hex(), "", private=False)
     key = await keychain.get_key(TEST_KEY_3.fingerprint, include_secrets=False)
     assert key is not None
-    assert key.public_key == TEST_KEY_3.public_key
+    assert key.observation_root == TEST_KEY_3.observation_root
     assert key.secrets is None
 
     pk = await keychain.get_key_for_fingerprint(TEST_KEY_3.fingerprint, private=False)
     assert pk is not None
-    assert pk == TEST_KEY_3.public_key
+    assert pk == TEST_KEY_3.observation_root
 
     pk = await keychain.get_key_for_fingerprint(None, private=False)
     assert pk is not None
-    assert pk == TEST_KEY_3.public_key
+    assert pk == TEST_KEY_3.observation_root
 
     with pytest.raises(KeychainKeyNotFound):
         pk = await keychain.get_key_for_fingerprint(1234567890, private=False)
@@ -83,8 +85,8 @@ async def test_get_key_for_fingerprint(keychain_proxy: KeychainProxy) -> None:
     with pytest.raises(KeychainIsEmpty):
         await keychain.get_key_for_fingerprint(None, private=False)
     await keychain_proxy.add_key(TEST_KEY_1.mnemonic_str(), TEST_KEY_1.label)
-    assert await keychain.get_key_for_fingerprint(TEST_KEY_1.fingerprint, private=False) == TEST_KEY_1.public_key
-    assert await keychain.get_key_for_fingerprint(None, private=False) == TEST_KEY_1.public_key
+    assert await keychain.get_key_for_fingerprint(TEST_KEY_1.fingerprint, private=False) == TEST_KEY_1.observation_root
+    assert await keychain.get_key_for_fingerprint(None, private=False) == TEST_KEY_1.observation_root
     with pytest.raises(KeychainKeyNotFound):
         await keychain.get_key_for_fingerprint(1234567890, private=False)
 
@@ -99,3 +101,15 @@ async def test_get_keys(keychain_proxy_with_keys: KeychainProxy, include_secrets
     else:
         expected_keys = [replace(TEST_KEY_1, secrets=None), replace(TEST_KEY_2, secrets=None)]
     assert keys == expected_keys
+
+
+@pytest.mark.anyio
+async def test_get_first_private_key(keychain_proxy_with_keys: KeychainProxy) -> None:
+    assert TEST_KEY_1.private_key == await keychain_proxy_with_keys.get_first_private_key()
+
+
+@pytest.mark.anyio
+async def test_get_all_private_keys(keychain_proxy_with_keys: KeychainProxy) -> None:
+    assert [TEST_KEY_1.private_key, TEST_KEY_2.private_key] == [
+        k for k, e in await keychain_proxy_with_keys.get_all_private_keys()
+    ]

chia/_tests/core/data_layer/test_data_rpc.py

@@ -70,8 +70,8 @@
 from chia.util.timing import adjusted_timeout, backoff_times
 from chia.wallet.trading.offer import Offer as TradingOffer
 from chia.wallet.transaction_record import TransactionRecord
-from chia.wallet.wallet import Wallet
 from chia.wallet.wallet_node import WalletNode
+from chia.wallet.wallet_protocol import MainWalletProtocol
 
 pytestmark = pytest.mark.data_layer
 nodes = tuple[WalletNode, FullNodeSimulator]
@@ -825,7 +825,7 @@ async def offer_setup_fixture(
     [full_node_service], wallet_services, bt = two_wallet_nodes_services
     enable_batch_autoinsertion_settings = getattr(request, "param", (True, True))
     full_node_api = full_node_service._api
-    wallets: list[Wallet] = []
+    wallets: list[MainWalletProtocol] = []
     for wallet_service in wallet_services:
         wallet_node = wallet_service._node
         assert wallet_node.server is not None
@@ -2373,8 +2373,8 @@ async def test_wallet_log_in_changes_active_fingerprint(
 
     mnemonic = create_mnemonic()
     assert wallet_rpc_api.service.local_keychain is not None
-    private_key = wallet_rpc_api.service.local_keychain.add_key(mnemonic_or_pk=mnemonic)
-    secondary_fingerprint: int = private_key.get_g1().get_fingerprint()
+    private_key, _ = wallet_rpc_api.service.local_keychain.add_key(mnemonic_or_pk=mnemonic)
+    secondary_fingerprint: int = private_key.public_key().get_fingerprint()
 
     await wallet_rpc_api.log_in(request={"fingerprint": primary_fingerprint})