Repository: fetchai/tools-pocketbook

Issue:
Provide delete option
Whilst this can be done manually, it would be nice to be able to remove an address/key with a specific delete option.


Discussion:

Patch:
diff --git a/pocketbook/address_book.py b/pocketbook/address_book.py
index aba8bbb..2aecadb 100644
--- a/pocketbook/address_book.py
+++ b/pocketbook/address_book.py
@@ -38,6 +38,20 @@ class AddressBook:
 
         return success
 
+    def remove(self, name: str) -> bool:
+        """
+        Remove an address from the address book
+
+        :param name: The name of the address to remove
+        :return: True if successful, otherwise False
+        """
+        success = False
+        if name in self._address_book:
+            del self._address_book[name]
+            self._save()
+            success = True
+        return success
+
     def keys(self):
         return self._address_book.keys()
 
diff --git a/pocketbook/cli.py b/pocketbook/cli.py
index 21453e0..19c0cd9 100644
--- a/pocketbook/cli.py
+++ b/pocketbook/cli.py
@@ -8,6 +8,7 @@ from .commands.display import run_display
 from .commands.list import run_list
 from .commands.rename import run_rename
 from .commands.transfer import run_transfer
+from .commands.delete import run_delete
 from .disclaimer import display_disclaimer
 from .utils import NetworkUnavailableError, checked_address
 
@@ -47,6 +48,10 @@ def parse_commandline():
     parser_rename.add_argument('new', help='The new name of the account')
     parser_rename.set_defaults(handler=run_rename)
 
+    parser_delete = subparsers.add_parser('delete', aliases=['rm'])
+    parser_delete.add_argument('name', help='The name of the account to remove')
+    parser_delete.set_defaults(handler=run_delete)
+
     return parser, parser.parse_args()
 
 
@@ -76,6 +81,9 @@ def main():
             else:
                 exit_code = 0
 
+    except KeyboardInterrupt:
+        print('Stopped!')
+
     except NetworkUnavailableError:
         print('The network appears to be unavailable at the moment. Please try again later')
 
diff --git a/pocketbook/commands/delete.py b/pocketbook/commands/delete.py
new file mode 100644
index 0000000..93b493a
--- /dev/null
+++ b/pocketbook/commands/delete.py
@@ -0,0 +1,56 @@
+KEY_WARNING_TEMPLATE = """
+### WARNING ###
+
+You are about to delete the key {} from the database. This is a PERMANENT change and can
+no be reverted. This might mean that you loose access to funds forever.
+
+If you want to continue please confirm by entering the address of wallet you want to remove
+below:
+
+Address (Base58): """
+
+
+def run_delete(args):
+    from pocketbook.key_store import KeyStore
+    from pocketbook.address_book import AddressBook
+
+    key_store = KeyStore()
+    address_book = AddressBook()
+
+    is_key = args.name in key_store.list_keys()
+    is_address = args.name in address_book.keys()
+
+    def address_delete():
+        return address_book.remove(args.name)
+
+    def key_delete():
+        return key_store.remove_key(args.name)
+
+    if is_key and is_address:
+        raise RuntimeError('Corrected database, account is both an address and a key')
+    elif is_key:
+        warning = True
+        handler = key_delete
+    elif is_address:
+        warning = False
+        handler = address_delete
+    else:
+        print('Unknown key or address: {}. Please check and try again'.format(args.name))
+        return 1
+
+    # double check that the user really want to remove the key
+    if warning:
+        input_address = input(KEY_WARNING_TEMPLATE.format(args.name))
+
+        address = key_store.lookup_address(args.name)
+
+        if str(address) != input_address:
+            print('The input address for key {} does not match. Please double check and try again')
+            return 1
+
+    # perform the removal
+    if not handler():
+        print('Failed to remove the specified key or address')
+        return 1
+
+    return 0
diff --git a/pocketbook/key_store.py b/pocketbook/key_store.py
index 88cec46..01cec17 100644
--- a/pocketbook/key_store.py
+++ b/pocketbook/key_store.py
@@ -107,7 +107,7 @@ class KeyStore:
         os.rename(old_key_file_path, new_key_file_path)
 
         # remove the old entries to the metadata
-        self._index['key'] = list(filter(lambda x: x['name'] != old_name, self._index['key']))
+        self._remove_key_from_index(old_name)
 
         # update and insert the new meta data into place
         new_metadata = deepcopy(old_metadata)
@@ -117,6 +117,25 @@ class KeyStore:
 
         return True
 
