Update app.py

app.py

@@ -3,13 +3,11 @@ import websockets
 import sqlite3
 import secrets
 
-# SQLite database setup
-conn = sqlite3.connect('pairing.db')
+# Connect to SQLite database (create it if not exists)
+conn = sqlite3.connect('client_pairs.db')
 cursor = conn.cursor()
-
-# Create Pairing table if not exists
 cursor.execute('''
-    CREATE TABLE IF NOT EXISTS Pairing (
+    CREATE TABLE IF NOT EXISTS pairs (
         key TEXT PRIMARY KEY,
         client_address TEXT
     )
@@ -18,54 +16,58 @@ conn.commit()
 
 async def handle_client(websocket, path):
     try:
-        # Receive a command from the client
-        command = await websocket.recv()
-
-        if command == "create_pair":
-            # Create a unique key
-            key = secrets.token_hex(16)
-
-            # Save the key and client address in the database
-            cursor.execute('INSERT INTO Pairing (key, client_address) VALUES (?, ?)', (key, str(websocket.remote_address)))
-            conn.commit()
-
-            # Inform the client about the generated key
-            await websocket.send(f"Your pair key is: {key}")
-
-        elif command.startswith("connect_to_key"):
-            # Extract the key from the command
-            key_to_connect = command.split(" ")[1]
-
-            # Check if the key exists in the database
-            cursor.execute('SELECT client_address FROM Pairing WHERE key = ?', (key_to_connect,))
-            result = cursor.fetchone()
-
-            if result:
-                other_client_address = result[0]
-
-                # Inform both clients about the successful pairing
-                await websocket.send(f"Connected to {other_client_address}")
-                await websocket.send(f"You are paired with {other_client_address}")
+        # Inform the client about successful connection
+        await websocket.send("Connected to the server")
 
-                # Create a bidirectional bridge
-                asyncio.ensure_future(bridge_clients(websocket, other_client_address))
-
-    except websockets.exceptions.ConnectionClosed:
-        pass  # Connection closed, handle appropriately
-
-async def bridge_clients(source, destination_address):
-    destination = next(c for c in connected_clients if str(c.remote_address) == destination_address)
-
-    try:
         while True:
-            # Receive data from the source client
-            data_from_source = await source.recv()
-
-            # Forward the data to the destination client
-            await destination.send(data_from_source)
+            command = await websocket.recv()
+
+            if command == "create_pair_key":
+                # Generate a unique key
+                pair_key = secrets.token_urlsafe(8)
+
+                # Save the key and client address in the database
+                cursor.execute("INSERT INTO pairs (key, client_address) VALUES (?, ?)", (pair_key, str(websocket.remote_address)))
+                conn.commit()
+
+                # Send the generated key to the client
+                await websocket.send(f"Pair key created: {pair_key}")
+
+            elif command.startswith("connect_to_key "):
+                key_to_connect = command.split(" ")[1]
+
+                # Check if the key exists in the database
+                cursor.execute("SELECT client_address FROM pairs WHERE key=?", (key_to_connect,))
+                result = cursor.fetchone()
+
+                if result:
+                    other_client_address = result[0]
+
+                    # Inform both clients about the successful pairing
+                    await websocket.send(f"Paired with {other_client_address}")
+                    await websocket.send(f"Connected to {websocket.remote_address}")
+                    other_client = next((client for client in connected_clients if str(client.remote_address) == other_client_address), None)
+                    
+                    if other_client:
+                        await other_client.send(f"Paired with {websocket.remote_address}")
+                        await other_client.send(f"Connected to {other_client.remote_address}")
+                else:
+                    await websocket.send("Invalid pair key")
+
+            else:
+                # Broadcast messages to all connected clients
+                for client in connected_clients:
+                    if client != websocket:
+                        await client.send(f"Client says: {command}")
 
     except websockets.exceptions.ConnectionClosed:
-        pass  # Exit the loop if connection is closed
+        pass
+    finally:
+        # Remove the client from the set upon disconnection
+        connected_clients.remove(websocket)
+
+# Set to store connected clients
+connected_clients = set()
 
 start_server = websockets.serve(handle_client, "localhost", 8765)
 
@@ -75,5 +77,4 @@ async def main():
     await asyncio.Event().wait()
 
 if __name__ == "__main__":
-    connected_clients = set()
     asyncio.run(main())