Pasted

#!/usr/bin/env python3
"""
VTTGrabber v4.04 - Stable USB Input Edition
- Adds scroll wheel zoom and middle-click pan for Foundry VTT
- Atomic message IDs for reliable CDP communication
- Retains all v4.03 functionality: strict USB filtering, modifier keys, device selection
"""
 
import asyncio
import evdev
from evdev import ecodes, InputDevice
import json
import os
import signal
import sys
import urllib.request
import websockets
from pathlib import Path
 
CONFIG_FILE = Path.home() / '.vttgrabber.json'
CDP_URL = "ws://localhost:9222/devtools/browser"
SCREEN_WIDTH = 1920
SCREEN_HEIGHT = 1080
 
 
class DeviceScanner:
    """Strict USB device discovery with ASUS hardware blacklist"""
    BUILT_IN_PATTERNS = ['AT Translated', 'Lid Switch', 'Power Button', 'Sleep Button']
    
    def is_forbidden(self, device):
        name = device.name.lower()
        info = device.info
        
        if info.vendor == 0x0b05 and info.product == 0x19b6:
            return True, 'ASUS built-in hardware (Hardcoded block)'
        if 'asus' in name and 'keyboard' in name:
            return True, 'ASUS built-in hardware (Name block)'
        if any(p.lower() in name for p in self.BUILT_IN_PATTERNS):
            return True, 'System built-in'
        return False, ''
 
    def is_usb(self, device):
        return device.phys and 'usb' in device.phys.lower()
 
    def get_all_devices(self):
        devices = {'remotes': [], 'keyboards': [], 'mice': [], 'forbidden': []}
        
        for path in evdev.list_devices():
            try:
                dev = InputDevice(path)
                caps = dev.capabilities()
                
                device_info = {
                    'name': dev.name,
                    'path': path,
                    'usb_path': self._get_usb_path(path)
                }
                
                forbidden, reason = self.is_forbidden(dev)
                if forbidden:
                    device_info['reason'] = reason
                    devices['forbidden'].append(device_info)
                    continue
                    
                if not self.is_usb(dev):
                    continue
 
                # Hardware heuristics: mice have REL_X, keyboards have >20 keys
                keys = caps.get(ecodes.EV_KEY, [])
                rel_axes = caps.get(ecodes.EV_REL, [])
                abs_axes = caps.get(ecodes.EV_ABS, [])
                
                is_mouse = ecodes.REL_X in rel_axes or ecodes.ABS_X in abs_axes
                is_keyboard = len(keys) > 20
 
                # Combo devices (wireless dongles) register as both
                if is_mouse and is_keyboard:
                    devices['remotes'].append(device_info)
                    devices['keyboards'].append(device_info)
                    devices['mice'].append(device_info)
                elif is_mouse:
                    devices['mice'].append(device_info)
                elif is_keyboard:
                    devices['keyboards'].append(device_info)
                    
            except (PermissionError, OSError):
                continue
        
        return devices
    
    def _get_usb_path(self, path):
        try:
            real = os.path.realpath(path)
            parts = real.split('/')
            for i, part in enumerate(parts):
                if 'usb' in part.lower():
                    return '/'.join(parts[:i+2])
        except: 
            pass
        return path
    
    def find_combo_devices(self, devices):
        usb_groups = {}
        for dev in devices:
            usb = dev.get('usb_path', dev['path'])
            usb_groups.setdefault(usb, []).append(dev)
        
        combos = []
        for usb_path, devs in usb_groups.items():
            if len(devs) > 1:
                combos.append({
                    'name': devs[0]['name'],
                    'path': devs[0]['path'],
                    'usb_path': usb_path,
                    'interfaces': devs
                })
        return combos
    
    def prompt_device_selection(self, devices_list, device_type, forbidden_list=None, saved_path=None):
        if not devices_list:
            print(f"\nNo USB {device_type} devices found.")
            return None
        
        print(f"\n--- Select USB {device_type.upper()} (Enter to skip) ---")
        
        for i, dev in enumerate(devices_list, 1):
            marker = " [saved]" if dev['path'] == saved_path else ""
            print(f"  [{i}] {dev['name']}{marker}")
            print(f"      {dev['path']}")
        
        if forbidden_list:
            print(f"\n  Forbidden/Blocked Hardware:")
            for dev in forbidden_list:
                print(f"    - {dev['name']} ({dev['reason']})")
        
        choice = input(f"\nEnter 1-{len(devices_list)} or press Enter to skip: ").strip()
        
        if choice.isdigit():
            idx = int(choice) - 1
            if 0 <= idx < len(devices_list):
                return devices_list[idx]['path']
        
        if saved_path:
            print(f"Using saved: {saved_path}")
            return saved_path
        
        return None
 
 
