merkle-tree-stream/merkle_tree_stream/generator.py

107 lines
2.8 KiB
Python
Raw Normal View History

"""The merkle tree stream generator."""
from typing import Any, Callable, List, Optional
import attr
from flat_tree import FlatTreeAccessor
Hash = str
__all__ = ['MerkleTreeGenerator', 'MerkleTreeNode']
flat_tree = FlatTreeAccessor()
@attr.s(auto_attribs=True)
class MerkleTreeNode:
"""A node in a merkle tree.
:param index: The index of node
:param parent: The parent of the node
:param size: The size of the data
:param data: The data of the node
:param hash: The hash of the data
"""
index: int
parent: int
size: int
data: Optional[bytes]
hash: Optional[str] = None
def __attrs_post_init__(self) -> Any:
"""Initialise the parent index."""
self.parent = flat_tree.parent(self.index)
@attr.s(auto_attribs=True)
class MerkleTreeGenerator:
"""A stream that generates a merkle tree based on the incoming data.
:param leaf: The leaf hash generation function
:param parent: The parent hash generation function
:param roots: The tree roots
"""
leaf: Callable[[MerkleTreeNode, List[MerkleTreeNode]], Hash]
parent: Callable[[MerkleTreeNode, List[MerkleTreeNode]], Hash]
roots: List[MerkleTreeNode] = attr.Factory(list)
def next(
self, data: bytes, nodes: Optional[List[MerkleTreeNode]] = None
) -> List[MerkleTreeNode]:
"""Further generate the tree based on the incoming data.
:param data: Incoming data
:param nodes: Pre-existing nodes
"""
nodes = nodes or []
index = 2 * (self.blocks + 1)
leaf_node = MerkleTreeNode(
index=index,
parent=flat_tree.parent(index),
hash=None,
data=data,
size=len(data),
)
leaf_node.hash = self.leaf(leaf_node, self.roots)
self.roots.append(leaf_node)
nodes.append(leaf_node)
while len(self.roots) > 1:
left = self.roots[len(self.roots) - 2]
right = self.roots[len(self.roots) - 1]
if left.parent != right.parent:
break
self.roots.pop()
new_node = MerkleTreeNode(
index=left.parent,
parent=flat_tree.parent(left.parent),
hash=self.parent(left, [right]),
size=left.size + right.size,
data=None,
)
self.roots[len(self.roots) - 1] = new_node
nodes.append(new_node)
return nodes
def __attrs_post_init__(self) -> Any:
"""Initialise parent and block defaults."""
index = self.roots[len(self.roots) - 1].index
right_span = flat_tree.right_span(index)
self.blocks = (1 + (right_span / 2)) if self.roots else 0
for root in self.roots:
if not root.parent:
root.parent = flat_tree.parent(root.index)