Issue #89 - Initial work on adding checksum option to get_new_addresses

Author: plenarius

examples/address_generator.py

@@ -16,8 +16,8 @@
 from six import binary_type, moves as compat, text_type
 
 
-def main(uri, index, count):
-  # type: (Text, int, Optional[int], bool) -> None
+def main(uri, index, count, security, checksum):
+  # type: (Text, int, Optional[int], Optional[int], Optional[bool]) -> None
   seed = get_seed()
 
   # Create the API instance.
@@ -34,7 +34,7 @@ def main(uri, index, count):
   print('')
 
   # Here's where all the magic happens!
-  api_response = api.get_new_addresses(index, count)
+  api_response = api.get_new_addresses(index, count, security, checksum)
   for addy in api_response['addresses']:
     print(binary_type(addy).decode('ascii'))
 
@@ -111,4 +111,19 @@ def output_seed(seed):
         'If not specified, the first unused address will be returned.'
   )
 
+  parser.add_argument(
+    '--security',
+      type      = int,
+      default   = 2,
+      help      = 'Security level to be used for the private key / address. Can be 1, 2 or 3',
+  )
+
+  parser.add_argument(
+    '--with-checksum',
+      action    = 'store_true',
+      default   = False,
+      dest      = 'checksum',
+      help      = 'List the address with the checksum.',
+  )
+
   main(**vars(parser.parse_args(argv[1:])))

iota/api.py

@@ -611,8 +611,9 @@ def get_new_addresses(
       index = 0,
       count = 1,
       security_level = AddressGenerator.DEFAULT_SECURITY_LEVEL,
+      checksum = False,
   ):
-    # type: (int, Optional[int], int) -> dict
+    # type: (int, Optional[int], int, Optional[bool]) -> dict
     """
     Generates one or more new addresses from the seed.
 
@@ -636,6 +637,10 @@ def get_new_addresses(
 
       This value must be between 1 and 3, inclusive.
 
+    :param checksum:
+      Specify whether to return the address with the checksum.
+      Defaults to False.
+
     :return:
       Dict with the following items::
 
@@ -651,6 +656,7 @@ def get_new_addresses(
       count         = count,
       index         = index,
       securityLevel = security_level,
+      checksum      = checksum,
       seed          = self.seed,
     )
 

iota/commands/extended/get_new_addresses.py

@@ -6,7 +6,7 @@
 
 import filters as f
 
-from iota import Address
+from iota import Address, AddressChecksum
 from iota.commands import FilterCommand, RequestFilter
 from iota.commands.core.find_transactions import FindTransactionsCommand
 from iota.crypto.addresses import AddressGenerator
@@ -36,25 +36,35 @@ def _execute(self, request):
     count           = request['count'] # type: Optional[int]
     index           = request['index'] # type: int
     security_level  = request['securityLevel'] # type: int
+    checksum        = request['checksum'] # type: Optional[bool]
     seed            = request['seed'] # type: Seed
 
     return {
-      'addresses': self._find_addresses(seed, index, count, security_level),
+      'addresses': 
+         self._find_addresses(seed, index, count, security_level, checksum),
     }
 
-  def _find_addresses(self, seed, index, count, security_level):
-    # type: (Seed, int, Optional[int], int) -> List[Address]
+  def _find_addresses(self, seed, index, count, security_level, checksum):
+    # type: (Seed, int, Optional[int], int, Optional[bool]) -> List[Address]
     """
     Find addresses matching the command parameters.
     """
-    # type: (Seed, int, Optional[int]) -> List[Address]
-    generator = AddressGenerator(seed, security_level)
+    # type: (Seed, int, Optional[bool]) -> List[Address]
+    generator = AddressGenerator(seed, security_level, checksum)
 
     if count is None:
       # Connect to Tangle and find the first address without any
       # transactions.
       for addy in generator.create_iterator(start=index):
-        response = FindTransactionsCommand(self.adapter)(addresses=[addy])
+        # If we're generating addresses with checksums we need to check
+        # for transactions on the address without the checksum
+        if not checksum:
+          response = FindTransactionsCommand(self.adapter)(addresses=[addy])
+        else:
+          response = FindTransactionsCommand(self.adapter)(
+            addresses=[addy][:-AddressChecksum.LEN]
+        )
+
 
         if not response.get('hashes'):
           return [addy]
@@ -84,12 +94,15 @@ def __init__(self):
             | f.Max(self.MAX_SECURITY_LEVEL)
             | f.Optional(default=AddressGenerator.DEFAULT_SECURITY_LEVEL),
 
+        'checksum':  f.Type(bool) | f.Optional(default=False),
+
         'seed': f.Required | Trytes(result_type=Seed),
       },
 
       allow_missing_keys = {
         'count',
         'index',
         'securityLevel',
+        'checksum',
       },
     )

iota/crypto/addresses.py

@@ -40,11 +40,12 @@ class AddressGenerator(Iterable[Address]):
     - :py:class:`iota.transaction.BundleValidator`
   """
 