+    def remove_key(self, name: str) -> bool:
+        if name not in self.list_keys():
+            return False
+
+        # remove the key path from the disk
+        key_file_path = self._format_key_path(name)
+        if not os.path.exists(key_file_path):
+            raise RuntimeError('Key not present on disk, unable to remove')
+        os.remove(key_file_path)
+
+        # update the index
+        self._remove_key_from_index(name)
+        self._flush_index()
+
+        return True
+
+    def _remove_key_from_index(self, name: str):
+        self._index['key'] = list(filter(lambda x: x['name'] != name, self._index['key']))
+
     def _format_key_path(self, name: str):
         return os.path.join(self._root, '{}.key'.format(name))
 


Test patch:
diff --git a/tests/commands/test_delete.py b/tests/commands/test_delete.py
new file mode 100644
index 0000000..b88fbd4
--- /dev/null
+++ b/tests/commands/test_delete.py
@@ -0,0 +1,136 @@
+import unittest
+from unittest.mock import patch, Mock
+
+from pocketbook.address_book import AddressBook
+from pocketbook.key_store import KeyStore
+from tests.utils import SAMPLE_ADDRESS
+
+
+class DeleteCommandTests(unittest.TestCase):
+
+    @patch('pocketbook.address_book.AddressBook', spec=AddressBook)
+    @patch('pocketbook.key_store.KeyStore', spec=KeyStore)
+    def test_delete_address(self, MockKeyStore, MockAddressBook):
+        key_store = MockKeyStore()
+        key_store.list_keys.return_value = []
+
+        address_book = MockAddressBook()
+        address_book.keys.return_value = ['foo']
+
+        args = Mock()
+        args.name = 'foo'
+
+        from pocketbook.commands.delete import run_delete
+        self.assertEqual(0, run_delete(args))
+
+        address_book.remove.assert_called_once_with('foo')
+
+    @patch('builtins.input', return_value='foo-address')
+    @patch('pocketbook.address_book.AddressBook', spec=AddressBook)
+    @patch('pocketbook.key_store.KeyStore', spec=KeyStore)
+    def test_delete_key(self, MockKeyStore, MockAddressBook, *args):
+        key_store = MockKeyStore()
+        key_store.list_keys.return_value = ['foo']
+        key_store.lookup_address.return_value = 'foo-address'
+
+        address_book = MockAddressBook()
+        address_book.keys.return_value = []
+
+        args = Mock()
+        args.name = 'foo'
+
+        from pocketbook.commands.delete import run_delete
+        self.assertEqual(0, run_delete(args))
+
+        key_store.lookup_address.assert_called_once_with('foo')
+        key_store.remove_key.assert_called_once_with('foo')
+
+    @patch('builtins.input', return_value='foo-address')
+    @patch('pocketbook.address_book.AddressBook', spec=AddressBook)
+    @patch('pocketbook.key_store.KeyStore', spec=KeyStore)
+    def test_delete_unknown(self, MockKeyStore, MockAddressBook, *args):
+        key_store = MockKeyStore()
+        key_store.list_keys.return_value = []
+
+        address_book = MockAddressBook()
+        address_book.keys.return_value = []
+
+        args = Mock()
+        args.name = 'foo'
+
+        from pocketbook.commands.delete import run_delete
+        self.assertEqual(1, run_delete(args))
+
+    @patch('builtins.input', return_value='not-sample-address')
+    @patch('pocketbook.address_book.AddressBook', spec=AddressBook)
+    @patch('pocketbook.key_store.KeyStore', spec=KeyStore)
+    def test_delete_key_invalid_confirmation(self, MockKeyStore, MockAddressBook, *args):
+        key_store = MockKeyStore()
+        key_store.list_keys.return_value = ['foo']
+        key_store.lookup_address.return_value = SAMPLE_ADDRESS
+
+        address_book = MockAddressBook()
+        address_book.keys.return_value = []
+
+        args = Mock()
+        args.name = 'foo'
+
+        from pocketbook.commands.delete import run_delete
+        self.assertEqual(1, run_delete(args))
+
+        key_store.lookup_address.assert_called_once_with('foo')
+
+    @patch('pocketbook.address_book.AddressBook', spec=AddressBook)
+    @patch('pocketbook.key_store.KeyStore', spec=KeyStore)
+    def test_delete_address_failed(self, MockKeyStore, MockAddressBook):
+        key_store = MockKeyStore()
+        key_store.list_keys.return_value = []
+
+        address_book = MockAddressBook()
+        address_book.keys.return_value = ['foo']
+        address_book.remove.return_value = False
+
+        args = Mock()
+        args.name = 'foo'
+
+        from pocketbook.commands.delete import run_delete
+        self.assertEqual(1, run_delete(args))
+
+        address_book.remove.assert_called_once_with('foo')
+
+    @patch('builtins.input', return_value='foo-address')
+    @patch('pocketbook.address_book.AddressBook', spec=AddressBook)
+    @patch('pocketbook.key_store.KeyStore', spec=KeyStore)
+    def test_delete_key_failed(self, MockKeyStore, MockAddressBook, *args):
+        key_store = MockKeyStore()
+        key_store.list_keys.return_value = ['foo']
+        key_store.lookup_address.return_value = 'foo-address'
+        key_store.remove_key.return_value = False
+
+        address_book = MockAddressBook()
+        address_book.keys.return_value = []
+
+        args = Mock()
+        args.name = 'foo'
+
+        from pocketbook.commands.delete import run_delete
+        self.assertEqual(1, run_delete(args))
+
+        key_store.lookup_address.assert_called_once_with('foo')
+        key_store.remove_key.assert_called_once_with('foo')
+
+    @patch('pocketbook.address_book.AddressBook', spec=AddressBook)
+    @patch('pocketbook.key_store.KeyStore', spec=KeyStore)
+    def test_db_corruption(self, MockKeyStore, MockAddressBook):
+        key_store = MockKeyStore()
+        key_store.list_keys.return_value = ['foo']
+
+        address_book = MockAddressBook()
+        address_book.keys.return_value = ['foo']
+
+        args = Mock()
+        args.name = 'foo'
+
+        from pocketbook.commands.delete import run_delete
+        with self.assertRaises(RuntimeError):
+            run_delete(args)
diff --git a/tests/test_address_book.py b/tests/test_address_book.py
index 694db76..e76c6e0 100644
--- a/tests/test_address_book.py
+++ b/tests/test_address_book.py
@@ -8,6 +8,12 @@ from .utils import TemporaryPocketBookRoot, SAMPLE_ADDRESS
 
 
 class AddressBookTests(unittest.TestCase):
