client.py

import asyncio
import websockets
import socket
import os
import pyautogui
from PIL import Image
import io
import base64
import json

SERVER_URL = "ws://localhost:8000/ws/"  # Change to your server IP if needed

def get_computer_name():
    return socket.gethostname()

async def send_screenshot(websocket):
    screenshot = pyautogui.screenshot()
    buf = io.BytesIO()
    screenshot.save(buf, format='JPEG', quality=50)
    img_bytes = buf.getvalue()
    img_b64 = base64.b64encode(img_bytes).decode('utf-8')
    await websocket.send(json.dumps({"type": "screen", "data": img_b64}))

async def handle_input(event):
    if event['input_type'] == 'mouse_move':
        pyautogui.moveTo(event['x'], event['y'])
    elif event['input_type'] == 'mouse_click':
        pyautogui.click(event['x'], event['y'])
    elif event['input_type'] == 'key':
        pyautogui.press(event['key'])

async def listen():
    name = get_computer_name()
    uri = SERVER_URL + name
    while True:
        try:
            async with websockets.connect(uri, max_size=2**23) as websocket:
                print(f"Connected as {name}")
                while True:
                    # Send screenshot
                    await send_screenshot(websocket)
                    # Listen for input events
                    try:
                        msg = await asyncio.wait_for(websocket.recv(), timeout=0.1)
                        try:
                            event = json.loads(msg)
                            if event.get('type') == 'input':
                                await handle_input(event['data'])
                        except Exception as e:
                            print(f"Error handling input: {e}")
                    except asyncio.TimeoutError:
                        pass
                    await asyncio.sleep(0.5)
        except Exception as e:
            print(f"Disconnected, retrying in 5s: {e}")
            await asyncio.sleep(5)

if __name__ == "__main__":
    asyncio.run(listen())