Add SIGINT signal handling to wssshd.py for clean Ctrl+C exit

- Implement asyncio Event-based signal handling for WebSocket daemon
- Add graceful shutdown with proper resource cleanup
- Terminate active terminal processes cleanly on shutdown
- Clean up WebSocket connections, tunnels, and client registrations
- Provide user feedback during shutdown process
- Ensure compatibility with existing daemon functionality

wssshd.py

@@ -35,6 +35,7 @@ import fcntl
 import termios
 import stat
 import configparser
+import signal
 from flask import Flask, render_template, request, redirect, url_for, flash, jsonify, send_from_directory
 from flask_login import LoginManager, UserMixin, login_user, login_required, logout_user, current_user
 from flask_sqlalchemy import SQLAlchemy
@@ -575,6 +576,16 @@ async def main():
     global server_password
     server_password = args.password
 
+    # Set up signal handling for clean exit
+    shutdown_event = asyncio.Event()
+
+    def signal_handler(signum, frame):
+        if debug: print(f"[DEBUG] Received signal {signum}, initiating shutdown")
+        shutdown_event.set()
+
+    # Register signal handler for SIGINT (Ctrl+C)
+    signal.signal(signal.SIGINT, signal_handler)
+
     # Load certificate
     if getattr(sys, 'frozen', False):
         # Running as bundled executable
@@ -593,11 +604,13 @@ async def main():
     ws_server = await websockets.serve(handle_websocket, args.host, args.port, ssl=ssl_context)
 
     print(f"WebSocket SSH Daemon running on {args.host}:{args.port}")
+    print("Press Ctrl+C to stop the server")
 
     # Start cleanup task
     cleanup_coro = asyncio.create_task(cleanup_task())
 
     # Start web interface if specified
+    flask_thread = None
     if args.web_host and args.web_port:
         # Handle HTTPS setup
         ssl_context = None
@@ -623,10 +636,73 @@ async def main():
         flask_thread.start()
         print(f"Web interface available at {protocol}://{args.web_host}:{args.web_port}")
 
-    # Wait for server to close
+    try:
+        # Create tasks for waiting
+        server_wait_task = asyncio.create_task(ws_server.wait_closed())
+        shutdown_wait_task = asyncio.create_task(shutdown_event.wait())
+
+        # Wait for either server to close or shutdown signal
+        done, pending = await asyncio.wait(
+            [server_wait_task, shutdown_wait_task],
+            return_when=asyncio.FIRST_COMPLETED
+        )
+
+        # Cancel pending tasks
+        for task in pending:
+            task.cancel()
+
+        print("\nShutting down WebSocket SSH Daemon...")
+
+        # Close WebSocket server if it's still running
+        if not ws_server.closed:
+            ws_server.close()
             await ws_server.wait_closed()
 
         # Cancel cleanup task
+        if not cleanup_coro.done():
+            cleanup_coro.cancel()
+            try:
+                await cleanup_coro
+            except asyncio.CancelledError:
+                pass
+
+        # Clean up active terminals
+        for request_id, terminal in list(active_terminals.items()):
+            proc = terminal['proc']
+            if proc.poll() is None:
+                if debug: print(f"[DEBUG] Terminating terminal process {request_id}")
+                proc.terminate()
+                try:
+                    # Wait up to 5 seconds for process to terminate
+                    await asyncio.wait_for(
+                        asyncio.get_event_loop().run_in_executor(None, proc.wait),
+                        timeout=5.0
+                    )
+                except asyncio.TimeoutError:
+                    if debug: print(f"[DEBUG] Force killing terminal process {request_id}")
+                    proc.kill()
+                    try:
+                        await asyncio.get_event_loop().run_in_executor(None, proc.wait)
+                    except:
+                        pass
+            del active_terminals[request_id]
+
+        # Clean up active tunnels
+        for request_id in list(active_tunnels.keys()):
+            if debug: print(f"[DEBUG] Cleaning up tunnel {request_id}")
+            del active_tunnels[request_id]
+
+        # Clean up clients
+        for client_id in list(clients.keys()):
+            if debug: print(f"[DEBUG] Cleaning up client {client_id}")
+            del clients[client_id]
+
+        print("WebSocket SSH Daemon stopped cleanly")
+
+    except Exception as e:
+        print(f"Error during shutdown: {e}")
+        # Ensure cleanup happens even if there's an error
+        if not cleanup_coro.done():
             cleanup_coro.cancel()
 
 if __name__ == '__main__':