rescriptoon/rescriptoon/__init__.py

import copy
import logging
import os
import select
import time
import typing

import Xlib.display
from Xlib import XK, X, Xatom

from rescriptoon._actions import (
    RewriteKeyEventAction,
    SelectGagAction,
    ToggleOverlayAction,
)

EXTENDED_CONTROLS_DEFAULT_TOGGLE_KEYSYM_NAME = "grave"
TOONTOWN_WINDOW_NAME = "Toontown Rewritten"

EXTENDED_CONTROLS_DEFAULT_KEYSYM_MAPPINGS = {
    XK.XK_w: RewriteKeyEventAction(keysym=XK.XK_Up, target_engine_index=0),
    XK.XK_a: RewriteKeyEventAction(keysym=XK.XK_Left, target_engine_index=0),
    XK.XK_s: RewriteKeyEventAction(keysym=XK.XK_Down, target_engine_index=0),
    XK.XK_d: RewriteKeyEventAction(keysym=XK.XK_Right, target_engine_index=0),
    XK.XK_Control_L: RewriteKeyEventAction(
        keysym=XK.XK_Control_L, target_engine_index=0
    ),
    XK.XK_v: RewriteKeyEventAction(keysym=XK.XK_Delete, target_engine_index=0),
    XK.XK_i: RewriteKeyEventAction(keysym=XK.XK_Up, target_engine_index=1),
    XK.XK_j: RewriteKeyEventAction(keysym=XK.XK_Left, target_engine_index=1),
    XK.XK_k: RewriteKeyEventAction(keysym=XK.XK_Down, target_engine_index=1),
    XK.XK_l: RewriteKeyEventAction(keysym=XK.XK_Right, target_engine_index=1),
    XK.XK_slash: RewriteKeyEventAction(keysym=XK.XK_Control_L, target_engine_index=1),
    XK.XK_n: RewriteKeyEventAction(keysym=XK.XK_Delete, target_engine_index=1),
    XK.XK_space: RewriteKeyEventAction(keysym=XK.XK_Control_L),
    XK.XK_e: SelectGagAction(
        target_engine_index=0, column_index=4, factor_y=-0.047
    ),  # elephant trunk
    XK.XK_o: SelectGagAction(
        target_engine_index=1, column_index=4, factor_y=-0.047
    ),  # elephant trunk
    XK.XK_f: SelectGagAction(
        target_engine_index=0, column_index=5, factor_y=-0.047
    ),  # foghorn
    XK.XK_semicolon: SelectGagAction(
        target_engine_index=1, column_index=5, factor_y=-0.047
    ),  # foghorn
}


def _x_walk_children_windows(
    parent_window: "Xlib.display.Window",
) -> typing.Iterator["Xlib.display.Window"]:
    yield parent_window
    for child_window in parent_window.query_tree().children:
        for subchild_window in _x_walk_children_windows(child_window):
            yield subchild_window


def _x_wait_for_event(xdisplay, timeout_seconds):
    """ Wait up to `timeout_seconds` seconds for a xevent.
        Return True, if a xevent is available.
        Return False, if the timeout was reached. """
    rlist = select.select(
        [xdisplay.display.socket],  # rlist
        [],  # wlist
        [],  # xlist
        timeout_seconds,  # timeout [seconds]
    )[0]
    return len(rlist) > 0


