app.py

import asyncio
import json
import os
import shutil
import subprocess
import time
import logging
import concurrent.futures
import threading
import numpy as np
import uuid
import psutil
from aiohttp import web
from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack, RTCIceServer, RTCConfiguration
from av import VideoFrame
import mss

HOST = "0.0.0.0"
PORT = 7860
DISPLAY_NUM = ":99"

MAX_WIDTH = 3840
MAX_HEIGHT = 2160

DEFAULT_WIDTH = 1280
DEFAULT_HEIGHT = 720

TURN_USER = "g08abe68c81a07f098bb5f0914549bb32440e5aad0b216c7fba2b61e76fd62c6"
TURN_PASS = "aed1a10dd10eba9401ad9d99e5c66036d8a970eab5ba8e6dc9845ab57c771a7d"

logging.basicConfig(level=logging.WARNING) 
logger = logging.getLogger("WebRTC-Antigravity")

executor = concurrent.futures.ThreadPoolExecutor(max_workers=16)
thread_local_storage = threading.local()

config = {
    "width": DEFAULT_WIDTH,
    "height": DEFAULT_HEIGHT
}

class InputManager:
    def __init__(self):
        self.process = None
        self.lock = threading.Lock()
        self.scroll_accum = 0

    def start_process(self):
        if not os.environ.get("DISPLAY"):
            return

        try:
            self.process = subprocess.Popen(
                ['xdotool', '-'], 
                stdin=subprocess.PIPE, 
                encoding='utf-8', 
                bufsize=0 
            )
        except Exception as e:
            logger.error(f"Failed to start xdotool: {e}")

    def _send_raw(self, command):
        if self.process is None or self.process.poll() is not None:
            self.start_process()
        
        if self.process:
            try:
                self.process.stdin.write(command + "\n")
                self.process.stdin.flush()
            except Exception:
                try: self.process.kill()
                except: pass
                self.process = None

    def send(self, command):
        with self.lock:
            self._send_raw(command)

    def scroll(self, dy):
        with self.lock:
            self.scroll_accum += dy
            THRESHOLD = 40 
            while self.scroll_accum >= THRESHOLD:
                self._send_raw("click 5")
                self.scroll_accum -= THRESHOLD
            while self.scroll_accum <= -THRESHOLD:
                self._send_raw("click 4")
                self.scroll_accum += THRESHOLD

    def mouse_move(self, x, y): self.send(f"mousemove {x} {y}")
    def mouse_down(self, btn): self.send(f"mousedown {btn}")
    def mouse_up(self, btn): self.send(f"mouseup {btn}")
    def click(self, btn, repeat=1): self.send(f"click --repeat {repeat} {btn}")
    def key_down(self, key): self.send(f"keydown {key}")
    def key_up(self, key): self.send(f"keyup {key}")

input_manager = InputManager()

def start_system():
    os.environ["DISPLAY"] = DISPLAY_NUM

    if not shutil.which("Xvfb"): raise FileNotFoundError("Xvfb missing")
    
    subprocess.Popen([
        "Xvfb", DISPLAY_NUM, 
        "-screen", "0", f"{MAX_WIDTH}x{MAX_HEIGHT}x24",
        "-ac", "-noreset"
    ])
    
    time.sleep(3)

    input_manager.start_process()

    set_resolution(DEFAULT_WIDTH, DEFAULT_HEIGHT)

    if shutil.which("matchbox-window-manager"):
        subprocess.Popen("matchbox-window-manager -use_titlebar no", shell=True)

    threading.Thread(target=maintain_antigravity, daemon=True).start()

def maintain_antigravity():
    while True:
        antigravity_running = False
        for proc in psutil.process_iter(['name']):
            try:
                # Check for exact name 'antigravity'
                if 'antigravity' in proc.info['name']:
                    antigravity_running = True
                    break
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                pass
        
        if not antigravity_running:
            logger.warning("Antigravity not found. Launching...")
            # We assume the binary is simply named 'antigravity' and available in PATH
            try:
                subprocess.Popen(["antigravity"])
                time.sleep(5)
            except Exception as e:
                logger.error(f"Failed to launch antigravity: {e}")
        
        time.sleep(2)

