Remove "clean" duplicate

clean

-#!/bin/bash
-
-# WebSocket SSH Tools Clean Script
-# Clean script for removing build artifacts from WebSocket SSH tools
-#
-# Copyright (C) 2024 Stefy Lanza <stefy@nexlab.net> and SexHack.me
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program.  If not, see <https://www.gnu.org/licenses/>.
-
-# Remove PyInstaller build artifacts
-rm -rf build/
-rm -rf dist/
-rm -f *.spec
-
-# Remove C version build artifacts
-if [ -d "wssshtools" ]; then
-    cd wssshtools
-    make clean 2>/dev/null || true
-    rm -f Makefile
-    rm -f configure.sh.stamp
-    rm -f man/*.1.gz 2>/dev/null || true
-    cd ..
-fi
-
-# Remove Debian packaging artifacts
-rm -f dist/wsssh-tools*.deb
-rm -f dist/wsssh-tools*.dsc
-rm -f dist/wsssh-tools*.tar.gz
-rm -f dist/wsssh-tools*.changes
-rm -f dist/wsssh-tools*.buildinfo
-rm -f wsssh-tools*.deb
-rm -f wsssh-tools*.dsc
-rm -f wsssh-tools*.tar.gz
-rm -f wsssh-tools*.changes
-rm -f wsssh-tools*.buildinfo
-
-# Remove Debian build directory and artifacts
-if [ -d "wssshtools" ]; then
-    rm -rf wssshtools/debian/wsssh-tools/
-    rm -f wssshtools/debian/files
-    rm -f wssshtools/debian/*.debhelper*
-    rm -f wssshtools/debian/*.substvars
-    rm -f wssshtools/debian/debhelper-build-stamp
-fi
-
-# Optionally remove SSL certificates (uncomment if needed)
-# rm -f cert.pem key.pem
-
-echo "Clean complete. Build artifacts removed."
-echo "Note: SSL certificates (cert.pem, key.pem) preserved. Uncomment lines in clean.sh to remove them."
\ No newline at end of file

wssshd.py

@@ -40,7 +40,7 @@ from flask_login import LoginManager, UserMixin, login_user, login_required, log
 from flask_sqlalchemy import SQLAlchemy
 from werkzeug.security import generate_password_hash, check_password_hash
 
-# Client registry: id -> websocket
+# Client registry: id -> {'websocket': ws, 'last_seen': timestamp, 'status': 'active'|'disconnected'}
 clients = {}
 # Active tunnels: request_id -> {'client_ws': ws, 'wsssh_ws': ws, 'client_id': id}
 active_tunnels = {}
@@ -49,6 +49,7 @@ active_terminals = {}
 debug = False
 server_password = None
 args = None
+import time
 
 # Flask app for web interface
 app = Flask(__name__)
@@ -73,6 +74,20 @@ class User(UserMixin, db.Model):
 def load_user(user_id):
     return db.session.get(User, int(user_id))
 
+def cleanup_expired_clients():
+    """Remove clients that have been disconnected for more than 30 seconds"""
+    current_time = time.time()
+    expired_clients = []
+
+    for client_id, client_info in clients.items():
+        if client_info['status'] == 'disconnected':
+            if current_time - client_info['last_seen'] > 30:
+                expired_clients.append(client_id)
+                if debug: print(f"[DEBUG] [WebSocket] Client {client_id} expired and removed")
+
+    for client_id in expired_clients:
+        del clients[client_id]
+
 def openpty_with_fallback():
     """Open a PTY with fallback to different device paths for systems where /dev/pty doesn't exist"""
     # First try the standard pty.openpty()
@@ -152,8 +167,15 @@ with app.app_context():
 @login_required
 def index():
     global args
+    # Get client information with status
+    client_info = {}
+    for client_id, client_data in clients.items():
+        client_info[client_id] = {
+            'status': client_data['status'],
+            'last_seen': client_data['last_seen']
+        }
     return render_template('index.html',
-                         clients=list(clients.keys()),
+                          clients=client_info,
                           websocket_port=args.port,
                           domain=args.domain)
 
@@ -241,8 +263,14 @@ def terminal(client_id):
 @app.route('/api/clients')
 @login_required
 def get_clients():
+    client_info = {}
+    for client_id, client_data in clients.items():
+        client_info[client_id] = {
+            'status': client_data['status'],
+            'last_seen': client_data['last_seen']
+        }
     return jsonify({
-        'clients': list(clients.keys()),
+        'clients': client_info,
         'count': len(clients)
     })
 
@@ -377,30 +405,44 @@ def disconnect_terminal(client_id):
 async def handle_websocket(websocket, path=None):
     try:
         async for message in websocket:
-            if debug: print(f"[DEBUG] WebSocket message received: {message[:100]}...")
+            if debug: print(f"[DEBUG] [WebSocket] Message received: {message[:100]}...")
             data = json.loads(message)
             if data.get('type') == 'register':
                 client_id = data['id']
                 client_password = data.get('password', '')
                 if client_password == server_password:
-                    clients[client_id] = websocket
+                    # Check if client was previously disconnected
+                    was_disconnected = False
+                    if client_id in clients and clients[client_id]['status'] == 'disconnected':
+                        was_disconnected = True
+                        print(f"[DEBUG] [WebSocket] Client {client_id} reconnecting (was disconnected)")
+
+                    clients[client_id] = {
+                        'websocket': websocket,
+                        'last_seen': time.time(),
+                        'status': 'active'
+                    }
+
+                    if was_disconnected:
+                        print(f"Client {client_id} reconnected")
+                    else:
                         print(f"Client {client_id} registered")
                     await websocket.send(json.dumps({"type": "registered", "id": client_id}))
                 else:
-                    print(f"Client {client_id} registration failed: invalid password")
+                    print(f"[DEBUG] [WebSocket] Client {client_id} registration failed: invalid password")
                     await websocket.send(json.dumps({"type": "registration_error", "error": "Invalid password"}))
             elif data.get('type') == 'tunnel_request':
                 client_id = data['client_id']
                 request_id = data['request_id']
-                if client_id in clients:
+                if client_id in clients and clients[client_id]['status'] == 'active':
                     # Store tunnel mapping
                     active_tunnels[request_id] = {
-                        'client_ws': clients[client_id],
+                        'client_ws': clients[client_id]['websocket'],
                         'wsssh_ws': websocket,
                         'client_id': client_id
                     }
                     # Forward tunnel request to client
-                    await clients[client_id].send(json.dumps({
+                    await clients[client_id]['websocket'].send(json.dumps({
                         "type": "tunnel_request",
                         "request_id": request_id
                     }))
@@ -412,7 +454,7 @@ async def handle_websocket(websocket, path=None):
                     await websocket.send(json.dumps({
                         "type": "tunnel_error",
                         "request_id": request_id,
-                        "error": "Client not registered"
+                        "error": "Client not registered or disconnected"
                     }))
             elif data.get('type') == 'tunnel_data':
                 # Forward tunnel data using active tunnel mapping
@@ -420,11 +462,14 @@ async def handle_websocket(websocket, path=None):
                 if request_id in active_tunnels:
                     tunnel = active_tunnels[request_id]
                     # Forward to client
+                    if tunnel['client_id'] in clients and clients[tunnel['client_id']]['status'] == 'active':
                         await tunnel['client_ws'].send(json.dumps({
                             "type": "tunnel_data",
                             "request_id": request_id,
                             "data": data['data']
                         }))
+                    else:
+                        if debug: print(f"[DEBUG] [WebSocket] Cannot forward tunnel_data: client {tunnel['client_id']} not active")
             elif data.get('type') == 'tunnel_response':
                 # Forward tunnel response from client to wsssh
                 request_id = data['request_id']
@@ -439,21 +484,24 @@ async def handle_websocket(websocket, path=None):
                 request_id = data['request_id']
                 if request_id in active_tunnels:
                     tunnel = active_tunnels[request_id]
-                    # Forward close to client
+                    # Forward close to client if still active
+                    if tunnel['client_id'] in clients and clients[tunnel['client_id']]['status'] == 'active':
                         await tunnel['client_ws'].send(json.dumps({
                             "type": "tunnel_close",
                             "request_id": request_id
                         }))
                     # Clean up tunnel
                     del active_tunnels[request_id]
+                    if debug: print(f"[DEBUG] [WebSocket] Tunnel {request_id} closed")
     except websockets.exceptions.ConnectionClosed:
-        # Remove from registry and clean up tunnels
+        # Mark client as disconnected instead of removing immediately
         disconnected_client = None
-        for cid, ws in clients.items():
-            if ws == websocket:
+        for cid, client_info in clients.items():
+            if client_info['websocket'] == websocket:
                 disconnected_client = cid
-                del clients[cid]
-                print(f"Client {cid} disconnected")
+                clients[cid]['status'] = 'disconnected'
+                clients[cid]['last_seen'] = time.time()
+                print(f"[DEBUG] [WebSocket] Client {cid} disconnected (marked for timeout)")
                 break
 
         # Clean up active tunnels for this client
@@ -464,6 +512,13 @@ async def handle_websocket(websocket, path=None):
                     tunnels_to_remove.append(request_id)
             for request_id in tunnels_to_remove:
                 del active_tunnels[request_id]
+                if debug: print(f"[DEBUG] [WebSocket] Tunnel {request_id} cleaned up due to client disconnect")
+
+async def cleanup_task():
+    """Periodic task to clean up expired clients"""
+    while True:
+        await asyncio.sleep(10)  # Run every 10 seconds
+        cleanup_expired_clients()
 
 async def main():
     config = configparser.ConfigParser()
@@ -539,6 +594,9 @@ async def main():
 
     print(f"WebSocket SSH Daemon running on {args.host}:{args.port}")
 
+    # Start cleanup task
+    cleanup_coro = asyncio.create_task(cleanup_task())
+
     # Start web interface if specified
     if args.web_host and args.web_port:
         # Handle HTTPS setup
@@ -565,7 +623,11 @@ async def main():
         flask_thread.start()
         print(f"Web interface available at {protocol}://{args.web_host}:{args.web_port}")
 
+    # Wait for server to close
     await ws_server.wait_closed()
 
+    # Cancel cleanup task
+    cleanup_coro.cancel()
+
 if __name__ == '__main__':
     asyncio.run(main())
\ No newline at end of file