Storage recovery

main.py

import json
import plyvel
import glob
import binascii
import argparse
import re

pattern = re.compile('dump-block-([\d]+)')
STORAGE_PREFIX = b'\x70'

def main(audit_root, output_path, target_block):
    db = plyvel.DB(output_path, create_if_missing=True)
    storage = {}

    files = glob.glob(f"{audit_root}/**/*.json", recursive=True)
    sorted_files = sorted(files, key=lambda name: int(re.search(pattern, name).group(1)))

    # replay the storage changes in memory
    for filename in sorted_files:
        with open(filename,'r') as f:
            content = json.load(f)
            for block in content:
                if block['block'] <= target_block and block['size'] > 0:
                    for entry in block['storage']:
                        try:
                            storage.update({entry['key']:entry['value']})
                        except KeyError:
                            # KeyError happens when there is no 'value', thus indicating a Delete action
                            del storage[entry['key']]

    # persist to disk
    with db.write_batch() as wb:
        for k, v in storage.items():
            new_key = STORAGE_PREFIX + binascii.unhexlify(k.encode())
            new_value = binascii.unhexlify(v.encode())
            wb.put(new_key, new_value)

    db.close()

# def read_output():
#     db = plyvel.DB('./temp_output/', create_if_missing=False)
#     for k,v in db.iterator(prefix=STORAGE_PREFIX):
#         print(binascii.hexlify(k).decode(), binascii.hexlify(v).decode())

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--storage", action="store",
                        help="Absolute path to use for loading the Storage audit JSON files")
    parser.add_argument("-o", "--output", action="store",
                        help="Absolute path to use for storing the new chain")
    parser.add_argument("-b", "--target_block", action="store", type=int,
                        help="Blockheight to restore Storage to")
    args = parser.parse_args()
    if not args.storage or not args.output or not args.target_block:
        parser.print_help()
    else:
        main(args.storage, args.output, args.target_block)