Prefix unittest and refactoring

graphenebase/account.py

@@ -8,29 +8,19 @@
 from .base58 import ripemd160, Base58, doublesha256
 from .dictionary import words as BrainKeyDictionary
 from .utils import _bytes
+from .prefix import Prefix
 
 import ecdsa
 
 
-class Prefix:
-    """ This class is meant to allow changing the prefix.
-        The prefix is used to link a public key to a specific blockchain.
-    """
-    prefix = "GPH"
-
-    def set_prefix(self, prefix):
-        self.prefix = prefix
-
-
 class PasswordKey(Prefix):
     """ This class derives a private key given the account name, the
         role and a password. It leverages the technology of Brainkeys
         and allows people to have a secure private key by providing a
         passphrase only.
     """
     def __init__(self, account, password, role="active", prefix=None):
-        if prefix:
-            self.set_prefix(prefix)
+        self.set_prefix(prefix)
         self.account = account
         self.role = role
         self.password = password
@@ -77,8 +67,7 @@ class BrainKey(Prefix):
     """
 
     def __init__(self, brainkey=None, sequence=0, prefix=None):
-        if prefix:
-            self.set_prefix(prefix)
+        self.set_prefix(prefix)
         if not brainkey:
             self.brainkey = BrainKey.suggest()
         else:
@@ -168,8 +157,7 @@ class Address(Prefix):
 
     """
     def __init__(self, address, prefix=None):
-        if prefix:
-            self.set_prefix(prefix)
+        self.set_prefix(prefix)
         self._address = Base58(address, prefix=self.prefix)
 
     @classmethod
@@ -250,8 +238,7 @@ class PublicKey(Prefix):
 
     """
     def __init__(self, pk, prefix=None):
-        if prefix:
-            self.set_prefix(prefix)
+        self.set_prefix(prefix)
         if isinstance(pk, PublicKey):
             pk = format(pk, self.prefix)
 
@@ -398,6 +385,7 @@ class PrivateKey(Prefix):
 
     """
     def __init__(self, wif=None, prefix=None):
+        self.set_prefix(prefix)
         if wif is None:
             import os
             self._wif = Base58(hexlify(os.urandom(32)).decode('ascii'))
@@ -407,7 +395,6 @@ def __init__(self, wif=None, prefix=None):
             self._wif = wif
         else:
             self._wif = Base58(wif)
-        self.set_prefix(prefix)
 
         # test for valid key by trying to obtain a public key
         assert len(repr(self._wif)) == 64

graphenebase/base58.py

@@ -4,24 +4,15 @@
 
 from binascii import hexlify, unhexlify
 from .utils import _bytes
+from .prefix import Prefix
 
 log = logging.getLogger(__name__)
 
 """ Default Prefix """
 PREFIX = "GPH"
 
-known_prefixes = [
-    PREFIX,
-    "BTS",
-    "MUSE",
-    "TEST",
-    "STM",
-    "GLX",
-    "GLS",
-]
 
-
-class Base58(object):
+class Base58(Prefix):
     """Base58 base class
 
     This class serves as an abstraction layer to deal with base58 encoded
@@ -46,8 +37,8 @@ class Base58(object):
         * etc.
 
     """
-    def __init__(self, data, prefix=PREFIX):
-        self._prefix = prefix
+    def __init__(self, data, prefix=None):
+        self.set_prefix(prefix)
         if isinstance(data, Base58):
             data = repr(data)
         if all(c in string.hexdigits for c in data):
@@ -58,8 +49,8 @@ def __init__(self, data, prefix=PREFIX):
             raise NotImplementedError(
                 "Private Keys starting with L or K are not supported!"
             )
-        elif data[:len(self._prefix)] == self._prefix:
-            self._hex = gphBase58CheckDecode(data[len(self._prefix):])
+        elif data[:len(self.prefix)] == self.prefix:
+            self._hex = gphBase58CheckDecode(data[len(self.prefix):])
         else:
             raise ValueError("Error loading Base58 object")
 
@@ -77,13 +68,7 @@ def __format__(self, _format):
             return base58encode(self._hex)
         elif _format.upper() == "BTC":
             return base58CheckEncode(0x00, self._hex)
-        elif _format.upper() in known_prefixes:
-            return _format.upper() + str(self)
         else:
-            log.warning(
-                "Format {} unkown. You've been warned!\n".format(
-                    _format
-                ))
             return _format.upper() + str(self)
 
     def __repr__(self):

graphenebase/prefix.py

@@ -0,0 +1,9 @@
+class Prefix:
+    """ This class is meant to allow changing the prefix.
+        The prefix is used to link a public key to a specific blockchain.
+    """
+    prefix = "GPH"
+
+    def set_prefix(self, prefix):
+        if prefix:
+            self.prefix = prefix

tests/test_crypto.py

@@ -114,6 +114,22 @@ def test_btsprivkey(self):
                             'BTS5tTDDR6M3mkcyVv16edsw8dGUyNQZrvKU',
                             'BTS8G9ATJbJewVjTgTGmLGLNe1uP5XDWzaKX'])
 
+    def test_btsprivkey_change_prefix(self):
+        class P(PrivateKey):
+            prefix = "GGG"
+
+        self.assertEqual([str(P(Base58("5HqUkGuo62BfcJU5vNhTXKJRXuUi9QSE6jp8C3uBJ2BVHtB8WSd")).address),
+                          str(P("5JWcdkhL3w4RkVPcZMdJsjos22yB5cSkPExerktvKnRNZR5gx1S").address),
+                          str(P("5HvVz6XMx84aC5KaaBbwYrRLvWE46cH6zVnv4827SBPLorg76oq").address),
+                          str(P("5Jete5oFNjjk3aUMkKuxgAXsp7ZyhgJbYNiNjHLvq5xzXkiqw7R").address),
+                          str(P("5KDT58ksNsVKjYShG4Ls5ZtredybSxzmKec8juj7CojZj6LPRF7").address)
+                          ], [
+                            'GGGGu2U7Q3rmkCUCkQH2SToLMjEVUr86GrpA',
+                            'GGG9YgTfC8EfkgDG7DoRXJpMVKRougo64Lop',
+                            'GGGBXqRucGm7nRkk6jm7BNspTJTWRtNcx7k5',
+                            'GGG5tTDDR6M3mkcyVv16edsw8dGUyNQZrvKU',
+                            'GGG8G9ATJbJewVjTgTGmLGLNe1uP5XDWzaKX'])
+
     def test_uncompressed_privatekey(self):
         self.assertEqual([format(PrivateKey(Base58("5HqUkGuo62BfcJU5vNhTXKJRXuUi9QSE6jp8C3uBJ2BVHtB8WSd")).uncompressed, "BTS"),
                           format(PrivateKey("5JWcdkhL3w4RkVPcZMdJsjos22yB5cSkPExerktvKnRNZR5gx1S").uncompressed, "BTS"),