Skip to content

Commit

Permalink
python3
Browse files Browse the repository at this point in the history
  • Loading branch information
EderSantana committed Sep 9, 2016
1 parent 6d51ab5 commit 0b63e0d
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 15 deletions.
35 changes: 24 additions & 11 deletions merkletools.py
@@ -1,4 +1,5 @@
import hashlib
import binascii
try:
import sha3
except:
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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:
Expand All @@ -109,25 +122,25 @@ 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:
proof_hash = target_hash
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
7 changes: 3 additions & 4 deletions tests/test_merkle_tools.py
Expand Up @@ -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)
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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])
Expand All @@ -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)
Expand Down

0 comments on commit 0b63e0d

Please sign in to comment.