diff --git a/silkaj/crypto_tools.py b/silkaj/crypto_tools.py index 637d2c36a3f2192c631693d4144a3ec7b77d3e92..4ebd8bd1969a979e46c679f85e93ea45172ce90c 100644 --- a/silkaj/crypto_tools.py +++ b/silkaj/crypto_tools.py @@ -35,12 +35,7 @@ def check_public_key(pubkey, display_error): return pubkey elif re.search(re.compile(PUBKEY_CHECKSUM_PATTERN), pubkey): pubkey, checksum = pubkey.split(":") - pubkey_byte = b58_decode(pubkey) - checksum_calculed = b58_encode( - hash.sha256( - hash.sha256(pubkey_byte, encoding.RawEncoder), encoding.RawEncoder - ) - )[:3] + checksum_calculed = gen_checksum(pubkey) if checksum_calculed == checksum: return pubkey else: @@ -52,6 +47,16 @@ def check_public_key(pubkey, display_error): return False +def gen_checksum(pubkey): + """ + Returns the checksum of the input pubkey (encoded in b58) + """ + pubkey_byte = b58_decode(pubkey) + return b58_encode( + hash.sha256(hash.sha256(pubkey_byte, encoding.RawEncoder), encoding.RawEncoder) + )[:3] + + b58_digits = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz" diff --git a/tests/test_crypto_tools.py b/tests/test_crypto_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..3634e04fdf8e05cb882f6a38c0598447f19fd5c8 --- /dev/null +++ b/tests/test_crypto_tools.py @@ -0,0 +1,30 @@ +""" +Copyright 2016-2020 Maël Azimi <m.a@moul.re> + +Silkaj is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as published by +the Free Software Foundation, either version 3 of the License, or +(at your option) any later version. + +Silkaj is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with Silkaj. If not, see <https://www.gnu.org/licenses/>. +""" + +import pytest + +from silkaj import crypto_tools + + +@pytest.mark.parametrize( + "pubkey, checksum", + [ + ("J4c8CARmP9vAFNGtHRuzx14zvxojyRWHW2darguVqjtX", "KAv"), + ], +) +def test_gen_checksum(pubkey, checksum): + assert checksum == crypto_tools.gen_checksum(pubkey)