64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
|
import os
|
||
|
|
import argparse
|
||
|
|
from math import ceil
|
||
|
|
from pydantic import BaseModel
|
||
|
|
from typing import List
|
||
|
|
from reedsolo import RSCodec
|
||
|
|
|
||
|
|
class Args(BaseModel):
|
||
|
|
file_path: str
|
||
|
|
n: int
|
||
|
|
|
||
|
|
def parse_arguments():
|
||
|
|
parser = argparse.ArgumentParser(description='Split file into multiple pieces using Reed-Solomon coding.')
|
||
|
|
parser.add_argument('file_path', type=str, help='Path to the input file.')
|
||
|
|
parser.add_argument('n', type=int, help='Number of pieces to split the file into.')
|
||
|
|
return Args(**vars(parser.parse_args()))
|
||
|
|
|
||
|
|
def split_file_into_pieces(file_path: str, n: int):
|
||
|
|
# Load the file into memory
|
||
|
|
with open(file_path, 'rb') as f:
|
||
|
|
file_data = f.read()
|
||
|
|
|
||
|
|
# Calculate the size of each piece
|
||
|
|
piece_size = ceil(len(file_data) / n)
|
||
|
|
|
||
|
|
# Create the folders for the pieces
|
||
|
|
folder_names = [f'piece_{i+1}' for i in range(n)]
|
||
|
|
for folder_name in folder_names:
|
||
|
|
os.makedirs(folder_name, exist_ok=True)
|
||
|
|
|
||
|
|
# Encode each piece using Reed-Solomon coding and write it to the corresponding folder
|
||
|
|
codec = RSCodec(n)
|
||
|
|
for i, folder_name in enumerate(folder_names):
|
||
|
|
piece_data = file_data[i*piece_size:(i+1)*piece_size]
|
||
|
|
encoded_data = codec.encode(piece_data)
|
||
|
|
with open(os.path.join(folder_name, f'{i+1}.dat'), 'wb') as f:
|
||
|
|
f.write(encoded_data)
|
||
|
|
|
||
|
|
def reassemble_file(folder_names: List[str], output_file_path: str):
|
||
|
|
# Load the encoded pieces into memory
|
||
|
|
encoded_pieces = []
|
||
|
|
for i, folder_name in enumerate(folder_names):
|
||
|
|
with open(os.path.join(folder_name, f'{i+1}.dat'), 'rb') as f:
|
||
|
|
encoded_piece_data = f.read()
|
||
|
|
encoded_pieces.append(encoded_piece_data)
|
||
|
|
|
||
|
|
# Decode the pieces using Reed-Solomon coding
|
||
|
|
codec = RSCodec(len(encoded_pieces))
|
||
|
|
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
|
||
|
|
with open(output_file_path, 'wb') as f:
|
||
|
|
f.write(file_data)
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
args = parse_arguments()
|
||
|
|
split_file_into_pieces(args.file_path, args.n)
|
||
|
|
folder_names = [f'piece_{i+1}' for i in range(args.n)]
|
||
|
|
reassemble_file(folder_names, 'reconstructed_file.png')
|
||
|
|
|