import pytest import os import json from pymerkle import MerkleTree # Clean exports dir before running the test for file in os.listdir(os.path.join(os.path.dirname(__file__), 'exports')): os.remove(os.path.join(os.path.dirname(__file__), 'exports', file)) tree = MerkleTree() for i in range(12): tree.update(record='{}-th record'.format(i)) export_path = os.path.join(os.path.dirname(__file__), 'exports', '{}.json'.format(tree.uuid)) tree.export(file_path=export_path) with open(export_path, 'r') as f: exported_version = json.load(f) def test_export(): assert exported_version == { "header": { "encoding": "utf_8", "hash_type": "sha256", "security": True }, "hashes": [ "a08665f5138f40a07987234ec9821e5be05ecbf5d7792cd4155c4222618029b6", "3dbbc4898d7e909de7fc7bb1c0af36feba78abc802102556e4ea52c28ccb517f", "45c44059cf0f5a447933f57d851a6024ac78b44a41603738f563bcbf83f35d20", "b5db666b0b34e92c2e6c1d55ba83e98ff37d6a98dda532b125f049b43d67f802",
class Tree: """Tree class wraps and uses the defaults for pymerkle. It accepts SHA-256 hashes, and UTF-8 encoding. For testing you might benefit from turning off preimage protection""" def __init__(self, file=None): """sets up the underlying merkle tree Parameters: file (str): Optional recovery of tree state dumped by export(). """ if file is None: self.merkle = MerkleTree( b"hello world", b"Hello world", b"hello World", b"Hello World", b"hello world!", b"Hello World!", # include enough hello worlds to construct a path raw_bytes=False, ) else: self.merkle = MerkleTree.loadFromFile(file) def export(self): """export the tree as json to a file called tree_<calendar>.json""" filestring = "tree_" + str(calendar.timegm(time.gmtime())) + ".json" logger.debug("exporting tree to " + filestring) self.merkle.export(filestring) @logger.catch def stamp(self, checksum): """this method will add the checksum to the merkle tree and return the full proof serialized. Errors: ChecksumFormatError, ChecksumExistsError""" #The checksum comes in as a string, so we've got to be sure it can be converted into bytes valid = re.match("^[A-Fa-f0-9]{64}$", checksum) if valid == None: raise ChecksumFormatError self.merkle.update(digest=checksum) logger.info("Checksum: {} added to the tree", checksum) return True return False @logger.catch def proofFor(self, checksum): """this method will not add a new checksum, but will check it exists and return the proof. Assumes that the checksum has already been looked up and isn't in the cache. Errors: ChecksumNotFoundError """ # dont send a proof it it doesn't exist exists = self.merkle.find_index(checksum) if not exists: raise ChecksumNotFoundError else: logger.info("Checksum: {} proof found in the tree", checksum) return self.merkle.auditProof(checksum, commit=True).serialize() return None def validate(self, proof): """this method will return whether or not the proof submitted is valid. Assumes proof was generated by this service, with pymerkle.toJSONString(). Note the proof could be considered valid even if it's checksum isn't in this service's merkle tree. Returns: Errors: ValidationError""" try: tmp = Proof.deserialize(proof) logger.info("The proof to validate: {}", tmp) return self.merkle.validateProof(tmp, get_receipt=True).serialize() except: raise ValidationError def get_current_root(self): return self.merkle.get_commitment() def consistency_proof(self, subhash): """Returns a consistency proof that the subhash is a valid ancestor root of the current one Returns: A pymerkle.Proof object serialized to JSON """ return self.merkle.merkleProof({"subhash": subhash}).serialize()