#!/usr/bin/python3
import base64
import glob
import struct
import traceback


files_to_scan = '/mountpoint/lost+found/**'


def from_base128(data, cursor):
    numbers = []
    while True:
        number = data[cursor]
        cursor += 1
        numbers.append(number & 127)
        if number & 128 == 0:
            break
    result = 0
    for number in reversed(numbers):
        result = result * 128 + number
    return result, cursor


def iter_protobuf(data):  # https://developers.google.com/protocol-buffers/docs/encoding
    cursor = 0
    while cursor < len(data):
        key = data[cursor]
        cursor += 1

        wire_type = key & 7
        field_number = key >> 3

        if wire_type == 0:  # varint
            value, cursor = from_base128(data, cursor)
            yield 'varint', field_number, value
        elif wire_type == 2:  # length-delimited
            length, cursor = from_base128(data, cursor)
            value = data[cursor:cursor + length]
            cursor += length
            yield 'length-delimited', field_number, value
        else:
            print(f'wire_type = {wire_type}, field_number = {field_number}; unknown wire_type, sorry')
            break


def id_to_string(v):  # https://github.com/storj/storj/blob/bd36a41a9ebe4d855ad654daf3856728e2796d98/storage/filestore/dir.go#L37
    return base64.b32encode(v).lower().strip(b'=').decode('ascii')


def extract(file_name):
    with open(file_name, 'rb') as fh:
        data = fh.read(512)  # V1PieceHeaderReservedArea, https://github.com/storj/storj/blob/bd36a41a9ebe4d855ad654daf3856728e2796d98/storagenode/pieces/readwrite.go#L259

    header_size, = struct.unpack('>h', data[:2])
    assert header_size <= 510

    for wire_type, field_number, value in iter_protobuf(data[2:2 + header_size]):
        if field_number == 1:  # https://github.com/storj/proto/blob/c713232da63d9e8221f057591edd7fbf91909e33/pkg/pb/piecestore2.proto#L119
            assert wire_type == 'varint'
            assert value == 1
        if field_number == 5:  # https://github.com/storj/proto/blob/c713232da63d9e8221f057591edd7fbf91909e33/pkg/pb/piecestore2.proto#L119
            for wire_type, field_number, value in iter_protobuf(value):
                if field_number == 2:  # https://github.com/storj/proto/blob/c713232da63d9e8221f057591edd7fbf91909e33/pkg/pb/orders.proto#L63
                    assert wire_type == 'length-delimited'
                    extracted_satellite_id = id_to_string(value)
                elif field_number == 5:  # https://github.com/storj/proto/blob/c713232da63d9e8221f057591edd7fbf91909e33/pkg/pb/orders.proto#L72
                    assert wire_type == 'length-delimited'
                    extracted_piece_id = id_to_string(value)

    return extracted_satellite_id, extracted_piece_id


for file_name in glob.iglob(files_to_scan, recursive=True):
    try:
        extracted_satellite_id, extracted_piece_id = extract(file_name)
        print(f'File {file_name} should probably go do {extracted_satellite_id}/{extracted_piece_id[:2]}/{extracted_piece_id[2:]}.sj1')
    except:
        print(f'Error for file {file_name}, maybe it is not a piece file?')
        traceback.print_exc()
        print()