125 lines
4.4 KiB
Python
125 lines
4.4 KiB
Python
import json
|
|
from quart import Blueprint, jsonify, request, abort, current_app
|
|
from werkzeug.exceptions import HTTPException
|
|
from enum import Enum
|
|
from internal.types import ActiveClient, NodeCommands
|
|
|
|
nodes_bp = Blueprint('nodes', __name__)
|
|
|
|
|
|
async def register_client(websocket, client_id):
|
|
"""Registers a new client connection."""
|
|
current_app.active_clients[client_id] = ActiveClient(websocket)
|
|
print(f"Client {client_id} connected.")
|
|
|
|
|
|
async def unregister_client(client_id):
|
|
"""Unregisters a disconnected client."""
|
|
if client_id in current_app.active_clients:
|
|
del current_app.active_clients[client_id]
|
|
print(f"Client {client_id} disconnected.")
|
|
|
|
|
|
async def send_command_to_client(client_id, command_name, *args):
|
|
"""Sends a command to a specific client."""
|
|
if client_id in current_app.active_clients:
|
|
websocket = current_app.active_clients[client_id].websocket
|
|
message = json.dumps({"type": "command", "name": command_name, "args": args})
|
|
try:
|
|
await websocket.send(message)
|
|
print(f"Sent command '{command_name}' to client {client_id}")
|
|
except websockets.exceptions.ConnectionClosedError:
|
|
print(f"Failed to send to client {client_id}: connection closed.")
|
|
await unregister_client(client_id)
|
|
else:
|
|
print(f"Client {client_id} not found.")
|
|
|
|
|
|
async def send_command_to_all_clients(command_name, *args):
|
|
"""Sends a command to all connected clients."""
|
|
message = json.dumps({"type": "command", "name": command_name, "args": args})
|
|
# Use a list of items to avoid issues if clients disconnect during iteration
|
|
clients_to_send = list(current_app.active_clients.items())
|
|
for client_id, websocket in clients_to_send:
|
|
try:
|
|
await websocket.send(message)
|
|
print(f"Sent command '{command_name}' to client {client_id}")
|
|
except websockets.exceptions.ConnectionClosedError:
|
|
print(f"Failed to send to client {client_id}: connection closed.")
|
|
await unregister_client(client_id)
|
|
|
|
|
|
@nodes_bp.route("/", methods=['GET'])
|
|
async def get_nodes():
|
|
"""API endpoint to list currently connected client IDs."""
|
|
return jsonify(list(current_app.active_clients.keys()))
|
|
|
|
|
|
@nodes_bp.route("/join", methods=['POST'])
|
|
async def join():
|
|
"""
|
|
Send a join command to the specific system specified
|
|
"""
|
|
data = await request.get_json()
|
|
|
|
client_id = data.get("client_id")
|
|
system_id = data.get("system_id")
|
|
guild_id = data.get("guild_id")
|
|
channel_id = data.get("channel_id")
|
|
|
|
# Check to make sure the client is online
|
|
if client_id not in current_app.active_clients:
|
|
return jsonify({"error": f"Client {client_id} not found, it might be offline"}), 404
|
|
|
|
try:
|
|
args = [system_id, guild_id, channel_id]
|
|
|
|
if not isinstance(args, list):
|
|
return jsonify({"error": "'args' must be a list"}), 400
|
|
|
|
# Send the command asynchronously
|
|
await send_command_to_client(client_id, NodeCommands.JOIN, *args)
|
|
|
|
return jsonify({"status": "command sent", "client_id": client_id, "command": NodeCommands.JOIN}), 200
|
|
|
|
except Exception as e:
|
|
return jsonify({"error": f"Failed to send command: {e}"}), 500
|
|
|
|
|
|
@nodes_bp.route("/leave", methods=['POST'])
|
|
async def leave():
|
|
"""
|
|
Send a join command to the specific system specified
|
|
"""
|
|
data = await request.get_json()
|
|
|
|
client_id = data.get("client_id")
|
|
guild_id = data.get("guild_id")
|
|
|
|
# Check to make sure the client is online
|
|
if client_id not in current_app.active_clients:
|
|
return jsonify({"error": f"Client {client_id} not found, it might be offline"}), 404
|
|
|
|
try:
|
|
args = [guild_id]
|
|
|
|
if not isinstance(args, list):
|
|
return jsonify({"error": "'args' must be a list"}), 400
|
|
|
|
# Send the command asynchronously
|
|
await send_command_to_client(client_id, NodeCommands.LEAVE, *args)
|
|
|
|
return jsonify({"status": "command sent", "client_id": client_id, "command": NodeCommands.LEAVE}), 200
|
|
|
|
except Exception as e:
|
|
return jsonify({"error": f"Failed to send command: {e}"}), 500
|
|
|
|
|
|
@nodes_bp.route("/get_online_bots", methods=['GET'])
|
|
async def get_online_bots():
|
|
active_bots = []
|
|
for client in current_app.active_clients:
|
|
if current_app.active_clients[client].active_token:
|
|
active_bots.append(current_app.active_clients[client].active_token.to_dict())
|
|
|
|
return jsonify(active_bots) |