101 lines
4.1 KiB
Python
101 lines
4.1 KiB
Python
import sys, os
|
|
import math
|
|
from enum import Enum, auto
|
|
from pathlib import Path
|
|
import time
|
|
|
|
from libs.noSys.noSysModule import NoSysModule
|
|
from libs.app.common.logging import get_logger
|
|
from libs.app.common.paths import ROOT_DIR
|
|
from libs.fspn.utils.observable import Observable
|
|
from libs.fspn.utils.wrapper_util import threaded
|
|
from libs.fspn.utils.sha256_util import hash_bytes, hash_file
|
|
|
|
from .file import File, FileEvents
|
|
|
|
logger = get_logger()
|
|
|
|
class FileTransferEvents(Enum):
|
|
ON_RECEIVING_ = auto()
|
|
|
|
class FileTransfer(NoSysModule):
|
|
def __init__(self, nosys_core):
|
|
super().__init__(nosys_core)
|
|
self.files:dict[tuple[str,str], File] = {}
|
|
|
|
def send_file(self, file_path, connection_id, to_module):
|
|
path = Path(file_path)
|
|
file_folder = path.parent
|
|
file_name = path.name
|
|
file_size = path.stat().st_size
|
|
chunk_size = self.get_dynamic_chunk_size(file_size)
|
|
file_hash = hash_file(file_path)
|
|
file = File(file_folder, file_name, file_size, chunk_size, file_hash, connection_id, True, self, to_module)
|
|
file.subscribe_event(FileEvents.ON_FILE_APPROVED.name, self.on_file_approved)
|
|
self.files[file.id] = file
|
|
self.send_file_info(connection_id, file_name, file_size, chunk_size, file_hash, to_module)
|
|
|
|
def get_dynamic_chunk_size(self, file_size: int) -> int:
|
|
min_chunk = 64 * 1024 # 64 KB
|
|
max_chunk = 2 * 1024 * 1024 # 2 MB
|
|
target_chunks = 500
|
|
ideal_chunk = file_size // target_chunks
|
|
return max(min_chunk, min(ideal_chunk, max_chunk))
|
|
|
|
def on_module_message(self, event):
|
|
if event.meta:
|
|
action = event.meta["action"]
|
|
else:
|
|
action = event.data['action']
|
|
handler_action = getattr(self, 'on_'+action)
|
|
handler_action(event)
|
|
|
|
def send_file_info(self, connection_id, name, size, chunk_size, hash, to_module):
|
|
body = {"action":"file_info", "name":name, "size":size, "chunk_size":chunk_size, "hash":hash, "module":{"package": to_module[0], "name":to_module[1]}}
|
|
self.nosys_core.dispatcher.send_message(body, connection_id, self.id)
|
|
|
|
def on_file_info(self, event):
|
|
payload = event.data
|
|
folder = os.path.join(ROOT_DIR, "files")
|
|
to_module = (payload['module']['package'], payload['module']['name'])
|
|
file = File(folder, payload["name"], payload["size"], payload["chunk_size"], payload["hash"], event.connection.id, False, self, to_module)
|
|
file.subscribe_event(FileEvents.ON_FILE_APPROVED.name, self.on_file_approved)
|
|
self.files[file.id] = file
|
|
self.fire_event(f"{FileTransferEvents.ON_RECEIVING_.name}{to_module[0]}_{to_module[1]}", file=file)
|
|
|
|
def subscribe_module_file_events(self, module:tuple[str, str], callback):
|
|
self.subscribe_event(f"{FileTransferEvents.ON_RECEIVING_.name}{module[0]}_{module[1]}", callback)
|
|
|
|
def on_file_approved(self, event):
|
|
file:File = event.source
|
|
approved = event.approved
|
|
|
|
if file.sending:
|
|
logger.debug(f"Peer file approved {approved}")
|
|
else:
|
|
self.send_file_transfer_approved(file, approved)
|
|
|
|
def send_file_transfer_approved(self, file:File, approved):
|
|
body = {"action":"file_transfer_approved", "file_hash":file.hash, "approved":approved}
|
|
self.nosys_core.dispatcher.send_message(body, file.connection_id, self.id)
|
|
|
|
def on_file_transfer_approved(self, event):
|
|
payload = event.data
|
|
file_id = (payload["file_hash"], event.connection.id)
|
|
file:File = self.files[file_id]
|
|
file.approve_transfer(payload["approved"])
|
|
|
|
def send_file_part(self, file:File, part, part_hash, data):
|
|
meta = {"action":"file_part", "part":part, "part_hash":part_hash, "file_hash": file.hash}
|
|
self.nosys_core.dispatcher.send_binary(data, file.connection_id, self.id, meta)
|
|
|
|
def on_file_part(self, event):
|
|
meta = event.meta
|
|
data = event.data
|
|
file_id = (meta["file_hash"], event.connection.id)
|
|
file:File = self.files[file_id]
|
|
file.write_part(data, meta["part"], meta["part_hash"])
|
|
|
|
# TODO get_file_part, ack_file_part
|
|
|