Update app.py

app.py

@@ -1,56 +1,71 @@
 import asyncio
 import websockets
-import random
-import string
 import sqlite3
+import secrets
 
-# Database setup (using SQLite for simplicity)
-conn = sqlite3.connect('pairing_database.db')
+# SQLite database setup
+conn = sqlite3.connect('pairing.db')
 cursor = conn.cursor()
-cursor.execute('CREATE TABLE IF NOT EXISTS pairs (key TEXT, client_ws TEXT)')
+
+# Create Pairing table if not exists
+cursor.execute('''
+    CREATE TABLE IF NOT EXISTS Pairing (
+        key TEXT PRIMARY KEY,
+        client_address TEXT
+    )
+''')
 conn.commit()
 
 async def handle_client(websocket, path):
     try:
-        # Receive the client's request (create a pair or connect to a key)
-        request_type = await websocket.recv()
+        # Receive a command from the client
+        command = await websocket.recv()
 
-        if request_type == "create_pair":
-            # Generate a unique key
-            unique_key = generate_key()
+        if command == "create_pair":
+            # Create a unique key
+            key = secrets.token_hex(16)
 
-            # Save the key and associate it with the current client in the database
-            cursor.execute('INSERT INTO pairs VALUES (?, ?)', (unique_key, str(websocket)))
+            # 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 created key
-            await websocket.send(f"Pair key created: {unique_key}")
+            # Inform the client about the generated key
+            await websocket.send(f"Your pair key is: {key}")
 
-        elif request_type == "connect_to_key":
-            # Receive the key from the client
-            key_to_connect = await websocket.recv()
+        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_ws FROM pairs WHERE key=?', (key_to_connect,))
+            cursor.execute('SELECT client_address FROM Pairing WHERE key = ?', (key_to_connect,))
             result = cursor.fetchone()
 
             if result:
-                other_client_ws = result[0]
+                other_client_address = result[0]
 
-                # Pair the clients
-                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)}")
+                # 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}")
 
-            else:
-                # Inform the client that the key does not exist
-                await websocket.send(f"Key {key_to_connect} does not exist.")
+                # Create a bidirectional bridge
+                other_client = next(c for c in connected_clients if str(c.remote_address) == other_client_address)
+                asyncio.ensure_future(bridge_clients(websocket, other_client))
+                asyncio.ensure_future(bridge_clients(other_client, websocket))
 
     except websockets.exceptions.ConnectionClosed:
         pass  # Connection closed, handle appropriately
 
-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))
+async def bridge_clients(source, destination):
+    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)
+
+    except websockets.exceptions.ConnectionClosed:
+        pass  # Exit the loop if connection is closed
 
 start_server = websockets.serve(handle_client, "localhost", 8765)
 
@@ -60,4 +75,5 @@ async def main():
     await asyncio.Event().wait()
 
 if __name__ == "__main__":
+    connected_clients = set()
     asyncio.run(main())