1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2025-01-08 14:31:06 +00:00

trezor.crypto: refactor base58; allow custom digest function

This commit is contained in:
Pavol Rusnak 2016-10-25 14:07:52 +02:00
parent 9d5fe82b45
commit df5e770dec
No known key found for this signature in database
GPG Key ID: 91F3B339B9A02A3D
2 changed files with 29 additions and 11 deletions

View File

@ -5,6 +5,9 @@ import unittest
from ubinascii import unhexlify
from trezor.crypto import base58
from trezor.crypto.hashlib import ripemd160
digestfunc_graphene = lambda x: ripemd160(x).digest()[:4]
class TestCryptoBase58(unittest.TestCase):
@ -62,13 +65,24 @@ class TestCryptoBase58(unittest.TestCase):
('055ece0cadddc415b1980f001785947120acdb36fc', '3ALJH9Y951VCGcVZYAdpA3KchoP9McEj1G'),
]
vectors_graphene = [
('02e649f63f8e8121345fd7f47d0d185a3ccaa843115cd2e9392dcd9b82263bc680', '6dumtt9swxCqwdPZBGXh9YmHoEjFFnNfwHaTqRbQTghGAY2gRz'),
('021c7359cd885c0e319924d97e3980206ad64387aff54908241125b3a88b55ca16', '5725vivYpuFWbeyTifZ5KevnHyqXCi5hwHbNU9cYz1FHbFXCxX'),
('02f561e0b57a552df3fa1df2d87a906b7a9fc33a83d5d15fa68a644ecb0806b49a', '6kZKHSuxqAwdCYsMvwTcipoTsNE2jmEUNBQufGYywpniBKXWZK'),
('03e7595c3e6b58f907bee951dc29796f3757307e700ecf3d09307a0cc4a564eba3', '8b82mpnH8YX1E9RHnU2a2YgLTZ8ooevEGP9N15c1yFqhoBvJur'),
]
def test_decode_check(self):
for a, b in self.vectors:
self.assertEqual(base58.decode_check(b), unhexlify(a))
for a, b in self.vectors_graphene:
self.assertEqual(base58.decode_check(b, digestfunc=digestfunc_graphene), unhexlify(a))
def test_encode_check(self):
for a, b in self.vectors:
self.assertEqual(base58.encode_check(unhexlify(a)), b)
for a, b in self.vectors_graphene:
self.assertEqual(base58.encode_check(unhexlify(a), digestfunc=digestfunc_graphene), b)
if __name__ == '__main__':
unittest.main()

View File

@ -13,8 +13,6 @@
# This module adds shiny packaging and support for python3.
#
from .hashlib import sha256
# 58 character alphabet used
_alphabet = '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz'
@ -36,7 +34,7 @@ def encode(data: bytes) -> str:
acc, mod = divmod(acc, 58)
result += _alphabet[mod]
return ''.join([c for c in reversed(result + _alphabet[0] * (origlen - newlen))])
return ''.join((c for c in reversed(result + _alphabet[0] * (origlen - newlen))))
def decode(string: str) -> bytes:
@ -57,25 +55,31 @@ def decode(string: str) -> bytes:
acc, mod = divmod(acc, 256)
result.append(mod)
return bytes([b for b in reversed(result +[0] * (origlen - newlen))])
return bytes((b for b in reversed(result +[0] * (origlen - newlen))))
def encode_check(data: bytes) -> str:
def encode_check(data: bytes, digestfunc=None) -> str:
'''
Convert bytes to base58 encoded string, append checksum.
'''
digest = sha256(sha256(data).digest()).digest()
return encode(data + digest[:4])
if digestfunc is None:
from .hashlib import sha256
digestfunc = lambda x: sha256(sha256(x).digest()).digest()[:4]
return encode(data + digestfunc(data))
def decode_check(string: str) -> bytes:
def decode_check(string: str, digestfunc=None) -> bytes:
'''
Convert base58 encoded string to bytes and verify checksum.
'''
result = decode(string)
result, check = result[:-4], result[-4:]
digest = sha256(sha256(result).digest()).digest()
if check != digest[:4]:
if digestfunc is None:
from .hashlib import sha256
digestfunc = lambda x: sha256(sha256(x).digest()).digest()[:4]
digestlen = len(digestfunc(b''))
result, check = result[:-digestlen], result[-digestlen:]
if check != digestfunc(result):
raise ValueError('Invalid checksum')
return result