1# Copyright (c) 2012-2020 The Bitcoin Core developers
2# Distributed under the MIT software license, see the accompanying
3# file COPYING or http://www.opensource.org/licenses/mit-license.php.
4'''
5Bitcoin base58 encoding and decoding.
6
7Based on https://bitcointalk.org/index.php?topic=1026.0 (public domain)
8'''
9import hashlib
10
11# for compatibility with following code...
12class SHA256:
13    new = hashlib.sha256
14
15if str != bytes:
16    # Python 3.x
17    def ord(c):
18        return c
19    def chr(n):
20        return bytes( (n,) )
21
22__b58chars = '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz'
23__b58base = len(__b58chars)
24b58chars = __b58chars
25
26def b58encode(v):
27    """ encode v, which is a string of bytes, to base58.
28    """
29    long_value = 0
30    for (i, c) in enumerate(v[::-1]):
31        if isinstance(c, str):
32            c = ord(c)
33        long_value += (256**i) * c
34
35    result = ''
36    while long_value >= __b58base:
37        div, mod = divmod(long_value, __b58base)
38        result = __b58chars[mod] + result
39        long_value = div
40    result = __b58chars[long_value] + result
41
42    # Bitcoin does a little leading-zero-compression:
43    # leading 0-bytes in the input become leading-1s
44    nPad = 0
45    for c in v:
46        if c == 0:
47            nPad += 1
48        else:
49            break
50
51    return (__b58chars[0]*nPad) + result
52
53def b58decode(v, length = None):
54    """ decode v into a string of len bytes
55    """
56    long_value = 0
57    for i, c in enumerate(v[::-1]):
58        pos = __b58chars.find(c)
59        assert pos != -1
60        long_value += pos * (__b58base**i)
61
62    result = bytes()
63    while long_value >= 256:
64        div, mod = divmod(long_value, 256)
65        result = chr(mod) + result
66        long_value = div
67    result = chr(long_value) + result
68
69    nPad = 0
70    for c in v:
71        if c == __b58chars[0]:
72            nPad += 1
73            continue
74        break
75
76    result = bytes(nPad) + result
77    if length is not None and len(result) != length:
78        return None
79
80    return result
81
82def checksum(v):
83    """Return 32-bit checksum based on SHA256"""
84    return SHA256.new(SHA256.new(v).digest()).digest()[0:4]
85
86def b58encode_chk(v):
87    """b58encode a string, with 32-bit checksum"""
88    return b58encode(v + checksum(v))
89
90def b58decode_chk(v):
91    """decode a base58 string, check and remove checksum"""
92    result = b58decode(v)
93    if result is None:
94        return None
95    if result[-4:] == checksum(result[:-4]):
96        return result[:-4]
97    else:
98        return None
99
100def get_bcaddress_version(strAddress):
101    """ Returns None if strAddress is invalid.  Otherwise returns integer version of address. """
102    addr = b58decode_chk(strAddress)
103    if addr is None or len(addr)!=21:
104        return None
105    version = addr[0]
106    return ord(version)
107
108if __name__ == '__main__':
109    # Test case (from http://gitorious.org/bitcoin/python-base58.git)
110    assert get_bcaddress_version('15VjRaDX9zpbA8LVnbrCAFzrVzN7ixHNsC') == 0
111    _ohai = 'o hai'.encode('ascii')
112    _tmp = b58encode(_ohai)
113    assert _tmp == 'DYB3oMS'
114    assert b58decode(_tmp, 5) == _ohai
115    print("Tests passed")
116