forked from TrueCloudLab/frostfs-testlib
90 lines
2.7 KiB
Python
90 lines
2.7 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from typing import Optional
|
|
|
|
from frostfs_testlib.resources.common import DEFAULT_WALLET_CONFIG, DEFAULT_WALLET_PASS
|
|
from frostfs_testlib.shell import Shell
|
|
from frostfs_testlib.storage.cluster import Cluster, NodeBase
|
|
from frostfs_testlib.utils.wallet_utils import get_last_address_from_wallet, init_wallet
|
|
|
|
logger = logging.getLogger("frostfs.testlib.utils")
|
|
|
|
|
|
@dataclass
|
|
class WalletInfo:
|
|
path: str
|
|
password: str = DEFAULT_WALLET_PASS
|
|
config_path: str = DEFAULT_WALLET_CONFIG
|
|
|
|
@staticmethod
|
|
def from_node(node: NodeBase):
|
|
return WalletInfo(
|
|
node.get_wallet_path(), node.get_wallet_password(), node.get_wallet_config_path()
|
|
)
|
|
|
|
def get_address(self) -> str:
|
|
"""
|
|
Extracts the last address from wallet via neo3 lib.
|
|
|
|
Returns:
|
|
The address of the wallet.
|
|
"""
|
|
return get_last_address_from_wallet(self.path, self.password)
|
|
|
|
def get_address_from_json(self, account_id: int = 0) -> str:
|
|
"""
|
|
Extracts address of the given account id from wallet using json lookup.
|
|
(Useful if neo3 fails for some reason and can't be used).
|
|
|
|
Args:
|
|
account_id: id of the account to get address.
|
|
|
|
Returns:
|
|
address string.
|
|
"""
|
|
with open(self.path, "r") as wallet:
|
|
wallet_json = json.load(wallet)
|
|
assert abs(account_id) + 1 <= len(
|
|
wallet_json["accounts"]
|
|
), f"There is no index '{account_id}' in wallet: {wallet_json}"
|
|
|
|
return wallet_json["accounts"][account_id]["address"]
|
|
|
|
|
|
class WalletFactory:
|
|
def __init__(self, wallets_dir: str, shell: Shell, cluster: Cluster) -> None:
|
|
self.shell = shell
|
|
self.wallets_dir = wallets_dir
|
|
self.cluster = cluster
|
|
|
|
def create_wallet(
|
|
self, file_name: Optional[str] = None, password: Optional[str] = None
|
|
) -> WalletInfo:
|
|
"""
|
|
Creates new default wallet.
|
|
|
|
Args:
|
|
file_name: output wallet file name.
|
|
password: wallet password.
|
|
|
|
Returns:
|
|
WalletInfo object of new wallet.
|
|
"""
|
|
|
|
if file_name is None:
|
|
file_name = str(uuid.uuid4())
|
|
if password is None:
|
|
password = ""
|
|
|
|
base_path = os.path.join(self.wallets_dir, file_name)
|
|
wallet_path = f"{base_path}.json"
|
|
wallet_config_path = f"{base_path}.yaml"
|
|
init_wallet(wallet_path, password)
|
|
|
|
with open(wallet_config_path, "w") as config_file:
|
|
config_file.write(f'password: "{password}"')
|
|
|
|
return WalletInfo(wallet_path, password, wallet_config_path)
|