Get base58 support working

petertodd · petertodd · commit 1959525559ce · 2013-04-03T07:09:46.000-04:00
diff --git a/README b/README
@@ -2,3 +2,8 @@
 This python library provides an easy interface to the bitcoin
 data structures and protocol.
 
+
+Unit tests
+----------
+
+python -m unittest discover
diff --git a/bitcoin/base58.py b/bitcoin/base58.py
@@ -8,54 +8,92 @@
 # file COPYING or http://www.opensource.org/licenses/mit-license.php.
 #
 
-from bitcoin.serialize import Hash
+from bitcoin.serialize import Hash, ser_uint256
 
 b58_digits = '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz'
 
-def encode (n):
-    l = []
+from binascii import hexlify, unhexlify
+
+class Base58Error(Exception):
+    pass
+
+class InvalidBase58Error(Base58Error):
+    pass
+
+def encode(b):
+    """Encode bytes to a base58-encoded string"""
+
+    # Convert big-endian bytes to integer
+    n = int('0x0' + hexlify(b), 16)
+
+    # Divide that integer into bas58
+    res = []
     while n > 0:
         n, r = divmod (n, 58)
-        l.insert (0, (b58_digits[r]))
-    return ''.join (l)
-
-def decode (s):
-    n = 0
-    for ch in s:
-        n *= 58
-        digit = b58_digits.index (ch)
-        n += digit
-    return n
+        res.append(b58_digits[r])
+    res = ''.join(res[::-1])
 
-def encode_padded (s):
-    res = base58_encode (int ('0x' + s.encode ('hex'), 16))
+    # Encode leading zeros as base58 zeros
     pad = 0
-    for c in s:
+    for c in b:
         if c == chr(0): pad += 1
         else: break
     return b58_digits[0] * pad + res
 
-def decode_padded (s):
-    pad = 0
+def decode(s):
+    """Decode a base58-encoding string, returning bytes"""
+    if not s:
+        return b''
+
+    # Convert the string to an integer
+    n = 0
     for c in s:
-        if c == b58_digits[0]: pad += 1
-        else: break
-    h = '%x' % base58_decode (s)
+        n *= 58
+        if c not in b58_digits:
+            raise InvalidBase58Error('Character %r is not a valid base58 character' % c)
+        digit = b58_digits.index(c)
+        n += digit
+
+    # Convert the integer to bytes
+    h = '%x' % n
     if len(h) % 2:
         h = '0' + h
-    res = h.decode ('hex')
-    return chr(0) * pad + res
-
-def key_to_address (s):
-    vs = chr (addrtype) + s
-    check = Hash(vs)[:4]
-    return base58_encode_padded (vs + check)
-
-def address_to_key (s):
-    k = base58_decode_padded (s)
-    hash160, check0 = k[1:-4], k[-4:]
-    check1 = Hash (chr (addrtype) + hash160)[:4]
-    if check0 != check1:
-        return None
-    return hash160
+    res = unhexlify(h)
+
+    # Add padding back.
+    pad = 0
+    for c in s[:-1]:
+        if c == b58_digits[0]: pad += 1
+        else: break
+    return b'\x00' * pad + res
+
+
+class Base58ChecksumError(Base58Error):
+    pass
+
+class CBase58Data(bytes):
+    def __new__(cls, data, nVersion):
+        self = super(CBase58Data, cls).__new__(cls, data)
+        self.nVersion = nVersion
+        return self
+
+    def __str__(self):
+        vs = chr(self.nVersion) + self
+        check = ser_uint256(Hash(vs))[0:4]
+        return encode(vs + check)
+
+    @classmethod
+    def from_str(cls, s):
+        k = decode(s)
+        addrbyte, data, check0 = k[0], k[1:-4], k[-4:]
+        check1 = ser_uint256(Hash(addrbyte + data))[:4]
+        if check0 != check1:
+            raise Base58ChecksumError('Checksum mismatch: expected %r, calculated %r' % (check0, check1))
+        return cls(data, ord(addrbyte))
+
 
+class CBitcoinAddress(CBase58Data):
+    PUBKEY_ADDRESS = 0
+    SCRIPT_ADDRESS = 5
+    PUBKEY_ADDRESS_TEST = 111
+    SCRIPT_ADDRESS_TEST = 196
diff --git a/bitcoin/tests/__init__.py b/bitcoin/tests/__init__.py
@@ -0,0 +1,2 @@
+# Distributed under the MIT/X11 software license, see the accompanying
+# file COPYING or http://www.opensource.org/licenses/mit-license.php.
diff --git a/bitcoin/tests/data/base58_encode_decode.json b/bitcoin/tests/data/base58_encode_decode.json
@@ -0,0 +1,14 @@
+[
+["", ""],
+["61", "2g"],
+["626262", "a3gV"],
+["636363", "aPEr"],
+["73696d706c792061206c6f6e6720737472696e67", "2cFupjhnEsSn59qHXstmK2ffpLv2"],
+["00eb15231dfceb60925886b67d065299925915aeb172c06647", "1NS17iag9jJgTHD1VXjvLCEnZuQ3rJDE9L"],
+["516b6fcd0f", "ABnLTmg"],
+["bf4f89001e670274dd", "3SEo3LWLoPntC"],
+["572e4794", "3EFU7m"],
+["ecac89cad93923c02321", "EJDM8drfXA6uyA"],
+["10c8511e", "Rt5zm"],
+["00000000000000000000", "1111111111"]
+]
diff --git a/bitcoin/tests/test_base58.py b/bitcoin/tests/test_base58.py
@@ -0,0 +1,35 @@
+# Distributed under the MIT/X11 software license, see the accompanying
+# file COPYING or http://www.opensource.org/licenses/mit-license.php.
+
+from __future__ import print_function
+
+import json
+import os
+import unittest
+
+from binascii import unhexlify
+
+from bitcoin.base58 import *
+
+
+def load_test_vector(name):
+    with open(os.path.dirname(__file__) + '/data/' + name, 'r') as fd:
+        for testcase in json.load(fd):
+            yield testcase
+
+class Test_base58(unittest.TestCase):
+    def test_encode_decode(self):
+        for exp_bin, exp_base58 in load_test_vector('base58_encode_decode.json'):
+            exp_bin = unhexlify(exp_bin)
+
+            act_base58 = encode(exp_bin)
+            act_bin = decode(exp_base58)
+
+            self.assertEqual(act_base58, exp_base58)
+            self.assertEqual(act_bin, exp_bin)
+
+    def test_invalid_base58_exception(self):
+        with self.assertRaises(InvalidBase58Error):
+            decode('#')
+
+    # FIXME: need to test CBitcoinAddress

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+# Distributed under the MIT/X11 software license, see the accompanying`
	`2`	`+# file COPYING or http://www.opensource.org/licenses/mit-license.php.`