+    def assertAddressIsNotPresentOnDisk(self, name: str, ctx: TemporaryPocketBookRoot):
+        address_index_path = os.path.join(ctx.root, AddressBook.INDEX_FILE_NAME)
+        if os.path.isfile(address_index_path):
+            with open(address_index_path, 'r') as address_index_file:
+                address_index = toml.load(address_index_file)
+                self.assertNotIn(name, address_index)
 
     def assertAddressIsPresentOnDisk(self, name: str, address: str, ctx: TemporaryPocketBookRoot):
         address_index_path = os.path.join(ctx.root, AddressBook.INDEX_FILE_NAME)
@@ -92,9 +98,31 @@ class AddressBookTests(unittest.TestCase):
             address_book1.add('sample', SAMPLE_ADDRESS)
 
             self.assertTrue(address_book1.rename('sample', 'sample2'))
+            self.assertAddressIsNotPresentOnDisk('sample', ctx)
             self.assertAddressIsPresentOnDisk('sample2', SAMPLE_ADDRESS, ctx)
 
     def test_rename_fails(self):
         with TemporaryPocketBookRoot() as ctx:
             address_book1 = AddressBook(root=ctx.root)
             self.assertFalse(address_book1.rename('sample', 'sample2'))
+
+    def test_remove(self):
+        with TemporaryPocketBookRoot() as ctx:
+            address_book = AddressBook(root=ctx.root)
+            address_book.add('sample', SAMPLE_ADDRESS)
+            self.assertTrue(address_book.remove('sample'))
+
+            self.assertNotIn('sample', address_book.keys())
+
+    def test_remove_is_stored_on_disk(self):
+        with TemporaryPocketBookRoot() as ctx:
+            address_book = AddressBook(root=ctx.root)
+            address_book.add('sample', SAMPLE_ADDRESS)
+            self.assertTrue(address_book.remove('sample'))
+
+            self.assertAddressIsNotPresentOnDisk('sample', ctx)
+
+    def test_removed_failed(self):
+        with TemporaryPocketBookRoot() as ctx:
+            address_book = AddressBook(root=ctx.root)
+            self.assertFalse(address_book.remove('sample'))
diff --git a/tests/test_key_store.py b/tests/test_key_store.py
index 04fba86..8f21866 100644
--- a/tests/test_key_store.py
+++ b/tests/test_key_store.py
@@ -9,8 +9,24 @@ from pocketbook.key_store import KeyStore, DuplicateKeyNameError, KeyNotFoundErr
 from .utils import TemporaryPocketBookRoot, SUPER_SECURE_PASSWORD
 
 