def get_xrandr_output_name():
    try:
        out = subprocess.check_output(["xrandr"]).decode()
        for line in out.splitlines():
            if " connected" in line:
                return line.split()[0]
    except: pass
    return "screen"

def get_cvt_modeline(width, height, rate=60):
    H_BLANK = 160
    H_SYNC = 32
    H_FRONT_PORCH = 48
    V_FRONT_PORCH = 3
    V_SYNC = 5
    MIN_V_BLANK = 460
    
    frame_time_us = 1000000.0 / rate
    active_time_us = frame_time_us - MIN_V_BLANK
    if active_time_us <= 0: return None
    
    h_period_us = active_time_us / height
    v_blank_lines = int(MIN_V_BLANK / h_period_us) + 1
    v_total = height + v_blank_lines
    h_total = width + H_BLANK
    pclk = (h_total * v_total * rate) / 1000000.0
    
    h_sync_start = width + H_FRONT_PORCH
    h_sync_end = h_sync_start + H_SYNC
    v_sync_start = height + V_FRONT_PORCH
    v_sync_end = v_sync_start + V_SYNC
    
    return f'"{width}x{height}_60.00" {pclk:.2f} {width} {h_sync_start} {h_sync_end} {h_total} {height} {v_sync_start} {v_sync_end} {v_total} +hsync -vsync'

def set_resolution(w, h):
    try:
        if w % 2 != 0: w += 1
        if h % 2 != 0: h += 1

        output = get_xrandr_output_name()
        mode_name = f"WEB_{w}x{h}_{str(uuid.uuid4())[:4]}"
        modeline_str = get_cvt_modeline(w, h)
        if not modeline_str: return
            
        parts = modeline_str.split()
        mode_params = parts[1:]

        subprocess.run(["xrandr", "--newmode", mode_name] + mode_params, check=True)
        subprocess.run(["xrandr", "--addmode", output, mode_name], check=True)
        subprocess.run(["xrandr", "--output", output, "--mode", mode_name], check=True)
        
        config["width"] = w
        config["height"] = h
    except Exception as e:
        logger.error(f"Resolution setup failed: {e}")

class VirtualScreenTrack(VideoStreamTrack):
    kind = "video"
    def __init__(self):
        super().__init__()
        self.last_frame_time = 0
        self.frame_count = 0

    def _capture(self):
        try:
            if not hasattr(thread_local_storage, "sct"):
                thread_local_storage.sct = mss.mss()
            
            monitor = {"top": 0, "left": 0, "width": config["width"], "height": config["height"]}
            sct_img = thread_local_storage.sct.grab(monitor)
            img = np.array(sct_img)
            return img[..., :3] 
        except: return None

    async def recv(self):
        FPS = 30
        FRAME_TIME = 1.0 / FPS
        
        pts, time_base = await self.next_timestamp()
        
        current_time = time.time()
        wait = FRAME_TIME - (current_time - self.last_frame_time)
        if wait > 0:
            await asyncio.sleep(wait)
        
        self.last_frame_time = time.time()

        frame = await asyncio.get_event_loop().run_in_executor(executor, self._capture)
        
        if frame is None:
            blank = np.zeros((config["height"], config["width"], 3), dtype=np.uint8)
            av_frame = VideoFrame.from_ndarray(blank, format="bgr24")
        else:
            av_frame = VideoFrame.from_ndarray(frame, format="bgr24")

        av_frame.pts = pts
        av_frame.time_base = time_base
        return av_frame

