diff --git a/neo/Prompt/Commands/Wallet.py b/neo/Prompt/Commands/Wallet.py index 55639490b..ebb8e109e 100644 --- a/neo/Prompt/Commands/Wallet.py +++ b/neo/Prompt/Commands/Wallet.py @@ -56,7 +56,7 @@ def execute(self, arguments): if not item: print("Wallet %s " % json.dumps(wallet.ToJson(), indent=4)) - return + return wallet try: return self.execute_sub_command(item, arguments[1:]) @@ -75,45 +75,43 @@ def execute(self, arguments): PromptData.close_wallet() path = get_arg(arguments, 0) - if path: - if os.path.exists(path): - print("File already exists") - return - - passwd1 = prompt("[password]> ", is_password=True) - passwd2 = prompt("[password again]> ", is_password=True) - - if passwd1 != passwd2 or len(passwd1) < 10: - print("Please provide matching passwords that are at least 10 characters long") - return - - password_key = to_aes_key(passwd1) - - try: - PromptData.Wallet = UserWallet.Create(path=path, password=password_key) - contract = PromptData.Wallet.GetDefaultContract() - key = PromptData.Wallet.GetKey(contract.PublicKeyHash) - print("Wallet %s" % json.dumps(PromptData.Wallet.ToJson(), indent=4)) - print("Pubkey %s" % key.PublicKey.encode_point(True)) - return PromptData.Wallet - except Exception as e: - print("Exception creating wallet: %s" % e) - PromptData.Wallet = None - if os.path.isfile(path): - try: - os.remove(path) - except Exception as e: - print("Could not remove {}: {}".format(path, e)) - return - - if PromptData.Wallet: - PromptData.Prompt.start_wallet_loop() - return - - else: + if not path: print("Please specify a path") return + if os.path.exists(path): + print("File already exists") + return + + passwd1 = prompt("[password]> ", is_password=True) + passwd2 = prompt("[password again]> ", is_password=True) + + if passwd1 != passwd2 or len(passwd1) < 10: + print("Please provide matching passwords that are at least 10 characters long") + return + + password_key = to_aes_key(passwd1) + + try: + PromptData.Wallet = UserWallet.Create(path=path, password=password_key) + contract = PromptData.Wallet.GetDefaultContract() + key = PromptData.Wallet.GetKey(contract.PublicKeyHash) + print("Wallet %s" % json.dumps(PromptData.Wallet.ToJson(), indent=4)) + print("Pubkey %s" % key.PublicKey.encode_point(True)) + except Exception as e: + print("Exception creating wallet: %s" % e) + PromptData.Wallet = None + if os.path.isfile(path): + try: + os.remove(path) + except Exception as e: + print("Could not remove {}: {}".format(path, e)) + return + + if PromptData.Wallet: + PromptData.Prompt.start_wallet_loop() + return PromptData.Wallet + def command_desc(self): p1 = ParameterDesc('path', 'path to store the wallet file') return CommandDesc('create', 'create a new NEO wallet (with 1 address)', [p1]) @@ -130,28 +128,25 @@ def execute(self, arguments): path = get_arg(arguments, 0) - if path: - - if not os.path.exists(path): - print("Wallet file not found") - return + if not path: + print("Please specify a path") + return - passwd = prompt("[password]> ", is_password=True) - password_key = to_aes_key(passwd) + if not os.path.exists(path): + print("Wallet file not found") + return - try: - PromptData.Wallet = UserWallet.Open(path, password_key) + passwd = prompt("[password]> ", is_password=True) + password_key = to_aes_key(passwd) - PromptData.Prompt.start_wallet_loop() - print("Opened wallet at %s" % path) - return PromptData.Wallet - except Exception as e: - print("Could not open wallet: %s" % e) - return + try: + PromptData.Wallet = UserWallet.Open(path, password_key) - else: - print("Please specify a path") - return + PromptData.Prompt.start_wallet_loop() + print("Opened wallet at %s" % path) + return PromptData.Wallet + except Exception as e: + print("Could not open wallet: %s" % e) def command_desc(self): p1 = ParameterDesc('path', 'path to open the wallet file') @@ -175,7 +170,7 @@ class CommandWalletVerbose(CommandBase): def __init__(self): super().__init__() - def execute(self, arguments): + def execute(self, arguments=None): print("Wallet %s " % json.dumps(PromptData.Wallet.ToJson(verbose=True), indent=4)) return True @@ -188,12 +183,9 @@ class CommandWalletMigrate(CommandBase): def __init__(self): super().__init__() - def execute(self, arguments): - if PromptData.Wallet is not None: - PromptData.Wallet.Migrate() - print("Migrated wallet") - return True - return False + def execute(self, arguments=None): + PromptData.Wallet.Migrate() + return True def command_desc(self): return CommandDesc('migrate', 'migrate an old wallet to the new format') @@ -206,6 +198,11 @@ def __init__(self): def execute(self, arguments): addresses_to_create = get_arg(arguments, 0) + + if not addresses_to_create: + print("Please specify a number of addresses to create.") + return + return CreateAddress(PromptData.Wallet, addresses_to_create) def command_desc(self): diff --git a/neo/Prompt/Commands/tests/test_wallet_commands.py b/neo/Prompt/Commands/tests/test_wallet_commands.py index 10068ec79..54e209697 100644 --- a/neo/Prompt/Commands/tests/test_wallet_commands.py +++ b/neo/Prompt/Commands/tests/test_wallet_commands.py @@ -4,7 +4,10 @@ from neo.Core.Blockchain import Blockchain from neocore.UInt160 import UInt160 from neocore.Fixed8 import Fixed8 -from neo.Prompt.Commands.Wallet import CommandWallet, CreateAddress, DeleteAddress, ImportToken, ImportWatchAddr, ShowUnspentCoins, SplitUnspentCoin +from neo.Prompt.Commands.Wallet import CommandWallet +from neo.Prompt.Commands.Wallet import CreateAddress, DeleteAddress, ImportToken, ImportWatchAddr, ShowUnspentCoins, SplitUnspentCoin +from neo.Prompt.PromptData import PromptData +import os import shutil from mock import patch @@ -34,8 +37,109 @@ def GetWallet1(cls, recreate=False): to_aes_key(UserWalletTestCase.wallet_1_pass())) return cls._wallet1 + @classmethod + def OpenWallet(cls): + PromptData.Wallet = cls.GetWallet1(recreate=True) + + @classmethod + def tearDown(cls): + PromptData.Wallet = None + # Beginning with refactored tests + def test_wallet(self): + # without wallet opened + res = CommandWallet().execute(None) + self.assertFalse(res) + + # with wallet opened + self.OpenWallet() + res = CommandWallet().execute(None) + self.assertEqual(type(res), UserWallet) + + def test_wallet_wrong_command(self): + self.OpenWallet() + args = ['badcommand'] + res = CommandWallet().execute(args) + self.assertFalse(res) + + def test_wallet_create(self): + def remove_new_wallet(): + path = UserWalletTestCase.new_wallet_dest() + try: + if os.path.exists(path): + os.remove(path) + except Exception as e: + print("couldn't remove wallets %s " % e) + + with patch('neo.Prompt.PromptData.PromptData.Prompt'): + with patch('neo.Prompt.Commands.Wallet.prompt', side_effect=["testpassword", "testpassword"]): + # test wallet create successful + path = UserWalletTestCase.new_wallet_dest() + args = ['create', path] + self.assertFalse(os.path.isfile(path)) + res = CommandWallet().execute(args) + self.assertEqual(type(res), UserWallet) + self.assertTrue(os.path.isfile(path)) + remove_new_wallet() + + # test wallet create with no path + args = ['create'] + res = CommandWallet().execute(args) + self.assertFalse(res) + + # test wallet open with already existing path + with patch('neo.Prompt.Commands.Wallet.prompt', side_effect=["testpassword", "testpassword"]): + path = UserWalletTestCase.new_wallet_dest() + args = ['create', path] + self.assertFalse(os.path.isfile(path)) + res = CommandWallet().execute(args) + self.assertEqual(type(res), UserWallet) + self.assertTrue(os.path.isfile(path)) + + res = CommandWallet().execute(args) + self.assertFalse(res) + self.assertTrue(os.path.isfile(path)) + remove_new_wallet() + + # test wallet with different passwords + with patch('neo.Prompt.Commands.Wallet.prompt', side_effect=["testpassword", "bad"]): + path = UserWalletTestCase.new_wallet_dest() + args = ['create', path] + self.assertFalse(os.path.isfile(path)) + res = CommandWallet().execute(args) + self.assertFalse(res) + self.assertFalse(os.path.isfile(path)) + + # test wallet create unsuccessful + with patch('neo.Prompt.Commands.Wallet.prompt', side_effect=["testpassword", "testpassword"]): + with patch('neo.Implementations.Wallets.peewee.UserWallet.UserWallet.Create', side_effect=[Exception('test exception')]): + path = UserWalletTestCase.new_wallet_dest() + args = ['create', path] + res = CommandWallet().execute(args) + self.assertFalse(res) + self.assertFalse(os.path.isfile(path)) + + # test wallet create exception after creation + with patch('neo.Prompt.Commands.Wallet.prompt', side_effect=["testpassword", "testpassword"]): + with patch('neo.Wallets.Wallet.Wallet.GetKey', side_effect=[Exception('test exception')]): + path = UserWalletTestCase.new_wallet_dest() + args = ['create', path] + res = CommandWallet().execute(args) + self.assertFalse(res) + self.assertFalse(os.path.isfile(path)) + + # test wallet create exception after creation with file deletion failure + with patch('neo.Prompt.Commands.Wallet.prompt', side_effect=["testpassword", "testpassword"]): + with patch('neo.Wallets.Wallet.Wallet.GetKey', side_effect=[Exception('test exception')]): + with patch('os.remove', side_effect=[Exception('test exception')]): + path = UserWalletTestCase.new_wallet_dest() + args = ['create', path] + res = CommandWallet().execute(args) + self.assertFalse(res) + self.assertTrue(os.path.isfile(path)) + remove_new_wallet() + def test_wallet_open(self): with patch('neo.Prompt.PromptData.PromptData.Prompt'): with patch('neo.Prompt.Commands.Wallet.prompt', side_effect=["testpassword"]): @@ -44,7 +148,7 @@ def test_wallet_open(self): res = CommandWallet().execute(args) - self.assertEqual(str(type(res)), "") + self.assertEqual(type(res), UserWallet) # test wallet open with no path; this will also close the open wallet args = ['open'] @@ -84,7 +188,7 @@ def test_wallet_close(self): res = CommandWallet().execute(args) - self.assertEqual(str(type(res)), "") + self.assertEqual(type(res), UserWallet) # now close the open wallet manually args = ['close'] @@ -93,6 +197,49 @@ def test_wallet_close(self): self.assertTrue(res) + def test_wallet_verbose(self): + # test wallet verbose with no wallet opened + args = ['verbose'] + res = CommandWallet().execute(args) + self.assertFalse(res) + + # test wallet close with open wallet + self.OpenWallet() + args = ['verbose'] + res = CommandWallet().execute(args) + self.assertTrue(res) + + def test_wallet_create_address(self): + # test wallet create address with no wallet open + args = ['create_addr', 1] + res = CommandWallet().execute(args) + self.assertFalse(res) + + self.OpenWallet() + + # test wallet create address with no argument + args = ['create_addr'] + res = CommandWallet().execute(args) + self.assertFalse(res) + + # test wallet create address with negative number + args = ['create_addr', -1] + res = CommandWallet().execute(args) + self.assertFalse(res) + + # test wallet create successful + args = ['create_addr', 1] + res = CommandWallet().execute(args) + self.assertTrue(res) + self.assertEqual(type(res), UserWallet) + self.assertEqual(len(res.Addresses), 2) # Has one address when created. + + args = ['create_addr', 7] + res = CommandWallet().execute(args) + self.assertTrue(res) + self.assertEqual(type(res), UserWallet) + self.assertEqual(len(res.Addresses), 9) + ########################################################## ########################################################## def test_1_import_addr(self):