Initial YakPanel commit
This commit is contained in:
Niranjan
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