-class KeyStoreTests(unittest.TestCase):
-    def assertFlushedToDisk(self, name: str, password: str, entity: Entity, ctx: TemporaryPocketBookRoot):
+class KeyStoreTestCase(unittest.TestCase):
+    def assertKeyIsNotPresentOnDisk(self, name: str, ctx: TemporaryPocketBookRoot):
+        self.assertTrue(os.path.isdir(ctx.root))
+
+        # check the index record
+        index_path = os.path.join(ctx.root, KeyStore.INDEX_FILE_NAME)
+        self.assertTrue(os.path.isfile(index_path))
+        with open(index_path, 'r') as index_file:
+            index = toml.load(index_path)
+
+            # build up the key listings
+            key_addresses = {key['name']: key['address'] for key in index.get('key', [])}
+            self.assertNotIn(name, key_addresses)
+
+        key_path = os.path.join(ctx.root, '{}.key'.format(name))
+        self.assertFalse(os.path.isfile(key_path))
+
+    def assertKeyIsPresentOnDisk(self, name: str, password: str, entity: Entity, ctx: TemporaryPocketBookRoot):
         self.assertTrue(os.path.isdir(ctx.root))
 
         # check that the index record has been created correctly
@@ -36,6 +52,9 @@ class KeyStoreTests(unittest.TestCase):
             recovered = Entity.load(key_file, password)
             self.assertEqual(recovered.private_key, entity.private_key)
 
+
+class KeyStoreTests(KeyStoreTestCase):
+
     def test_init(self):
         with TemporaryPocketBookRoot() as ctx:
             key_store = KeyStore(root=ctx.root)
@@ -60,7 +79,7 @@ class KeyStoreTests(unittest.TestCase):
             key_store.add_key('sample', SUPER_SECURE_PASSWORD, entity)
 
             # check the files on have been set correctly
-            self.assertFlushedToDisk('sample', SUPER_SECURE_PASSWORD, entity, ctx)
+            self.assertKeyIsPresentOnDisk('sample', SUPER_SECURE_PASSWORD, entity, ctx)
 
     def test_address_lookup(self):
         with TemporaryPocketBookRoot() as ctx:
@@ -145,7 +164,7 @@ class KeyStoreTests(unittest.TestCase):
             key_store.add_key('sample', SUPER_SECURE_PASSWORD, entity)
 
             self.assertTrue(key_store.rename_key('sample', 'sample2'))
-            self.assertFlushedToDisk('sample2', SUPER_SECURE_PASSWORD, entity, ctx)
+            self.assertKeyIsPresentOnDisk('sample2', SUPER_SECURE_PASSWORD, entity, ctx)
 
     def test_rename_failure_doesnt_exist(self):
         with TemporaryPocketBookRoot() as ctx:
@@ -172,3 +191,43 @@ class KeyStoreTests(unittest.TestCase):
 
             with patch.object(key_store, '_lookup_meta_data', return_value=None):
                 self.assertFalse(key_store.rename_key('sample1', 'sample2'))
+
+    def test_remove_failed(self):
+        with TemporaryPocketBookRoot() as ctx:
+            key_store = KeyStore(root=ctx.root)
+            self.assertFalse(key_store.remove_key('sample'))
+
+    def test_remove(self):
+        with TemporaryPocketBookRoot() as ctx:
+            entity = Entity()
+
+            key_store = KeyStore(root=ctx.root)
+            key_store.add_key('sample', SUPER_SECURE_PASSWORD, entity)
+
+            self.assertTrue(key_store.remove_key('sample'))
+            self.assertNotIn('sample', key_store.list_keys())
+
+    def test_remove_reflected_on_disk(self):
+        with TemporaryPocketBookRoot() as ctx:
+            entity = Entity()
+
+            key_store = KeyStore(root=ctx.root)
+            key_store.add_key('sample', SUPER_SECURE_PASSWORD, entity)
+
+            self.assertTrue(key_store.remove_key('sample'))
+            self.assertKeyIsNotPresentOnDisk('sample', ctx)
+
+    def test_remove_with_corrupted_files(self):
+        with TemporaryPocketBookRoot() as ctx:
+            entity = Entity()
+
+            key_store = KeyStore(root=ctx.root)
+            key_store.add_key('sample', SUPER_SECURE_PASSWORD, entity)
+
+            # corrupt the files
+            key_path = key_store._format_key_path('sample')
+            self.assertTrue(os.path.isfile(key_path))
+            os.remove(key_path)
+
+            with self.assertRaises(RuntimeError):
+                key_store.remove_key('sample')