Update app.py

app.py

@@ -1,73 +1,56 @@
 import asyncio
 import websockets
+import random
+import string
 import sqlite3
-import secrets
 
-# Connect to SQLite database (create it if not exists)
-conn = sqlite3.connect('client_pairs.db')
+# Database setup (using SQLite for simplicity)
+conn = sqlite3.connect('pairing_database.db')
 cursor = conn.cursor()
-cursor.execute('''
-    CREATE TABLE IF NOT EXISTS pairs (
-        key TEXT PRIMARY KEY,
-        client_address TEXT
-    )
-''')
+cursor.execute('CREATE TABLE IF NOT EXISTS pairs (key TEXT, client_ws TEXT)')
 conn.commit()
 
 async def handle_client(websocket, path):
     try:
-        # Inform the client about successful connection
-        await websocket.send("Connected to the server")
-
-        while True:
-            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")
-
+        # Receive the client's request (create a pair or connect to a key)
+        request_type = await websocket.recv()
+
+        if request_type == "create_pair":
+            # Generate a unique key
+            unique_key = generate_key()
+
+            # Save the key and associate it with the current client in the database
+            cursor.execute('INSERT INTO pairs VALUES (?, ?)', (unique_key, str(websocket)))
+            conn.commit()
+
+            # Inform the client about the created key
+            await websocket.send(f"Pair key created: {unique_key}")
+
+        elif request_type == "connect_to_key":
+            # Receive the key from the client
+            key_to_connect = await websocket.recv()
+
+            # Retrieve all clients associated with the key from the database
+            cursor.execute('SELECT client_ws FROM pairs WHERE key=?', (key_to_connect,))
+            results = cursor.fetchall()
+
+            if results:
+                # Pair the current client with all clients associated with the key
+                for result in results:
+                    other_client_ws = result[0]
+                    if other_client_ws != str(websocket):
+                        await websocket.send(f"Connected to key {key_to_connect}. You are paired with {other_client_ws}")
+                        await other_client_ws.send(f"Connected to key {key_to_connect}. You are paired with {str(websocket)}")
             else:
-                # Broadcast messages to all connected clients
-                for client in connected_clients:
-                    if client != websocket:
-                        await client.send(f"Client says: {command}")
+                # Inform the client that the key does not exist
+                await websocket.send(f"Key {key_to_connect} does not exist.")
 
     except websockets.exceptions.ConnectionClosed:
-        pass
-    finally:
-        # Remove the client from the set upon disconnection
-        connected_clients.remove(websocket)
+        pass  # Connection closed, handle appropriately
 
-# Set to store connected clients
-connected_clients = set()
+def generate_key():
+    # Generate a random key (you might want to implement a more robust key generation logic)
+    return ''.join(random.choices(string.ascii_uppercase + string.digits, k=8))
 
 start_server = websockets.serve(handle_client, "localhost", 8765)