Initial YakPanel commit
This commit is contained in:
Niranjan
198
mod/project/node/filetransfer/__init__.py
Normal file
198
mod/project/node/filetransfer/__init__.py
Normal file
@@ -0,0 +1,198 @@
|
||||
import os
|
||||
import time
|
||||
import threading
|
||||
from datetime import datetime
|
||||
|
||||
from .socket_server import StatusServer, StatusClient, register_cleanup
|
||||
from mod.project.node.dbutil import FileTransferDB, FileTransfer, FileTransferTask
|
||||
from mod.project.node.nodeutil import ServerNode, LocalNode, LPanelNode
|
||||
from typing import Optional, Callable, Union
|
||||
|
||||
|
||||
class Filetransfer:
|
||||
SOCKET_FILE_DIR = "/tmp/filetransfer"
|
||||
if not os.path.exists(SOCKET_FILE_DIR):
|
||||
os.mkdir(SOCKET_FILE_DIR)
|
||||
|
||||
def __init__(self, task_id: int):
|
||||
self.ft_db = FileTransferDB()
|
||||
task_data, err = self.ft_db.get_task(task_id)
|
||||
if err is None:
|
||||
raise ValueError(err)
|
||||
|
||||
self.task = FileTransferTask.from_dict(task_data)
|
||||
|
||||
file_list = self.ft_db.get_task_file_transfers(task_id)
|
||||
if not file_list:
|
||||
raise ValueError("task_id:{} file_list is empty".format(task_id))
|
||||
|
||||
self.file_map = {file_data["transfer_id"]: FileTransfer.from_dict(file_data) for file_data in file_list}
|
||||
self.file_count = len(self.file_map)
|
||||
self.file_complete = sum([1 for file in self.file_map.values() if file.status == "complete"])
|
||||
self.file_error = sum([1 for file in self.file_map.values() if file.status == "error"])
|
||||
self.count_size = sum([file.file_size for file in self.file_map.values()])
|
||||
self.complete_size = sum([file.file_size for file in self.file_map.values() if file.status == "complete"])
|
||||
self.current_file_size = 0 # 记录当前文件完成的大小
|
||||
self._srv = StatusServer(self.get_task_status, self.SOCKET_FILE_DIR + "/task_" + str(task_id))
|
||||
|
||||
if self.task.task_action == "upload":
|
||||
self.sn = LocalNode()
|
||||
if self.task.target_node["lpver"]:
|
||||
self.dn = LPanelNode(self.task.target_node["address"], self.task.target_node["api_key"],
|
||||
self.task.target_node["lpver"])
|
||||
else:
|
||||
self.dn = ServerNode(self.task.target_node["address"], self.task.target_node["api_key"],
|
||||
self.task.target_node["app_key"])
|
||||
else:
|
||||
if self.task.source_node["lpver"]:
|
||||
self.sn = LPanelNode(self.task.source_node["address"], self.task.source_node["api_key"],
|
||||
self.task.source_node["lpver"])
|
||||
else:
|
||||
self.sn = ServerNode(self.task.source_node["address"], self.task.source_node["api_key"],
|
||||
self.task.source_node["app_key"])
|
||||
self.dn = LocalNode()
|
||||
|
||||
self._close_func: Optional[Callable]= None
|
||||
|
||||
def get_task_status(self, init: bool = False) -> dict:
|
||||
task_dict = self.task.to_dict()
|
||||
task_dict.update({
|
||||
"file_count": self.file_count,
|
||||
"file_complete": self.file_complete,
|
||||
"file_error": self.file_error,
|
||||
"count_size": self.count_size,
|
||||
"complete_size": self.complete_size,
|
||||
"progress": (self.complete_size + self.current_file_size) * 100 / self.count_size if self.count_size > 0 else 0,
|
||||
})
|
||||
return {
|
||||
"task": task_dict,
|
||||
"file_status_list": [{
|
||||
"source_path": file.src_file,
|
||||
"target_path": file.dst_file,
|
||||
"status": file.status,
|
||||
"progress": file.progress,
|
||||
"log": file.message,
|
||||
} for file in self.file_map.values()],
|
||||
}
|
||||
|
||||
def start_server(self):
|
||||
t = threading.Thread(target=self._srv.start_server, args=(), daemon=True)
|
||||
t.start()
|
||||
register_cleanup(self._srv)
|
||||
def close():
|
||||
self._srv.stop()
|
||||
|
||||
self._close_func = close
|
||||
|
||||
def close(self):
|
||||
if self._close_func is None:
|
||||
return
|
||||
self._close_func()
|
||||
|
||||
def update_status(self):
|
||||
self._srv.update_status()
|
||||
|
||||
def run(self):
|
||||
self.task.status = "running"
|
||||
self.ft_db.update_task(self.task)
|
||||
self.start_server()
|
||||
|
||||
pending_list = [file for file in self.file_map.values() if file.status == "pending"]
|
||||
for file in pending_list:
|
||||
if file.is_dir > 0:
|
||||
# 空文件夹处理部分
|
||||
exits, _ = self.dn.target_file_exits(file.dst_file)
|
||||
if exits:
|
||||
file.progress = 100
|
||||
file.status = "complete"
|
||||
self.ft_db.update_file_transfer(file)
|
||||
continue
|
||||
res, err = self.dn.create_dir(path=file.dst_file)
|
||||
if err:
|
||||
file.progress = 0
|
||||
file.status = "failed"
|
||||
file.message = err
|
||||
else:
|
||||
if res.get("status",False) or res.get("status",-1) == 0:
|
||||
file.progress = 100
|
||||
file.status = "complete"
|
||||
else:
|
||||
file.progress = 0
|
||||
file.status = "failed"
|
||||
file.message = res["msg"]
|
||||
self.ft_db.update_file_transfer(file)
|
||||
continue
|
||||
|
||||
|
||||
file.status = "running"
|
||||
file.started_at = datetime.now()
|
||||
self.ft_db.update_file_transfer(file)
|
||||
|
||||
def call_log(progress, log):
|
||||
file.progress = progress
|
||||
self.current_file_size = file.file_size * progress // 100
|
||||
self.ft_db.update_file_transfer(file)
|
||||
self.update_status()
|
||||
|
||||
if self.task.task_action == "upload":
|
||||
self.ft_db.update_file_transfer(file)
|
||||
res = self.dn.upload_file(
|
||||
filename=file.src_file,
|
||||
target_path=os.path.dirname(file.dst_file),
|
||||
mode=self.task.default_mode,
|
||||
call_log=call_log,
|
||||
)
|
||||
else:
|
||||
self.ft_db.update_file_transfer(file)
|
||||
res = self.sn.download_file(
|
||||
filename=file.src_file,
|
||||
target_path=os.path.dirname(file.dst_file),
|
||||
mode=self.task.default_mode,
|
||||
call_log=call_log,
|
||||
)
|
||||
|
||||
self.current_file_size = 0
|
||||
|
||||
if res:
|
||||
file.status = "failed"
|
||||
file.message = res
|
||||
self.file_error += 1
|
||||
else:
|
||||
file.status = "complete"
|
||||
file.progress = 100
|
||||
self.file_complete += 1
|
||||
self.complete_size += file.file_size
|
||||
|
||||
self.ft_db.update_file_transfer(file)
|
||||
self.update_status()
|
||||
|
||||
if self.file_error == 0:
|
||||
self.task.status = "complete"
|
||||
else:
|
||||
self.task.status = "failed"
|
||||
self.ft_db.update_task(self.task)
|
||||
self.update_status()
|
||||
time.sleep(10)
|
||||
self.close()
|
||||
|
||||
|
||||
def run_file_transfer_task(task_id: int):
|
||||
ft = Filetransfer(task_id)
|
||||
ft.run()
|
||||
|
||||
def task_running_log(task_id: int, call_log: Callable[[Union[str,dict]], None]):
|
||||
socket_file = Filetransfer.SOCKET_FILE_DIR + "/task_" + str(task_id)
|
||||
if not os.path.exists(socket_file):
|
||||
call_log("The task status link does not exist")
|
||||
return
|
||||
s_client = StatusClient(socket_file, callback=call_log)
|
||||
s_client.connect()
|
||||
|
||||
def wait_running(task_id: int, timeout:float = 3.0) -> str:
|
||||
socket_file = Filetransfer.SOCKET_FILE_DIR + "/task_" + str(task_id)
|
||||
while not os.path.exists(socket_file):
|
||||
if timeout <= 0:
|
||||
return "Task startup timeout"
|
||||
timeout -= 0.05
|
||||
time.sleep(0.05)
|
||||
return ""
|
||||
300
mod/project/node/filetransfer/socket_server.py
Normal file
300
mod/project/node/filetransfer/socket_server.py
Normal file
@@ -0,0 +1,300 @@
|
||||
import json
|
||||
import socket
|
||||
import struct
|
||||
import sys
|
||||
import threading
|
||||
import os
|
||||
import atexit
|
||||
from typing import Callable, Any, Union, Tuple, Optional, List
|
||||
|
||||
|
||||
class StatusServer:
|
||||
def __init__(self, get_status_func: Callable[[bool], Any], server_address: Union[str, Tuple[str, int]]):
|
||||
"""
|
||||
初始化服务端
|
||||
:param get_status_func: 获取状态的函数，返回当前进程状态字典, 支持一个参数 init，
|
||||
当init为True时，表示获取初始化状态，否则为更新状态
|
||||
:param server_address: 本地套接字文件路径（Unix域）或 (host, port)（TCP）
|
||||
"""
|
||||
self.get_status_func = get_status_func
|
||||
self.server_address = server_address
|
||||
self.clients: List[socket.socket] = []
|
||||
self.lock = threading.Lock() # 线程锁
|
||||
self.running = False
|
||||
self.server_socket = None
|
||||
|
||||
def handle_client(self, client_socket):
|
||||
"""处理客户端连接"""
|
||||
try:
|
||||
# 发送初始状态
|
||||
new_status = self.get_status_func(True)
|
||||
status_bytes = json.dumps(new_status).encode() # 使用 JSON 更安全
|
||||
packed_data = len(status_bytes).to_bytes(4, 'little') + status_bytes # 添加长度头
|
||||
|
||||
# 添加到客户端列表
|
||||
try:
|
||||
# 分块发送
|
||||
client_socket.sendall(packed_data) # 发送结束标志
|
||||
except Exception as e:
|
||||
print(f"Failed to send update to client: {e}")
|
||||
client_socket.close()
|
||||
return
|
||||
|
||||
with self.lock:
|
||||
self.clients.append(client_socket)
|
||||
|
||||
# 保持连接以支持后续更新
|
||||
while self.running:
|
||||
try:
|
||||
# 可选：接收客户端心跳或命令
|
||||
data = client_socket.recv(1024)
|
||||
if not data:
|
||||
break
|
||||
except:
|
||||
break
|
||||
|
||||
finally:
|
||||
# 关闭连接并从列表中移除
|
||||
client_socket.close()
|
||||
with self.lock:
|
||||
if client_socket in self.clients:
|
||||
self.clients.remove(client_socket)
|
||||
|
||||
def start_server(self):
|
||||
"""启动本地套接字服务端"""
|
||||
self.running = True
|
||||
|
||||
if isinstance(self.server_address, str):
|
||||
# Unix 域套接字
|
||||
self.server_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
try:
|
||||
os.unlink(self.server_address)
|
||||
except OSError:
|
||||
if os.path.exists(self.server_address):
|
||||
raise
|
||||
self.server_socket.bind(self.server_address)
|
||||
else:
|
||||
# TCP 套接字
|
||||
self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
self.server_socket.bind(self.server_address)
|
||||
|
||||
self.server_socket.listen(5)
|
||||
print(f"Server is listening on {self.server_address}...")
|
||||
|
||||
try:
|
||||
self.running = True
|
||||
while self.running:
|
||||
client_socket, _ = self.server_socket.accept()
|
||||
print("Client connected")
|
||||
|
||||
# 启动新线程处理客户端
|
||||
client_thread = threading.Thread(target=self.handle_client, args=(client_socket,))
|
||||
client_thread.start()
|
||||
except KeyboardInterrupt:
|
||||
print("Shutting down server...")
|
||||
finally:
|
||||
self.stop()
|
||||
|
||||
def stop(self):
|
||||
"""停止服务端并清理资源"""
|
||||
if not self.running:
|
||||
return
|
||||
self.running = False
|
||||
|
||||
with self.lock:
|
||||
for client in self.clients:
|
||||
client.close()
|
||||
self.clients.clear()
|
||||
|
||||
if self.server_socket:
|
||||
self.server_socket.close()
|
||||
self.server_socket = None
|
||||
|
||||
# 清理 Unix 套接字文件
|
||||
if isinstance(self.server_address, str) and os.path.exists(self.server_address):
|
||||
os.remove(self.server_address)
|
||||
print(f"Socket file removed: {self.server_address}")
|
||||
|
||||
def update_status(self, update_data: Optional[dict]=None):
|
||||
"""获取最新的状态并推送给所有客户端"""
|
||||
if not update_data:
|
||||
new_status = self.get_status_func(False)
|
||||
else:
|
||||
new_status = update_data
|
||||
status_bytes = json.dumps(new_status).encode() # 使用 JSON 更安全
|
||||
packed_data = len(status_bytes).to_bytes(4, 'little') + status_bytes # 添加长度头
|
||||
|
||||
with self.lock:
|
||||
for client in self.clients:
|
||||
print("Sending update to client...")
|
||||
print(len(status_bytes), status_bytes, packed_data)
|
||||
try:
|
||||
client.sendall(packed_data) # 直接发送完整数据
|
||||
except Exception as e:
|
||||
print(f"Failed to send update to client: {e}")
|
||||
client.close()
|
||||
if client in self.clients:
|
||||
self.clients.remove(client)
|
||||
|
||||
|
||||
class StatusClient:
|
||||
def __init__(self, server_address, callback=None):
|
||||
"""
|
||||
初始化客户端
|
||||
:param server_address: Unix 域路径（字符串） 或 TCP 地址元组 (host, port)
|
||||
:param callback: 接收到状态更新时的回调函数，接受一个 dict 参数
|
||||
"""
|
||||
self.server_address = server_address
|
||||
self.callback = callback
|
||||
self.sock: Optional[socket.socket] = None
|
||||
self.running = False
|
||||
self.receive_thread = None
|
||||
|
||||
def connect(self):
|
||||
"""连接到服务端"""
|
||||
if isinstance(self.server_address, str):
|
||||
print("Connecting to Unix socket...", self.server_address)
|
||||
# Unix 域套接字
|
||||
self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
self.sock.connect(self.server_address)
|
||||
else:
|
||||
# TCP 套接字
|
||||
self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
self.sock.connect(self.server_address)
|
||||
|
||||
print("Connected to server.")
|
||||
|
||||
# 启动接收线程
|
||||
self.running = True
|
||||
self.receive_thread = threading.Thread(target=self._receive_loop, daemon=True)
|
||||
self.receive_thread.start()
|
||||
|
||||
def _receive_loop(self):
|
||||
buffer = b''
|
||||
while self.running:
|
||||
try:
|
||||
# 读取长度头
|
||||
while len(buffer) < 4:
|
||||
data = self.sock.recv(4)
|
||||
if not data:
|
||||
raise ConnectionResetError("Server closed the connection")
|
||||
buffer += data
|
||||
length = int.from_bytes(buffer[:4], 'little')
|
||||
buffer = buffer[4:]
|
||||
|
||||
# 读取完整数据
|
||||
while len(buffer) < length:
|
||||
data = self.sock.recv(length - len(buffer))
|
||||
if not data:
|
||||
raise ConnectionResetError("Server closed the connection")
|
||||
buffer += data
|
||||
message = buffer[:length]
|
||||
buffer = buffer[length:]
|
||||
|
||||
# 解析JSON
|
||||
status = json.loads(message.decode())
|
||||
if self.callback:
|
||||
self.callback(status)
|
||||
except ConnectionResetError as e:
|
||||
print("Connection interrupted:", e)
|
||||
self.disconnect()
|
||||
break
|
||||
except json.JSONDecodeError as e:
|
||||
print("JSON parsing failed:", e)
|
||||
continue
|
||||
except Exception as e:
|
||||
print("reception error:", e)
|
||||
self.disconnect()
|
||||
break
|
||||
|
||||
def disconnect(self):
|
||||
"""断开连接"""
|
||||
self.running = False
|
||||
if self.sock:
|
||||
self.sock.close()
|
||||
self.sock = None
|
||||
print("Disconnected from server.")
|
||||
|
||||
def stop(self):
|
||||
"""停止客户端"""
|
||||
self.disconnect()
|
||||
if self.receive_thread and self.receive_thread.is_alive():
|
||||
self.receive_thread.join()
|
||||
print("Client stopped.")
|
||||
|
||||
def wait_receive(self):
|
||||
if self.receive_thread and self.receive_thread.is_alive():
|
||||
self.receive_thread.join()
|
||||
|
||||
|
||||
# 注册退出清理钩子
|
||||
def register_cleanup(server_instance):
|
||||
def cleanup():
|
||||
server_instance.stop()
|
||||
|
||||
atexit.register(cleanup)
|
||||
|
||||
# # 示例使用
|
||||
# if __name__ == '__main__' and "server" in sys.argv:
|
||||
#
|
||||
# import time
|
||||
#
|
||||
# # 模拟的状态存储
|
||||
# process_status = {
|
||||
# 'process1': 'running',
|
||||
# 'process2': 'stopped',
|
||||
# "big_data": "<AAAAAAAAAAAAAAAAAFFFFFFFFFFFFFFFFFFAAAAFFFFFFFFFFFFFFFAAAAAAAAAAAAAAAAAAAAAAAAA>"
|
||||
# }
|
||||
#
|
||||
# def get_status():
|
||||
# return process_status
|
||||
#
|
||||
# # 设置 Unix 域套接字地址
|
||||
# server_address = './socket_filetransfer.sock'
|
||||
#
|
||||
# # 创建服务端实例
|
||||
# server = StatusServer(get_status, server_address)
|
||||
# register_cleanup(server) # 注册退出时清理
|
||||
#
|
||||
# # 启动服务端线程
|
||||
# server_thread = threading.Thread(target=server.start_server)
|
||||
# server_thread.daemon = True
|
||||
# server_thread.start()
|
||||
#
|
||||
# # 模拟状态更新
|
||||
# try:
|
||||
# while True:
|
||||
# print(">>>>>>>change<<<<<<<<<<<<<<<")
|
||||
# time.sleep(5)
|
||||
# process_status['process1'] = 'stopped' if process_status['process1'] == 'running' else 'running'
|
||||
# server.update_status()
|
||||
#
|
||||
# time.sleep(5)
|
||||
# process_status['process2'] = 'running' if process_status['process2'] == 'stopped' else 'stopped'
|
||||
# server.update_status()
|
||||
# except KeyboardInterrupt:
|
||||
# pass
|
||||
#
|
||||
# # 示例使用
|
||||
# if __name__ == '__main__' and "client" in sys.argv:
|
||||
# # Unix 域套接字示例：
|
||||
# server_address = './socket_filetransfer.sock'
|
||||
#
|
||||
# # 示例回调函数
|
||||
# def on_status_update(status_dict):
|
||||
# print("[Callback] New status received:")
|
||||
# for k, v in status_dict.items():
|
||||
# print(f" - {k}: {v}")
|
||||
#
|
||||
#
|
||||
# client = StatusClient(server_address, callback=on_status_update)
|
||||
# try:
|
||||
# client.connect()
|
||||
#
|
||||
# # 主线程保持运行，防止程序退出
|
||||
# while client.running:
|
||||
# pass
|
||||
# except KeyboardInterrupt:
|
||||
# print("Client shutting down...")
|
||||
# client.stop()
|
||||
Reference in New Issue
Block a user