forked from TrueCloudLab/frostfs-testlib
91 lines
2.9 KiB
Python
91 lines
2.9 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import Optional
|
|
|
|
import yaml
|
|
|
|
from frostfs_testlib import reporter
|
|
from frostfs_testlib.resources.common import ASSETS_DIR, DEFAULT_WALLET_CONFIG, DEFAULT_WALLET_PASS
|
|
from frostfs_testlib.shell import Shell
|
|
from frostfs_testlib.storage.cluster import 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):
|
|
wallet_path = node.get_wallet_path()
|
|
wallet_password = node.get_wallet_password()
|
|
wallet_config_file = os.path.join(ASSETS_DIR, os.path.basename(node.get_wallet_config_path()))
|
|
with open(wallet_config_file, "w") as file:
|
|
file.write(yaml.dump({"wallet": wallet_path, "password": wallet_password}))
|
|
|
|
return WalletInfo(wallet_path, wallet_password, wallet_config_file)
|
|
|
|
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) -> None:
|
|
self.shell = shell
|
|
self.wallets_dir = wallets_dir
|
|
|
|
def create_wallet(self, file_name: str, 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 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'wallet: {wallet_path}\npassword: "{password}"')
|
|
|
|
reporter.attach(wallet_path, os.path.basename(wallet_path))
|
|
|
|
return WalletInfo(wallet_path, password, wallet_config_path)
|