Spaces:
Running
Running
import asyncio | |
import websockets | |
import random | |
import string | |
import sqlite3 | |
# 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, client_ws 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() | |
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: | |
# 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 # 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)) | |
start_server = websockets.serve(handle_client, "0.0.0.0", 7860) | |
async def main(): | |
print("WebSocket server is running. Waiting for connections...") | |
await start_server | |
await asyncio.Event().wait() | |
if __name__ == "__main__": | |
asyncio.run(main()) | |