class VTTGrabber:
    def __init__(self):
        self.devices = {}
        self.ws = None
        self.target_id = None
        self.running = False
        self.msg_counter = 0  # Atomic message ID counter
        
        self.mouse_x = SCREEN_WIDTH // 2
        self.mouse_y = SCREEN_HEIGHT // 2
        self.scroll_accumulator = 0  # For smooth scroll handling
        
        self.shift_pressed = False
        self.ctrl_pressed = False
        self.alt_pressed = False
        self.meta_pressed = False
        
    def load_saved(self):
        if CONFIG_FILE.exists():
            try: 
                return json.load(open(CONFIG_FILE))
            except: 
                pass
        return {}
    
    def save_config(self, remote, keyboard, mouse):
        try: 
            json.dump({'remote': remote, 'keyboard': keyboard, 'mouse': mouse}, open(CONFIG_FILE, 'w'))
        except: 
            pass
 
    async def select_and_grab_devices(self):
        scanner = DeviceScanner()
        devices = scanner.get_all_devices()
        saved = self.load_saved()
        
        all_remotes = devices['remotes'].copy()
        all_mice = devices['mice'].copy()
        combos = scanner.find_combo_devices(devices['keyboards'])
        for combo in combos:
            if combo not in all_remotes:
                all_remotes.append(combo)
        
        remote_path = scanner.prompt_device_selection(all_remotes, 'remote combo/trackpad', saved_path=saved.get('remote'))
        keyboard_path = scanner.prompt_device_selection(devices['keyboards'], 'keyboard', forbidden_list=devices['forbidden'], saved_path=saved.get('keyboard'))
        mouse_path = scanner.prompt_device_selection(all_mice, 'mouse/trackpad', saved_path=saved.get('mouse'))
        
        self.save_config(remote_path, keyboard_path, mouse_path)
        
        paths = {'remote': remote_path, 'keyboard': keyboard_path, 'mouse': mouse_path}
        for name, path in paths.items():
            if path:
                try:
                    dev = InputDevice(path)
                    dev.grab()
                    self.devices[name] = dev
                    print(f"Opened and grabbed {name}: {dev.name}")
                except Exception as e:
                    print(f"Failed to open {name}: {e}. Run with sudo.")
        return len(self.devices) > 0
 
    async def find_chromium_target(self):
        try:
            req = urllib.request.Request("http://localhost:9222/json", headers={"Host": "localhost:9222"})
            with urllib.request.urlopen(req, timeout=5) as response:
                targets = json.loads(response.read().decode())
            for target in targets:
                if target["type"] == "page":
                    self.target_id = target["id"]
                    print(f"Found target: {target.get('title', 'Unknown')}")
                    return target["webSocketDebuggerUrl"]
        except Exception as e:
            print(f"Error finding Chromium target: {e}")
            return None
 
    async def connect_cdp(self):
        ws_url = await self.find_chromium_target()
        if not ws_url:
            raise ConnectionError("No Chromium target found. Is Chrome running with --remote-debugging-port=9222?")
        self.ws = await websockets.connect(ws_url, ping_interval=None)
        await self.send_cdp("Input.enable", {})
        await self.send_cdp("Runtime.enable", {})
        asyncio.create_task(self._drain_ws())
 
    async def _drain_ws(self):
        """Background task to prevent WebSocket buffer overflow"""
        try:
            while self.running and self.ws:
                try: 
                    await asyncio.wait_for(self.ws.recv(), timeout=0.1)
                except asyncio.TimeoutError: 
                    continue
                except Exception: 
                    break
        except Exception: 
            pass
 
    async def send_cdp(self, method, params):
        """Send CDP command with atomic message ID"""
        if not self.ws: 
            return
        self.msg_counter += 1
        msg = {"id": self.msg_counter, "method": method, "params": params}
        try: 
            await self.ws.send(json.dumps(msg))
        except Exception: 
            pass
 
    async def inject_cursor(self):
        """Inject synthetic cursor for visibility"""
        cursor_script = """
        (function() {
            const existing = document.getElementById('vtt-cursor');
            if (existing) existing.remove();
            const c = document.createElement('div');
            c.id = 'vtt-cursor';
            c.style.cssText = `position: fixed; width: 24px; height: 24px; background: radial-gradient(circle, rgba(255,0,85,0.9) 0%, rgba(255,0,85,0.6) 40%, rgba(255,0,85,0) 70%); border: 3px solid #ff0055; border-radius: 50%; pointer-events: none; z-index: 2147483647; left: 0; top: 0; box-shadow: 0 0 10px #ff0055, 0 0 20px #ff0055; transition: transform 0.05s ease-out;`;
            document.body.appendChild(c);
            window.vttCursor = c;
            window.updateVTTCursor = function(x, y) { if (window.vttCursor) window.vttCursor.style.transform = `translate(${x}px, ${y}px)`; };
        })();
        """
        await self.send_cdp("Runtime.evaluate", {"expression": cursor_script, "awaitPromise": False})
 
    async def _get_modifiers(self):
        """Calculate modifier bitmask for CDP"""
        modifiers = 0
        if self.alt_pressed: 
            modifiers |= 1
        if self.ctrl_pressed: 
            modifiers |= 2
        if self.meta_pressed: 
            modifiers |= 4
        if self.shift_pressed: 
            modifiers |= 8
        return modifiers
 
    # =========================================================================
    # MOUSE HANDLING - Added scroll wheel and middle-click
    # =========================================================================
    async def route_mouse_device(self, device):
        """Handle mouse movement, buttons, and scroll"""
        async for event in device.async_read_loop():
            if not self.running: 
                break
            
            # Mouse movement
            if event.type == ecodes.EV_REL:
                if event.code == ecodes.REL_X: 
                    self.mouse_x = max(0, min(SCREEN_WIDTH - 1, self.mouse_x + event.value))
                elif event.code == ecodes.REL_Y: 
                    self.mouse_y = max(0, min(SCREEN_HEIGHT - 1, self.mouse_y + event.value))
                
                await self.send_cdp("Input.dispatchMouseEvent", {
                    "type": "mouseMoved", 
                    "x": self.mouse_x, 
                    "y": self.mouse_y, 
                    "modifiers": await self._get_modifiers()
                })
                
                # Update synthetic cursor position
                await self.send_cdp("Runtime.evaluate", {
                    "expression": f"if(window.updateVTTCursor) window.updateVTTCursor({self.mouse_x},{self.mouse_y});",
                    "awaitPromise": False
                })
            
            # Scroll wheel (NEW in v4.04)
            elif event.type == ecodes.EV_REL and event.code == ecodes.REL_WHEEL:
                self.scroll_accumulator += event.value
                if abs(self.scroll_accumulator) >= 1:
                    await self.send_cdp("Input.dispatchMouseEvent", {
                        "type": "mouseWheel",
                        "x": self.mouse_x,
                        "y": self.mouse_y,
                        "deltaX": 0,
                        "deltaY": int(self.scroll_accumulator * -50),  # Negative for natural scroll
                        "modifiers": await self._get_modifiers()
                    })
                    self.scroll_accumulator = 0
            
            # Mouse buttons (NEW: added middle click 274)
            elif event.type == ecodes.EV_KEY:
                if event.code in (272, 273, 274):  # 272=Left, 273=Right, 274=Middle
                    button_map = {272: "left", 273: "right", 274: "middle"}
                    button = button_map[event.code]
                    action = "mousePressed" if event.value == 1 else "mouseReleased"
                    
                    await self.send_cdp("Input.dispatchMouseEvent", {
                        "type": action,
                        "x": self.mouse_x,
                        "y": self.mouse_y,
                        "button": button,
                        "clickCount": 1,
                        "modifiers": await self._get_modifiers()
                    })
 
    # =========================================================================
    # KEYBOARD HANDLING - Unchanged from v4.03 (proven working)
    # =========================================================================
    def _parse_dynamic_key(self, keycode):
        """Extended dynamic parser for VTT functional keys"""
        key_name = ecodes.KEY.get(keycode)
        if not key_name: 
            return None, None, None
        if isinstance(key_name, list): 
            key_name = key_name[0]
        
        base = key_name.replace('KEY_', '')
        
        # Standard Alphanumeric
        if len(base) == 1 and base.isalpha():
            return (base if self.shift_pressed else base.lower(), f"Key{base}", base if self.shift_pressed else base.lower())
        if len(base) == 1 and base.isdigit():
            symbols = {'1':'!', '2':'@', '3':'#', '4':'$', '5':'%', '6':'^', '7':'&', '8':'*', '9':'(', '0':')'}
            return (symbols[base] if self.shift_pressed else base, f"Digit{base}", symbols[base] if self.shift_pressed else base)
            
        # Expanded Functional Mapping
        specials = {
            'SPACE': (' ', 'Space', ' '), 
            'ENTER': ('\r', 'Enter', None), 
            'KPENTER': ('\r', 'NumpadEnter', None),
            'BACKSPACE': ('Backspace', 'Backspace', None),
            'DELETE': ('Delete', 'Delete', None),
            'PAGEUP': ('PageUp', 'PageUp', None),
            'PAGEDOWN': ('PageDown', 'PageDown', None),
            'UP': ('ArrowUp', 'ArrowUp', None), 
            'DOWN': ('ArrowDown', 'ArrowDown', None), 
            'LEFT': ('ArrowLeft', 'ArrowLeft', None), 
            'RIGHT': ('ArrowRight', 'ArrowRight', None),
            'MINUS': ('_' if self.shift_pressed else '-', 'Minus', '_' if self.shift_pressed else '-'), 
            'EQUAL': ('+' if self.shift_pressed else '=', 'Equal', '+' if self.shift_pressed else '='),
            'SLASH': ('?' if self.shift_pressed else '/', 'Slash', '?' if self.shift_pressed else '/'), 
            'DOT': ('>' if self.shift_pressed else '.', 'Period', '>' if self.shift_pressed else '.'),
            'COMMA': ('<' if self.shift_pressed else ',', 'Comma', '<' if self.shift_pressed else ','),
            'ESC': ('Escape', 'Escape', None)
        }
        return specials.get(base, (None, None, None))
 
    async def _send_modifier_keys(self, keycode, value):
        """Send modifier key state changes"""
        mod_map = {
            ecodes.KEY_LEFTSHIFT: ("Shift", "ShiftLeft"), 
            ecodes.KEY_RIGHTSHIFT: ("Shift", "ShiftRight"),
            ecodes.KEY_LEFTCTRL: ("Control", "ControlLeft"), 
            ecodes.KEY_RIGHTCTRL: ("Control", "ControlRight"),
            ecodes.KEY_LEFTALT: ("Alt", "AltLeft"), 
            ecodes.KEY_RIGHTALT: ("Alt", "AltRight"),
        }
        if keycode in mod_map:
            key, code = mod_map[keycode]
            await self.send_cdp("Input.dispatchKeyEvent", {
                "type": "keyDown" if value > 0 else "keyUp", 
                "key": key, 
                "code": code, 
                "modifiers": await self._get_modifiers()
            })
 
    async def route_keyboard_device(self, device):
        """Handle keyboard input with strict event type checking"""
        async for event in device.async_read_loop():
            if not self.running or event.type != ecodes.EV_KEY: 
                continue
                
            keycode, value = event.code, event.value
            
            # Update modifier state
            if keycode in (ecodes.KEY_LEFTSHIFT, ecodes.KEY_RIGHTSHIFT): 
                self.shift_pressed = value > 0
                await self._send_modifier_keys(keycode, value)
                continue
            elif keycode in (ecodes.KEY_LEFTCTRL, ecodes.KEY_RIGHTCTRL): 
                self.ctrl_pressed = value > 0
                await self._send_modifier_keys(keycode, value)
                continue
            elif keycode in (ecodes.KEY_LEFTALT, ecodes.KEY_RIGHTALT): 
                self.alt_pressed = value > 0
                await self._send_modifier_keys(keycode, value)
                continue
            
            # Special key handling
            if keycode == ecodes.KEY_TAB and value == 1:
                await self.send_cdp("Runtime.evaluate", {
                    "expression": "const evt = new KeyboardEvent('keydown', {key: 'Tab', code: 'Tab', keyCode: 9, which: 9, bubbles: true}); (document.activeElement || document.body).dispatchEvent(evt);"
                })
                continue
            if keycode == ecodes.KEY_ESC and value == 1:
                await self.send_cdp("Runtime.evaluate", {
                    "expression": "document.querySelectorAll('.window-header .close, #chat-controls .close').forEach(b => b.click()); const c = document.querySelector('canvas#board'); if(c)c.focus(); document.dispatchEvent(new KeyboardEvent('keydown', {key: 'Escape', code: 'Escape', keyCode: 27, which: 27, bubbles: true}));"
                })
                continue
            
            # Regular key handling
            display_key, code, text_char = self._parse_dynamic_key(keycode)
            if not display_key: 
                continue
            
            if value in (1, 2):  # Key down or repeat
                await self.send_cdp("Input.dispatchKeyEvent", {
                    "type": "keyDown", 
                    "key": display_key, 
                    "code": code, 
                    "modifiers": await self._get_modifiers()
                })
                if value == 1 and text_char:  # Only send char on initial press
                    await self.send_cdp("Input.dispatchKeyEvent", {
                        "type": "char", 
                        "text": text_char, 
                        "modifiers": await self._get_modifiers()
                    })
            elif value == 0:  # Key up
                await self.send_cdp("Input.dispatchKeyEvent", {
                    "type": "keyUp", 
                    "key": display_key, 
                    "code": code, 
                    "modifiers": await self._get_modifiers()
                })
 
    # =========================================================================
    # MAIN
    # =========================================================================
    async def run(self):
        self.running = True
        
        if not await self.select_and_grab_devices():
            print("No USB devices selected. Exiting.")
            return
            
        print("\nConnecting to Chromium...")
        await self.connect_cdp()
        await self.inject_cursor()
        
        # Create tasks for all grabbed devices
        tasks = []
        for name, device in self.devices.items():
            if 'mouse' in name or 'remote' in name:
                tasks.append(asyncio.create_task(self.route_mouse_device(device)))
            if 'keyboard' in name or 'remote' in name:
                tasks.append(asyncio.create_task(self.route_keyboard_device(device)))
        
        await asyncio.gather(*tasks)
 
    def cleanup(self):
        """Release all grabbed devices"""
        self.running = False
        for name, device in self.devices.items():
            try: 
                device.ungrab()
                print(f"Released {name}")
            except: 
                pass
 
 
if __name__ == "__main__":
    grabber = VTTGrabber()
    
    # Signal handling for clean shutdown
    def signal_handler(signum, frame):
        print(f"\nReceived signal {signum}, shutting down...")
        grabber.cleanup()
        sys.exit(0)
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    try: 
        asyncio.run(grabber.run())
    except (KeyboardInterrupt, Exception) as e: 
        print(f"Error: {e}")
        grabber.cleanup()