[refactor-prompt] add wallet import watchaddr

neo/Implementations/Wallets/peewee/UserWallet.py

@@ -178,7 +178,7 @@ def AddWatchOnly(self, script_hash):
             address.save()
             return address
         else:
-            raise Exception("Address already exists in wallet")
+            raise ValueError("Address already exists in wallet")
 
     def AddNEP5Token(self, token):
 

neo/Implementations/Wallets/peewee/test_user_wallet.py

@@ -16,7 +16,6 @@
 
 
 class UserWalletTestCase(WalletFixtureTestCase):
-
     wallet_1_script_hash = UInt160(data=b'\x1c\xc9\xc0\\\xef\xff\xe6\xcd\xd7\xb1\x82\x81j\x91R\xec!\x8d.\xc0')
 
     wallet_1_addr = 'AJQ6FoaSXDFzA6wLnyZ1nFN7SGSN2oNTc3'
@@ -100,7 +99,7 @@ def test_3_import_watch_addr(self):
 
         # now add it again
 
-        self.assertRaises(Exception, wallet.AddWatchOnly, self.import_watch_addr)
+        self.assertRaises(ValueError, wallet.AddWatchOnly, self.import_watch_addr)
 
     def test_4_get_change_address(self):
 

neo/Prompt/Commands/Wallet.py

@@ -23,6 +23,7 @@
 from neo.Prompt.Commands.Send import CommandWalletSend, CommandWalletSendMany, CommandWalletSign
 from neo.Prompt.Commands.Tokens import CommandWalletToken
 from neo.logging import log_manager
+from neocore.Utils import isValidPublicAddress
 
 logger = log_manager.getLogger()
 
@@ -287,6 +288,7 @@ def __init__(self):
         super().__init__()
         self.register_sub_command(CommandWalletImportWIF())
         self.register_sub_command(CommandWalletImportNEP2())
+        self.register_sub_command(CommandWalletImportWatchAddr())
 
     def command_desc(self):
         return CommandDesc('import', 'import wallet items')
@@ -379,6 +381,36 @@ def command_desc(self):
         return CommandDesc('nep2', 'import a passphrase protected private key record (NEP-2 format)', [p1])
 
 
+class CommandWalletImportWatchAddr(CommandBase):
+    def __init__(self):
+        super().__init__()
+
+    def execute(self, arguments):
+        wallet = PromptData.Wallet
+
+        if len(arguments) != 1:
+            print("Please specify the required parameters")
+            return False
+
+        addr = arguments[0]
+        if not isValidPublicAddress(addr):
+            print("Invalid address specified")
+            return False
+
+        try:
+            addr_script_hash = wallet.ToScriptHash(addr)
+            wallet.AddWatchOnly(addr_script_hash)
+        except ValueError as e:
+            print(str(e))
+            return False
+
+        print(f"Added address {addr} as watch-only")
+        return True
+
+    def command_desc(self):
+        p1 = ParameterDesc('address', 'a public NEO address to watch')
+        return CommandDesc('watch_addr', 'import a public address as watch only', [p1])
+
 #########################################################################
 #########################################################################
 
@@ -425,32 +457,6 @@ def DeleteAddress(wallet, addr):
     return success
 
 
-def ImportWatchAddr(wallet, addr):
-    if wallet is None:
-        print("Please open a wallet")
-        return False
-
-    script_hash = None
-    try:
-        script_hash = wallet.ToScriptHash(addr)
-    except Exception as e:
-        pass
-
-    if not script_hash:
-        try:
-            data = bytearray(binascii.unhexlify(addr.encode('utf-8')))
-            data.reverse()
-            script_hash = UInt160(data=data)
-        except Exception as e:
-            pass
-
-    if script_hash:
-        wallet.AddWatchOnly(script_hash)
-        print("added watch address")
-    else:
-        print("incorrect format for watch address")
-
-
 def ImportToken(wallet, contract_hash):
     if wallet is None:
         print("please open a wallet")

neo/Prompt/Commands/tests/test_wallet_commands.py

@@ -6,7 +6,7 @@
 from neocore.Fixed8 import Fixed8
 from neo.Core.TX.ClaimTransaction import ClaimTransaction
 from neo.Prompt.Commands.Wallet import CommandWallet
-from neo.Prompt.Commands.Wallet import CreateAddress, DeleteAddress, ImportToken, ImportWatchAddr, ShowUnspentCoins, SplitUnspentCoin
+from neo.Prompt.Commands.Wallet import CreateAddress, DeleteAddress, ImportToken, ShowUnspentCoins, SplitUnspentCoin
 from neo.Prompt.PromptData import PromptData
 import os
 import shutil
@@ -386,26 +386,6 @@ def test_wallet_alias(self):
     ##########################################################
     ##########################################################
 
-    def test_1_import_addr(self):
-        wallet = self.GetWallet1(recreate=True)
-
-        self.assertEqual(len(wallet.LoadWatchOnly()), 0)
-
-        result = ImportWatchAddr(wallet, self.watch_addr_str)
-
-        self.assertEqual(len(wallet.LoadWatchOnly()), 1)
-
-    def test_2_import_addr(self):
-        wallet = self.GetWallet1()
-
-        self.assertEqual(len(wallet.LoadWatchOnly()), 1)
-
-        success = DeleteAddress(wallet, self.watch_addr_str)
-
-        self.assertTrue(success)
-
-        self.assertEqual(len(wallet.LoadWatchOnly()), 0)
-
     def test_3_import_token(self):
         wallet = self.GetWallet1(recreate=True)
 
@@ -606,3 +586,38 @@ def test_wallet_import_nep2(self):
                 self.assertFalse(res)
                 # if we imported successfully we get the wallet1 address
                 self.assertIn(self.wallet_1_addr, mock_print.getvalue())
+
+    def test_wallet_import_watchaddr(self):
+        self.OpenWallet1()
+
+        # test missing wif key argument
+        with patch('sys.stdout', new=StringIO()) as mock_print:
+            args = ['import', 'watch_addr']
+            res = CommandWallet().execute(args)
+            self.assertFalse(res)
+            self.assertIn("specify the required parameter", mock_print.getvalue())
+
+        # test with bad address
+        with patch('sys.stdout', new=StringIO()) as mock_print:
+            args = ['import', 'watch_addr', 'too_short_addr']
+            res = CommandWallet().execute(args)
+            self.assertFalse(res)
+            self.assertIn("Invalid address specified", mock_print.getvalue())
+
+        # test with good address
+        with patch('sys.stdout', new=StringIO()) as mock_print:
+            address = 'AZfFBeBqtJvaTK9JqG8uk6N7FppQY6byEg'
+            args = ['import', 'watch_addr', address]
+            res = CommandWallet().execute(args)
+            self.assertTrue(res)
+            self.assertIn("Added address", mock_print.getvalue())
+            self.assertIn("watch-only", mock_print.getvalue())
+            self.assertIn(PromptData.Wallet.ToScriptHash(address), PromptData.Wallet.LoadWatchOnly())
+
+        # test address already exists
+        with patch('sys.stdout', new=StringIO()) as mock_print:
+            address = 'AZfFBeBqtJvaTK9JqG8uk6N7FppQY6byEg'
+            args = ['import', 'watch_addr', address]
+            res = CommandWallet().execute(args)
+            self.assertFalse(res)
+            self.assertIn("Address already exists in wallet", mock_print.getvalue())