教育用Botnet演示程序：WebSocket与用户输入冲突及退出异常问题

问题描述 #

我正在开发一款用于教育演示（视频制作、展示Botnet工作原理）的Botnet控制端程序，但遇到两个核心问题：

• WebSocket服务器启动后程序卡住，无法接收用户输入
• 无法通过KeyboardInterrupt（Ctrl+C）正常退出，必须通过任务管理器终止Python进程

原代码 #

主程序代码 #

from colorama import Fore, Style
import socket
import title
import os
import asyncio
import websockets
from aioconsole import ainput

connections = 0
started = False
started_lock = asyncio.Lock()

# ----------- TITLE ----------- #
print(title.generate_title())

# ----------- COMMANDS ----------- #
hostname = socket.gethostname()
ip = socket.gethostbyname(hostname)

async def handler(websocket, path):
    global connections
    connections += 1

    try:
        async for message in websocket:
            pass
    except websockets.exceptions.ConnectionClosed:
        connections -= 1

async def start_websocket_server():
    global started

    try:
        async with started_lock:
            started = True
        start_server = await websockets.serve(handler, "localhost", 5914)
        print(f"{Fore.BLUE}[*]: {Style.RESET_ALL}Websocket server started!")
        await start_server.wait_closed()
    except Exception as e:
        print(f"{Fore.RED}[ERR]: {Style.RESET_ALL}An error occurred while starting websocket server: {e}")

def clear():
    if os.name == 'nt':
        os.system('cls')
    else:
        os.system('clear')

async def user_input():
    while True:
        if started:
            try:
                user_input = await ainput(f"{ip}:{hostname}> ")
                if user_input.strip() == "help":
                    print("""List of available commands:
                    connections - number of connected bots.
                    sessions - get all bots and ids
                    getinfo {id} - get info about bot by id
                    exit - exit the program""")
                elif user_input.strip() == "cls" or user_input.strip() == "clear":
                    clear()
                elif user_input.strip() == "connections":
                    print(f"{Fore.BLUE}[*]: {Style.RESET_ALL}There are currently connected {connections} bots.")
                elif user_input.strip() == "exit":
                    print("Exiting...")
                    return
                else:
                    print(f"Unknown command '{user_input}'\nView all commands using 'help'")
            except EOFError:
                print("EOFError. Exiting...")
                break
            except KeyboardInterrupt:
                print("User interrupted. Exiting...")
                break
            except Exception as e:
                print(f"An error occurred: {e}")

async def main():
    print(f"{Fore.BLUE}[*]: {Style.RESET_ALL}Starting websocket server...")

    # Start WebSocket server
    await start_websocket_server()

    # Start user input
    await user_input()

if __name__ == "__main__":
    asyncio.run(main())

依赖title库代码 #

from pyfiglet import Figlet
from colorama import Fore, Style
import random

title_fonts = ['cybermedium', 'rectangles', 'cyberlarge', '3-d', 'banner', 'banner3', 'banner4', 'chunky', 'colossal', 'computer', 'cosmic', 'crawford', 'cricket', 'doom', 'epic', 'poison']

try:
    import pyfiglet
except ImportError:
    import subprocess
    subprocess.run(['pip', 'install', 'pyfiglet'])
    import pyfiglet

# Generate and return the title with a random readable font
def generate_title():
    selected_font = random.choice(title_fonts)
    if selected_font == 'cosmic':
        fig = Figlet(font=selected_font, width=120)
        title = fig.renderText('Control Hub').strip()
    else:
        fig = Figlet(font=selected_font, width=100)
        title = fig.renderText('ControlHub').strip()
    return Fore.MAGENTA + title + f"\n" + Style.RESET_ALL

问题根源分析 #

1. 输入阻塞问题：原main()函数中先await start_websocket_server()，而start_server.wait_closed()会一直阻塞事件循环，导致后续的user_input()永远无法执行，必须让WebSocket服务器和用户输入任务并发运行。
2. 退出异常问题：
   • KeyboardInterrupt无法被user_input()捕获，因为事件循环被wait_closed()阻塞时无法处理信号
   • 没有正确关闭WebSocket服务器的逻辑，导致进程无法正常终止

修复后的主程序代码 #

from colorama import Fore, Style
import socket
import title
import os
import asyncio
import websockets
from aioconsole import ainput

connections = 0
started = False
started_lock = asyncio.Lock()
server = None  # 保存服务器实例用于关闭

# ----------- TITLE ----------- #
print(title.generate_title())

# ----------- COMMANDS ----------- #
hostname = socket.gethostname()
ip = socket.gethostbyname(hostname)

async def handler(websocket, path):
    global connections
    connections += 1

    try:
        async for message in websocket:
            pass
    except websockets.exceptions.ConnectionClosed:
        connections -= 1

async def start_websocket_server():
    global started, server

    try:
        async with started_lock:
            started = True
        server = await websockets.serve(handler, "localhost", 5914)
        print(f"{Fore.BLUE}[*]: {Style.RESET_ALL}Websocket server started!")
        await server.wait_closed()
    except Exception as e:
        print(f"{Fore.RED}[ERR]: {Style.RESET_ALL}An error occurred while starting websocket server: {e}")

def clear():
    if os.name == 'nt':
        os.system('cls')
    else:
        os.system('clear')

async def user_input():
    while True:
        if started:
            try:
                user_input = await ainput(f"{ip}:{hostname}> ")
                if user_input.strip() == "help":
                    print("""List of available commands:
                    connections - number of connected bots.
                    sessions - get all bots and ids
                    getinfo {id} - get info about bot by id
                    exit - exit the program""")
                elif user_input.strip() == "cls" or user_input.strip() == "clear":
                    clear()
                elif user_input.strip() == "connections":
                    print(f"{Fore.BLUE}[*]: {Style.RESET_ALL}There are currently connected {connections} bots.")
                elif user_input.strip() == "exit":
                    print("Exiting...")
                    # 关闭WebSocket服务器
                    if server:
                        server.close()
                        await server.wait_closed()
                    return
                else:
                    print(f"Unknown command '{user_input}'\nView all commands using 'help'")
            except EOFError:
                print("EOFError. Exiting...")
                if server:
                    server.close()
                    await server.wait_closed()
                break
            except KeyboardInterrupt:
                print("User interrupted. Exiting...")
                if server:
                    server.close()
                    await server.wait_closed()
                break
            except Exception as e:
                print(f"An error occurred: {e}")

async def main():
    print(f"{Fore.BLUE}[*]: {Style.RESET_ALL}Starting websocket server...")

    # 并发运行WebSocket服务器和用户输入任务
    await asyncio.gather(
        start_websocket_server(),
        user_input()
    )

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nUser interrupted. Exiting...")

关键修复说明 #

1. 并发执行任务：用asyncio.gather()同时启动WebSocket服务器和用户输入任务，解决阻塞问题，让两者能同时运行。
2. 保存服务器实例：新增server全局变量保存WebSocket服务器对象，用于在退出时主动关闭服务器。
3. 完善退出逻辑：在exit命令、EOFError、KeyboardInterrupt分支中，主动调用server.close()和await server.wait_closed()，确保服务器资源被释放，进程能正常终止。
4. 顶层异常捕获：在__main__块中捕获KeyboardInterrupt，确保即使事件循环未处理到信号时也能触发退出逻辑。

内容的提问来源于stack exchange，提问作者Andrex Youtube