collab_v2.0.0 / app.py
Ramesh-vani's picture
Update app.py
ecc1b65 verified
raw
history blame
2.6 kB
import asyncio
import websockets
import sqlite3
import secrets
# SQLite database setup
conn = sqlite3.connect('pairing.db')
cursor = conn.cursor()
# 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 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}")
# 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)
except websockets.exceptions.ConnectionClosed:
pass # Exit the loop if connection is closed
start_server = websockets.serve(handle_client, "localhost", 8765)
async def main():
print("WebSocket server is running. Waiting for connections...")
await start_server
await asyncio.Event().wait()
if __name__ == "__main__":
connected_clients = set()
asyncio.run(main())