-  def __init__(self, seed, security_level=DEFAULT_SECURITY_LEVEL):
-    # type: (TrytesCompatible, int) -> None
+  def __init__(self, seed, security_level=DEFAULT_SECURITY_LEVEL, checksum=False):
+    # type: (TrytesCompatible, int, bool) -> None
     super(AddressGenerator, self).__init__()
 
     self.security_level = security_level
+    self.checksum       = checksum
     self.seed           = Seed(seed)
 
   def __iter__(self):
@@ -175,7 +176,10 @@ def _generate_address(self, key_iterator):
 
     Used in the event of a cache miss.
     """
-    return self.address_from_digest(self._get_digest(key_iterator))
+    if self.checksum:
+      return self.address_from_digest(self._get_digest(key_iterator)).with_valid_checksum()
+    else:
+      return self.address_from_digest(self._get_digest(key_iterator))
 
   @staticmethod
   def _get_digest(key_iterator):

test/commands/extended/get_new_addresses_test.py

@@ -35,6 +35,7 @@ def test_pass_happy_path(self):
       'index':          1,
       'count':          1,
       'securityLevel':  2,
+      'checksum':       False,
     }
 
     filter_ = self._filter(request)
@@ -59,6 +60,7 @@ def test_pass_optional_parameters_excluded(self):
         'index':          0,
         'count':          None,
         'securityLevel':  AddressGenerator.DEFAULT_SECURITY_LEVEL,
+        'checksum':       False,
       },
     )
 
@@ -75,6 +77,9 @@ def test_pass_compatible_types(self):
       'index':          100,
       'count':          8,
       'securityLevel':  2,
+
+      # ``checksum`` must be boolean.
+      'checksum':       False,
     })
 
     self.assertFilterPasses(filter_)
@@ -86,6 +91,7 @@ def test_pass_compatible_types(self):
         'index':          100,
         'count':          8,
         'securityLevel':  2,
+        'checksum':       False,
       },
     )
 
@@ -111,6 +117,7 @@ def test_fail_unexpected_parameters(self):
         'index':          None,
         'count':          1,
         'securityLevel':  2,
+        'checksum':       False,
 
         # Some men just want to watch the world burn.
         'foo': 'bar',
@@ -306,6 +313,21 @@ def test_fail_security_level_wrong_type(self):
       },
     )
 
+  def test_fail_checksum_wrong_type(self):
+    """
+    ``checksum`` is not a boolean.
+    """
+    self.assertFilterErrors(
+      {
+        'checksum':       '2',
+        'seed':           Seed(self.seed),
+      },
+
+      {
+        'checksum': [f.Type.CODE_WRONG_TYPE],
+      },
+    )
+
 
 class GetNewAddressesCommandTestCase(TestCase):
   # noinspection SpellCheckingInspection
@@ -333,6 +355,13 @@ def setUp(self):
         b'IWYTLQUUHDWSOVXLIKVJTYZBFKLABWRBFYVSMD9NB',
       )
 
+    self.addy_1_checksum =\
+      Address(
+        b'NYMWLBUJEISSACZZBRENC9HEHYQXHCGQHSNHVCEA'
+        b'ZDCTEVNGSDUEKTSYBSQGMVJRIEDHWDYSEYCFAZAH'
+        b'9T9FPJROTW',
+      )
+
   def test_wireup(self):
     """
     Verify that the command is wired up correctly.
@@ -446,3 +475,20 @@ def test_get_addresses_online(self):
         },
       ],
     )
+
+  def test_new_address_checksum(self):
+    """
+    Generate address with a checksum.
+    """
+    response =\
+      self.command(
+        count         = 1,
+        index         = 0,
+        seed          = self.seed,
+        checksum      = True,
+      )
+
+    self.assertDictEqual(
+      response,
+      {'addresses': [self.addy_1_checksum]},
+    )