39 lines
1.3 KiB
Python
39 lines
1.3 KiB
Python
|
import os
|
||
|
|
import argparse
|
||
|
|
from typing import List
|
||
|
|
from reedsolo import RSCodec
|
||
|
|
|
||
|
|
def parse_arguments():
|
||
|
|
parser = argparse.ArgumentParser(description='Reassemble file from multiple Reed-Solomon encoded pieces.')
|
||
|
|
parser.add_argument('input_files', type=str, nargs='+', help='List of files containing the encoded pieces, in order.')
|
||
|
|
return parser.parse_args()
|
||
|
|
|
||
|
|
def reassemble_file(input_files: List[str]):
|
||
|
|
# Load the encoded pieces into memory
|
||
|
|
encoded_pieces = []
|
||
|
|
for input_file in input_files[:-1]:
|
||
|
|
with open(input_file, 'rb') as f:
|
||
|
|
encoded_piece_data = f.read()
|
||
|
|
encoded_piece = bytearray(encoded_piece_data)
|
||
|
|
print(f"Read {len(encoded_piece)} bytes from {input_file}") # Print the size of the bytearray
|
||
|
|
encoded_pieces.append(encoded_piece)
|
||
|
|
|
||
|
|
# Decode the pieces using Reed-Solomon coding
|
||
|
|
n = len(encoded_pieces)
|
||
|
|
codec = RSCodec(n)
|
||
|
|
decoded_pieces = codec.decode(encoded_pieces)
|
||
|
|
|
||
|
|
# Concatenate the decoded pieces into the original file data
|
||
|
|
file_data = b''.join(decoded_pieces)
|
||
|
|
|
||
|
|
# Write the original file data to disk
|
||
|
|
output_file_path = input_files[-1]
|
||
|
|
with open(output_file_path, 'wb') as f:
|
||
|
|
f.write(file_data)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
args = parse_arguments()
|
||
|
|
reassemble_file(args.input_files)
|
||
|
|
|