diff --git a/examples/address_generator.py b/examples/address_generator.py index f48328a5..01043e63 100644 --- a/examples/address_generator.py +++ b/examples/address_generator.py @@ -12,12 +12,13 @@ from typing import Optional, Text from iota import __version__, Iota +from iota.crypto.addresses import AddressGenerator from iota.crypto.types import Seed 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], bool) -> None seed = get_seed() # Create the API instance. @@ -34,7 +35,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 +112,20 @@ def output_seed(seed): 'If not specified, the first unused address will be returned.' ) + parser.add_argument( + '--security', + type = int, + default = AddressGenerator.DEFAULT_SECURITY_LEVEL, + 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:]))) diff --git a/iota/api.py b/iota/api.py index 251bfa4f..3cf85b2e 100644 --- a/iota/api.py +++ b/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, 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, ) diff --git a/iota/commands/extended/get_new_addresses.py b/iota/commands/extended/get_new_addresses.py index b14c9a8a..7acf1a8a 100644 --- a/iota/commands/extended/get_new_addresses.py +++ b/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 @@ -33,28 +33,33 @@ def get_response_filter(self): pass def _execute(self, request): + checksum = request['checksum'] # type: bool count = request['count'] # type: Optional[int] index = request['index'] # type: int security_level = request['securityLevel'] # type: int 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, bool) -> List[Address] """ Find addresses matching the command parameters. """ - # type: (Seed, int, Optional[int]) -> List[Address] - generator = AddressGenerator(seed, security_level) + 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]) + # We use addy.address here because FindTransactions does + # not work on an address with a checksum + response = FindTransactionsCommand(self.adapter)( + addresses=[addy.address] + ) if not response.get('hashes'): return [addy] @@ -75,8 +80,10 @@ def __init__(self): super(GetNewAddressesRequestFilter, self).__init__( { # Everything except ``seed`` is optional. - 'count': f.Type(int) | f.Min(1), - 'index': f.Type(int) | f.Min(0) | f.Optional(default=0), + + 'checksum': f.Type(bool) | f.Optional(default=False), + 'count': f.Type(int) | f.Min(1), + 'index': f.Type(int) | f.Min(0) | f.Optional(default=0), 'securityLevel': f.Type(int) @@ -88,6 +95,7 @@ def __init__(self): }, allow_missing_keys = { + 'checksum', 'count', 'index', 'securityLevel', diff --git a/iota/crypto/addresses.py b/iota/crypto/addresses.py index 2c1a5703..a2a5ac66 100644 --- a/iota/crypto/addresses.py +++ b/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): diff --git a/test/commands/extended/get_new_addresses_test.py b/test/commands/extended/get_new_addresses_test.py index ea1a1200..735aa90e 100644 --- a/test/commands/extended/get_new_addresses_test.py +++ b/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( + checksum = True, + count = 1, + index = 0, + seed = self.seed, + ) + + self.assertDictEqual( + response, + {'addresses': [self.addy_1_checksum]}, + ) diff --git a/test/crypto/addresses_test.py b/test/crypto/addresses_test.py index 6c1773de..64487cae 100644 --- a/test/crypto/addresses_test.py +++ b/test/crypto/addresses_test.py @@ -279,3 +279,37 @@ def test_security_level_elevated(self): ), ], ) + + def test_generator_checksum(self): + """ + Creating a generator with checksums on the addresses. + """ + ag = AddressGenerator( + self.seed_2, + security_level=AddressGenerator.DEFAULT_SECURITY_LEVEL, + checksum=True + ) + + generator = ag.create_iterator() + + # noinspection SpellCheckingInspection + self.assertEqual( + next(generator), + + Address( + b'FNKCVJPUANHNWNBAHFBTCONMCUBC9KCZ9EKREBCJ' + b'AFMABCTEPLGGXDJXVGPXDCFOUCRBWFJFLEAVOEUPY' + b'ADHVCBXFD', + ), + ) + + # noinspection SpellCheckingInspection + self.assertEqual( + next(generator), + + Address( + b'MSYILYYZLSJ99TDMGQHDOBWGHTBARCBGJZE9PIMQ' + b'LTEXJXKTDREGVTPA9NDGGLQHTMGISGRAKSLYPGWMB' + b'WIKQRCIOD', + ), + )