From 0b63e0d80df1f8e4c5f15ea34ae58dc6e545ec3e Mon Sep 17 00:00:00 2001
From: EderSantana <eder@macbook.com>
Date: Fri, 9 Sep 2016 14:50:33 -0400
Subject: [PATCH] python3

---
 merkletools.py             | 35 ++++++++++++++++++++++++-----------
 tests/test_merkle_tools.py |  7 +++----
 2 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/merkletools.py b/merkletools.py
index 565c8b6..86c25c2 100644
--- a/merkletools.py
+++ b/merkletools.py
@@ -1,4 +1,5 @@
 import hashlib
+import binascii
 try:
     import sha3
 except:
@@ -32,6 +33,12 @@ def __init__(self, hash_type="sha256"):
 
         self.reset_tree()
 
+    def _to_hex(self, x):
+        try:  # python3
+            return x.hex()
+        except:  # python2
+            return binascii.hexlify(x)
+
     def reset_tree(self):
         self.leaves = list()
         self.levels = None
@@ -43,19 +50,25 @@ def add_leaf(self, values, do_hash=False):
         if isinstance(values, tuple) or isinstance(values, list):
             for v in values:
                 if do_hash:
-                    v = self.hash_function(v.encode("utf-8")).digest()
+                    v = v.encode('utf-8')
+                    v = self.hash_function(v).hexdigest()
+                    v = bytearray.fromhex(v)
                 else:
-                    v = v.decode('hex')
+                    # v = v.decode('hex')
+                    v = bytearray.fromhex(v)
                 self.leaves.append(v)
         else:
             if do_hash:
-                v = self.hash_function(values.encode("utf-8")).digest()
+                v = values.encode("utf-8")
+                v = self.hash_function(v).hexdigest()
+                v = bytearray.fromhex(v)
             else:
-                v = values.decode('hex')
+                # v = values.decode('hex')
+                v = bytearray.fromhex(values)
             self.leaves.append(v)
 
     def get_leaf(self, index):
-        return self.leaves[index].encode('hex')
+        return self._to_hex(self.leaves[index])  # .encode('hex')
 
     def get_leaf_count(self):
         return len(self.leaves)
@@ -88,7 +101,7 @@ def make_tree(self):
     def get_merkle_root(self):
         if self.is_ready:
             if self.levels is not None:
-                return self.levels[0][0].encode('hex')
+                return self._to_hex(self.levels[0][0])  # .encode('hex')
             else:
                 return None
         else:
@@ -109,14 +122,14 @@ def get_proof(self, index):
                 is_right_node = index % 2
                 sibling_index = index - 1 if is_right_node else index + 1
                 sibling_pos = "left" if is_right_node else "right"
-                sibling_value = self.levels[x][sibling_index].encode('hex')
+                sibling_value = self._to_hex(self.levels[x][sibling_index])  # .encode('hex')
                 proof.append({sibling_pos: sibling_value})
                 index = int(index / 2.)
             return proof
 
     def validate_proof(self, proof, target_hash, merkle_root):
-        merkle_root = merkle_root.decode('hex')
-        target_hash = target_hash.decode('hex')
+        merkle_root = bytearray.fromhex(merkle_root)  # merkle_root.decode('hex')
+        target_hash = bytearray.fromhex(target_hash)  # target_hash.decode('hex')
         if len(proof) == 0:
             return target_hash == merkle_root
         else:
@@ -124,10 +137,10 @@ def validate_proof(self, proof, target_hash, merkle_root):
             for p in proof:
                 try:
                     # the sibling is a left node
-                    sibling = p['left'].decode('hex')
+                    sibling = bytearray.fromhex(p['left'])  # p['left'].decode('hex')
                     proof_hash = self.hash_function(sibling + proof_hash).digest()
                 except:
                     # the sibling is a right node
-                    sibling = p['right'].decode('hex')
+                    sibling = bytearray.fromhex(p['right'])  # p['right'].decode('hex')
                     proof_hash = self.hash_function(proof_hash + sibling).digest()
             return proof_hash == merkle_root
diff --git a/tests/test_merkle_tools.py b/tests/test_merkle_tools.py
index a75b951..89089ca 100644
--- a/tests/test_merkle_tools.py
+++ b/tests/test_merkle_tools.py
@@ -2,7 +2,6 @@
 from math import sqrt
 from merkletools import MerkleTools
 
-
 def test_add_leaf():
     mt = MerkleTools()
     mt.add_leaf("tierion", do_hash=True)
@@ -36,7 +35,7 @@ def test_get_proof():
 def test_basics():
     bLeft = 'a292780cc748697cb499fdcc8cb89d835609f11e502281dfe3f6690b1cc23dcb'
     bRight = 'cb4990b9a8936bbc137ddeb6dcab4620897b099a450ecdc5f3e86ef4b3a7135c'
-    mRoot = hashlib.sha256(bLeft.decode('hex') + bRight.decode('hex')).hexdigest()
+    mRoot = hashlib.sha256(bytearray.fromhex(bLeft) + bytearray.fromhex(bRight)).hexdigest()
 
     # tree with no leaves
     mt = MerkleTools()
@@ -99,7 +98,7 @@ def test_unhashed_leaves():
 def test_md5_tree():
     bLeftmd5 = '0cc175b9c0f1b6a831c399e269772661'
     bRightmd5 = '92eb5ffee6ae2fec3ad71c777531578f'
-    mRootmd5 = hashlib.md5(bLeftmd5.decode('hex')+bRightmd5.decode('hex')).hexdigest()
+    mRootmd5 = hashlib.md5(bytearray.fromhex(bLeftmd5)+bytearray.fromhex(bRightmd5)).hexdigest()
 
     mt = MerkleTools('md5')
     mt.add_leaf([bLeftmd5, bRightmd5])
@@ -110,7 +109,7 @@ def test_md5_tree():
 def test_proof_nodes():
     bLeft = 'a292780cc748697cb499fdcc8cb89d835609f11e502281dfe3f6690b1cc23dcb'
     bRight = 'cb4990b9a8936bbc137ddeb6dcab4620897b099a450ecdc5f3e86ef4b3a7135c'
-    mRoot = hashlib.sha256(bLeft.decode('hex') + bRight.decode('hex')).hexdigest()
+    mRoot = hashlib.sha256(bytearray.fromhex(bLeft) + bytearray.fromhex(bRight)).hexdigest()
 
     mt = MerkleTools()
     mt.add_leaf(bLeft)