66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
import os
|
|
import shutil
|
|
import hashlib
|
|
import tempfile
|
|
import secrets
|
|
from fastapi import UploadFile
|
|
from typing import List
|
|
from .. import models
|
|
|
|
CHUNK_SIZE = 1024 * 1024 * 5 # 5MB
|
|
|
|
async def save_temp_file(file: UploadFile):
|
|
temp_dir = tempfile.mkdtemp()
|
|
file_path = os.path.join(temp_dir, file.filename)
|
|
sha256_hash = hashlib.sha256()
|
|
|
|
with open(file_path, "wb") as f:
|
|
while chunk := await file.read(CHUNK_SIZE):
|
|
f.write(chunk)
|
|
sha256_hash.update(chunk)
|
|
|
|
return file_path, sha256_hash.hexdigest()
|
|
|
|
def split_file(file_path: str, accounts: List[models.Account]):
|
|
# Simplified splitting logic for now
|
|
# In a real scenario, we'd split based on available space
|
|
parts = []
|
|
file_size = os.path.getsize(file_path)
|
|
num_accounts = len(accounts)
|
|
part_size = file_size // num_accounts
|
|
|
|
with open(file_path, "rb") as f:
|
|
for i, account in enumerate(accounts):
|
|
part_path = f"{file_path}.part{i}"
|
|
with open(part_path, "wb") as part_f:
|
|
remaining = part_size if i < num_accounts - 1 else -1
|
|
while remaining > 0 or remaining == -1:
|
|
read_size = min(CHUNK_SIZE, remaining) if remaining != -1 else CHUNK_SIZE
|
|
chunk = f.read(read_size)
|
|
if not chunk:
|
|
break
|
|
part_f.write(chunk)
|
|
if remaining != -1:
|
|
remaining -= len(chunk)
|
|
|
|
part_sha256 = hashlib.sha256(open(part_path, 'rb').read()).hexdigest()
|
|
|
|
parts.append({
|
|
"index": i,
|
|
"path": part_path,
|
|
"size": os.path.getsize(part_path),
|
|
"account": account,
|
|
"sha256": part_sha256
|
|
})
|
|
|
|
return parts
|
|
|
|
def generate_token():
|
|
return secrets.token_urlsafe(32)
|
|
|
|
def merge_files(part_paths: List[str], output_path: str):
|
|
with open(output_path, 'wb') as f_out:
|
|
for part_path in sorted(part_paths):
|
|
with open(part_path, 'rb') as f_in:
|
|
shutil.copyfileobj(f_in, f_out)
|