class Overlay:
    def __init__(
        self, toggle_keysym_name=EXTENDED_CONTROLS_DEFAULT_TOGGLE_KEYSYM_NAME,
    ):
        self._xdisplay = Xlib.display.Display()
        self._toggle_keysym = XK.string_to_keysym(toggle_keysym_name)
        if self._toggle_keysym == X.NoSymbol:
            raise ValueError(
                "rescriptoon controls toggle:"
                + " unknown keysym name '{}'".format(toggle_keysym_name)
            )
        self._keysym_mappings = copy.deepcopy(
            EXTENDED_CONTROLS_DEFAULT_KEYSYM_MAPPINGS,
        )
        if self._toggle_keysym in self._keysym_mappings:
            logging.warning("ignoring mapping for toggle key %s", toggle_keysym_name)
        self._keysym_mappings[self._toggle_keysym] = ToggleOverlayAction()
        self._active_key_registry = {}
        self._enabled = False
        self._engine_windows = None

    @property
    def xdisplay(self):
        return self._xdisplay

    @property
    def engine_windows(self) -> typing.List["Xlib.display.Window"]:
        return self._engine_windows

    @property
    def _engine_windows_open(self) -> bool:
        for window in self._engine_windows:
            try:
                window.get_wm_state()
            except Xlib.error.BadWindow:
                logging.info("engine window %x is no longer available", window.id)
                return False
        return True

    def run(self) -> None:
        self._engine_windows = self._find_engine_windows()
        logging.debug("engine window ids %r", [hex(w.id) for w in self._engine_windows])
        if not self._engine_windows:
            raise Exception("no toontown window found")
        self._grab_key(self.xdisplay.keysym_to_keycode(self._toggle_keysym),)
        self.enable()
        while self._engine_windows_open:
            while self.xdisplay.pending_events():
                self._handle_xevent(self.xdisplay.next_event())
            self._check_active_key_registry()
            # keep timeout low for _check_active_key_registry()
            # to be called frequently
            _x_wait_for_event(self.xdisplay, timeout_seconds=0.05)
        self._disable()

    def _handle_xevent(self, xevent: Xlib.protocol.rq.Event) -> None:
        if isinstance(
            xevent, (Xlib.protocol.event.KeyPress, Xlib.protocol.event.KeyRelease)
        ):
            self._handle_xkeyevent(xevent)

    def _handle_xkeyevent(
        self,
        xkeyevent: typing.Union[
            Xlib.protocol.event.KeyPress, Xlib.protocol.event.KeyRelease
        ],
    ) -> None:
        self._update_active_key_registry(xkeyevent)
        keysym_in = self.xdisplay.keycode_to_keysym(xkeyevent.detail, index=0,)
        action = self._keysym_mappings[keysym_in]
        action.execute(self, xkeyevent)

    @property
    def _toggle_keysym_name(self) -> str:
        return XK.keysym_to_string(self._toggle_keysym)

    def enable(self):
        for keysym in self._keysym_mappings.keys():
            if keysym != self._toggle_keysym:
                self._grab_key(self.xdisplay.keysym_to_keycode(keysym),)
        self._enabled = True
        logging.info(
            "rescriptoon is now enabled. press %s to disable.",
            self._toggle_keysym_name,
        )

    def _disable(self):
        for keysym in self._keysym_mappings.keys():
            if keysym != self._toggle_keysym:
                self._ungrab_key(self.xdisplay.keysym_to_keycode(keysym),)
        self._enabled = False

    def disable(self):
        self._disable()
        logging.info(
            "rescriptoon is now disabled. press %s to enable.",
            self._toggle_keysym_name,
        )

    @property
    def enabled(self):
        return self._enabled

    def toggle(self):
        if self.enabled:
            self.disable()
        else:
            self.enable()

    def _grab_key(self, keycode):
        for window in self._engine_windows:
            window.grab_key(
                keycode,
                X.AnyModifier,
                # owner_events
                # https://stackoverflow.com/questions/32122360/x11-will-xgrabpointer-prevent-other-apps-from-any-mouse-event
                # False,
                True,
                X.GrabModeAsync,
                X.GrabModeAsync,
            )

    def _ungrab_key(self, keycode):
        for window in self._engine_windows:
            window.ungrab_key(keycode, X.AnyModifier)

    def _find_engine_windows(self) -> typing.List["Xlib.display.Window"]:
        return list(
            filter(
                lambda w: w.get_wm_name() == TOONTOWN_WINDOW_NAME,
                _x_walk_children_windows(self.xdisplay.screen().root),
            )
        )

    def _update_active_key_registry(self, xkeyevent):
        # see self._check_active_key_registry
        keycode = xkeyevent.detail
        if isinstance(xkeyevent, Xlib.protocol.event.KeyPress):
            self._active_key_registry[keycode] = xkeyevent
        elif keycode in self._active_key_registry:
            del self._active_key_registry[keycode]

    def _check_active_key_registry(self):
        """
        WORKAROUND
        For an unknown reason some key release events don't get queued
        when multiple keys are being released simultaneously.
        So we keep a hashmap of supposedly currently pressed keys
        and periodically compare it with xdispaly.query_keymap().

        ref: https://stackoverflow.com/q/18160792/5894777
        """
        # https://tronche.com/gui/x/xlib/input/XQueryKeymap.html
        keymap = self.xdisplay.query_keymap()
        missed_releases = []
        for keycode, press_event in self._active_key_registry.items():
            byte_index = keycode >> 3
            bit_index = keycode & ((1 << 3) - 1)
            if not keymap[byte_index] & (1 << bit_index):
                print("DEBUG missed release event of key {}".format(keycode))
                missed_releases.append(
                    Xlib.protocol.event.KeyRelease(
                        window=press_event.window,
                        detail=press_event.detail,
                        state=press_event.state,
                        root_x=press_event.root_x,
                        root_y=press_event.root_y,
                        event_x=press_event.event_x,
                        event_y=press_event.event_y,
                        child=press_event.child,
                        root=press_event.root,
                        time=X.CurrentTime,
                        same_screen=press_event.same_screen,
                    )
                )
        for release_event in missed_releases:
            self._handle_xkeyevent(release_event)