def map_key(key):
    if key == " ": return "space"
    k = key.lower()
    charmap = {
        "control": "ctrl", "shift": "shift", "alt": "alt", "meta": "super", "cmd": "super",
        "enter": "Return", "backspace": "BackSpace", "tab": "Tab", "escape": "Escape",
        "arrowup": "Up", "arrowdown": "Down", "arrowleft": "Left", "arrowright": "Right",
        "home": "Home", "end": "End", "pageup": "Page_Up", "pagedown": "Page_Down",
        "delete": "Delete", "insert": "Insert",
        "f1": "F1", "f2": "F2", "f3": "F3", "f4": "F4", "f5": "F5", "f6": "F6",
        "f7": "F7", "f8": "F8", "f9": "F9", "f10": "F10", "f11": "F11", "f12": "F12",
        "!": "exclam", "@": "at", "#": "numbersign", "$": "dollar", "%": "percent",
        "^": "asciicircum", "&": "ampersand", "*": "asterisk", "(": "parenleft",
        ")": "parenright", "-": "minus", "_": "underscore", "=": "equal", "+": "plus",
        "[": "bracketleft", "{": "braceleft", "]": "bracketright", "}": "braceright",
        ";": "semicolon", ":": "colon", "'": "apostrophe", "\"": "quotedbl",
        ",": "comma", "<": "less", ".": "period", ">": "greater", "/": "slash",
        "?": "question", "\\": "backslash", "|": "bar", "`": "grave", "~": "asciitilde",
        " ": "space"
    }
    return charmap.get(k, k)

def process_input(data):
    try:
        msg = json.loads(data)
        t = msg.get("type")
        
        current_w = config["width"]
        current_h = config["height"]

        if t == "resize":
            target_w = int(msg.get("width"))
            target_h = int(msg.get("height"))
            set_resolution(target_w, target_h)

        elif t == "mousemove":
            input_manager.mouse_move(int(msg["x"] * current_w), int(msg["y"] * current_h))
        elif t == "mousedown":
            input_manager.mouse_down({0:1, 1:2, 2:3}.get(msg.get("button"), 1))
        elif t == "mouseup":
            input_manager.mouse_up({0:1, 1:2, 2:3}.get(msg.get("button"), 1))
        elif t == "wheel":
            input_manager.scroll(msg.get("deltaY", 0))
        elif t == "keydown":
            k = map_key(msg.get("key"))
            if k: input_manager.key_down(k)
        elif t == "keyup":
            k = map_key(msg.get("key"))
            if k: input_manager.key_up(k)

    except Exception:
        pass 

async def offer(request):
    try:
        params = await request.json()
        offer = RTCSessionDescription(sdp=params["sdp"], type=params["type"])
    except: return web.Response(status=400)

    pc = RTCPeerConnection(RTCConfiguration(iceServers=[
        RTCIceServer(urls=["turns:turn.cloudflare.com:443?transport=tcp", "turn:turn.cloudflare.com:3478?transport=udp"], username=TURN_USER, credential=TURN_PASS),
        RTCIceServer(urls=["stun:stun.l.google.com:19302"]) 
    ]))
    pcs.add(pc)

    @pc.on("connectionstatechange")
    async def on_state():
        if pc.connectionState in ["failed", "closed"]:
            await pc.close()
            pcs.discard(pc)

    @pc.on("datachannel")
    def on_dc(channel):
        channel.on("message", lambda m: asyncio.get_event_loop().run_in_executor(executor, process_input, m))

    pc.addTrack(VirtualScreenTrack())
    await pc.setRemoteDescription(offer)
    answer = await pc.createAnswer()
    await pc.setLocalDescription(answer)

    sdp = "\r\n".join([l for l in pc.localDescription.sdp.splitlines() if "a=candidate" not in l or "typ relay" in l]) + "\r\n"
    return web.Response(content_type="application/json", text=json.dumps({"sdp": sdp, "type": pc.localDescription.type}), headers={"Access-Control-Allow-Origin": "*"})

async def index(r): return web.Response(text="helloworld")
async def options(r): return web.Response(headers={"Access-Control-Allow-Origin": "*", "Access-Control-Allow-Methods": "POST, OPTIONS", "Access-Control-Allow-Headers": "Content-Type"})

pcs = set()
async def on_shutdown(app): await asyncio.gather(*[pc.close() for pc in pcs])

if __name__ == "__main__":
    start_system()
    app = web.Application()
    app.on_shutdown.append(on_shutdown)
    app.router.add_get("/", index)
    app.router.add_post("/offer", offer)
    app.router.add_options("/offer", options)
    web.run_app(app, host=HOST, port=PORT)