mirror of
https://github.com/fuckpiracyshield/api.git
synced 2024-05-20 05:56:19 +02:00
104 lines
3.9 KiB
Python
104 lines
3.9 KiB
Python
import os
|
|
import json
|
|
import tornado
|
|
|
|
from ioutils.protected import ProtectedHandler
|
|
from ioutils.errors import ErrorCode, ErrorMessage
|
|
from piracyshield_service.forensic.create_archive import ForensicCreateArchiveService
|
|
from piracyshield_component.environment import Environment
|
|
from piracyshield_component.exception import ApplicationException
|
|
|
|
class UploadForensicHandler(ProtectedHandler):
|
|
|
|
def get_received_chunks_info_path(self, filename):
|
|
return f"{Environment.CACHE_PATH}/{filename}_info.json"
|
|
|
|
def update_received_chunks_info(self, filename, chunk_index):
|
|
info_path = self.get_received_chunks_info_path(filename)
|
|
|
|
if os.path.exists(info_path):
|
|
with open(info_path, 'r') as file:
|
|
info = json.load(file)
|
|
|
|
received_chunks = set(info['received_chunks'])
|
|
else:
|
|
received_chunks = set()
|
|
|
|
received_chunks.add(chunk_index)
|
|
|
|
with open(info_path, 'w') as file:
|
|
json.dump({"received_chunks": list(received_chunks)}, file)
|
|
|
|
def check_all_chunks_received(self, filename, total_chunks):
|
|
info_path = self.get_received_chunks_info_path(filename)
|
|
|
|
if os.path.exists(info_path):
|
|
with open(info_path, 'r') as file:
|
|
info = json.load(file)
|
|
|
|
return len(info['received_chunks']) == total_chunks
|
|
|
|
return False
|
|
|
|
async def post(self, ticket_id):
|
|
if not self.initialize_account():
|
|
return
|
|
|
|
# verify permissions
|
|
self.permission_service.can_upload_ticket()
|
|
|
|
if 'archive' not in self.request.files or 'chunkIndex' not in self.request.arguments or 'totalChunks' not in self.request.arguments:
|
|
return self.error(status_code = 400, error_code = ErrorCode.MISSING_FILE, message = ErrorMessage.MISSING_FILE)
|
|
|
|
try:
|
|
archive_handler = self.request.files['archive'][0]
|
|
chunk_index = int(self.request.arguments['chunkIndex'][0])
|
|
total_chunks = int(self.request.arguments['totalChunks'][0])
|
|
original_filename = self.get_argument('originalFileName')
|
|
|
|
filename = f"{ticket_id}_{original_filename}"
|
|
|
|
temp_storage_path = f"{Environment.CACHE_PATH}/{filename}"
|
|
|
|
chunk_temp_path = temp_storage_path + f"_part{chunk_index}"
|
|
|
|
with open(chunk_temp_path, "wb") as temp_file:
|
|
temp_file.write(archive_handler['body'])
|
|
|
|
# update received chunks
|
|
self.update_received_chunks_info(filename, chunk_index)
|
|
|
|
# did we get every chunk?
|
|
if self.check_all_chunks_received(filename, total_chunks):
|
|
# proceed to assemble all the chunks into the a single file
|
|
with open(temp_storage_path, "wb") as final_file:
|
|
for i in range(total_chunks):
|
|
part_file_path = temp_storage_path + f"_part{i}"
|
|
|
|
if not os.path.exists(part_file_path):
|
|
raise ApplicationException("Missing file chunk: " + str(i))
|
|
|
|
with open(part_file_path, "rb") as part_file:
|
|
final_file.write(part_file.read())
|
|
|
|
os.remove(part_file_path)
|
|
|
|
# process the final file
|
|
forensic_create_archive_service = ForensicCreateArchiveService()
|
|
|
|
await tornado.ioloop.IOLoop.current().run_in_executor(
|
|
None,
|
|
forensic_create_archive_service.execute,
|
|
ticket_id,
|
|
filename
|
|
)
|
|
|
|
# clean stored info
|
|
os.remove(self.get_received_chunks_info_path(filename))
|
|
|
|
# always answer or the frontend won't send anything more
|
|
self.success()
|
|
|
|
except ApplicationException as e:
|
|
self.error(status_code = 400, error_code = e.code, message = e.message)
|