|
| 1 | +import hashlib |
| 2 | +import toml |
| 3 | +from typing import List, Dict |
| 4 | + |
| 5 | +class MerkleNode: |
| 6 | + """Represents a node in the Merkle ring.""" |
| 7 | + def __init__(self, data: str): |
| 8 | + self.data = data |
| 9 | + self.hash = self.compute_hash() |
| 10 | + |
| 11 | + def compute_hash(self) -> str: |
| 12 | + """Computes the hash for the node's data.""" |
| 13 | + return hashlib.sha256(self.data.encode()).hexdigest() |
| 14 | + |
| 15 | +class MerkleRing: |
| 16 | + """Represents a Merkle ring using TOML as the storage format.""" |
| 17 | + def __init__(self): |
| 18 | + self.nodes: List[MerkleNode] = [] |
| 19 | + |
| 20 | + def add_node(self, data: str): |
| 21 | + """Adds a new node to the ring.""" |
| 22 | + node = MerkleNode(data) |
| 23 | + self.nodes.append(node) |
| 24 | + |
| 25 | + def compute_root_hash(self) -> str: |
| 26 | + """Computes the Merkle root hash from the nodes in the ring.""" |
| 27 | + hashes = [node.hash for node in self.nodes] |
| 28 | + while len(hashes) > 1: |
| 29 | + if len(hashes) % 2 == 1: # If odd number of hashes, duplicate the last one |
| 30 | + hashes.append(hashes[-1]) |
| 31 | + hashes = [ |
| 32 | + hashlib.sha256((hashes[i] + hashes[i + 1]).encode()).hexdigest() |
| 33 | + for i in range(0, len(hashes), 2) |
| 34 | + ] |
| 35 | + return hashes[0] if hashes else "" |
| 36 | + |
| 37 | + def export_to_toml(self) -> str: |
| 38 | + """Exports the Merkle ring to a TOML-formatted string.""" |
| 39 | + toml_data = { |
| 40 | + "merkle_ring": [{"data": node.data, "hash": node.hash} for node in self.nodes], |
| 41 | + "root_hash": self.compute_root_hash() |
| 42 | + } |
| 43 | + return toml.dumps(toml_data) |
| 44 | + |
| 45 | + @staticmethod |
| 46 | + def import_from_toml(toml_string: str) -> 'MerkleRing': |
| 47 | + """Creates a MerkleRing instance from a TOML-formatted string.""" |
| 48 | + data = toml.loads(toml_string) |
| 49 | + ring = MerkleRing() |
| 50 | + for node_info in data.get("merkle_ring", []): |
| 51 | + ring.add_node(node_info["data"]) |
| 52 | + return ring |
| 53 | + |
| 54 | +# Example usage: |
| 55 | +if __name__ == "__main__": |
| 56 | + # Create a Merkle ring and add some nodes |
| 57 | + ring = MerkleRing() |
| 58 | + ring.add_node("Node 1 data") |
| 59 | + ring.add_node("Node 2 data") |
| 60 | + ring.add_node("Node 3 data") |
| 61 | + |
| 62 | + # Export to TOML |
| 63 | + toml_output = ring.export_to_toml() |
| 64 | + print("Exported TOML:") |
| 65 | + print(toml_output) |
| 66 | + |
| 67 | + # Import from TOML |
| 68 | + imported_ring = MerkleRing.import_from_toml(toml_output) |
| 69 | + print("\nImported Merkle Root Hash:", imported_ring.compute_root_hash()) |
0 commit comments