InputRemapper/inputremapper/focus/backends/sway.py

# -*- coding: utf-8 -*-
# input-remapper - GUI for device specific keyboard mappings
# Copyright (C) 2025 sezanzeb <b8x45ygc9@mozmail.com>
#
# This file is part of input-remapper.
#
# input-remapper is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# input-remapper is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with input-remapper.  If not, see <https://www.gnu.org/licenses/>.

"""Focus backend for Sway (and other i3-IPC compatible compositors)."""

from __future__ import annotations

import json
import os
import socket
import struct
from typing import Optional

import gi

from inputremapper.focus.focus_backend import FocusBackend, FocusEvent, OnFocus
from inputremapper.logging.logger import logger

gi.require_version("GLib", "2.0")
from gi.repository import GLib  # noqa: E402

# https://i3wm.org/docs/ipc.html
_MAGIC = b"i3-ipc"
_HEADER = struct.Struct("=6sII")  # magic, length, type (native byte order)

_MSG_SUBSCRIBE = 2
_MSG_GET_TREE = 4

_EVENT_MASK = 0x80000000


def _pack(message_type: int, payload: bytes = b"") -> bytes:
    return _HEADER.pack(_MAGIC, len(payload), message_type) + payload


def _recv_exactly(sock: socket.socket, length: int) -> bytes:
    chunks = []
    received = 0
    while received < length:
        chunk = sock.recv(length - received)
        if not chunk:
            raise ConnectionError("Socket closed while reading")
        chunks.append(chunk)
        received += len(chunk)
    return b"".join(chunks)


def _recv_message(sock: socket.socket) -> tuple[int, bytes]:
    header = _recv_exactly(sock, _HEADER.size)
    _, length, message_type = _HEADER.unpack(header)
    payload = _recv_exactly(sock, length) if length else b""
    return message_type, payload


def _container_to_event(container: dict) -> Optional[FocusEvent]:
    if not container:
        return None

    app_id = container.get("app_id")
    if not app_id:
        # XWayland windows expose their class via window_properties instead.
        window_properties = container.get("window_properties") or {}
        app_id = window_properties.get("class")

    title = container.get("name") or ""
    return FocusEvent(app_id=app_id or "", title=title, backend=SwayFocusBackend.name)


def _find_focused(node: dict) -> Optional[dict]:
    if node.get("focused"):
        return node
    for child in node.get("nodes", []) + node.get("floating_nodes", []):
        found = _find_focused(child)
        if found is not None:
            return found
    return None


class SwayFocusBackend(FocusBackend):
    """Subscribes to window events over the i3-IPC socket ($SWAYSOCK)."""

    name = "sway"

    def __init__(self) -> None:
        self._socket: Optional[socket.socket] = None
        self._watch_id: Optional[int] = None
        self._buffer = b""
        self._on_focus: Optional[OnFocus] = None

    @staticmethod
    def _socket_path() -> Optional[str]:
        return os.environ.get("SWAYSOCK") or os.environ.get("I3SOCK")

    @staticmethod
    def is_available() -> bool:
        path = SwayFocusBackend._socket_path()
        return path is not None and os.path.exists(path)

    def start(self, on_focus: OnFocus) -> None:
        path = self._socket_path()
        if not path:
            raise RuntimeError("SWAYSOCK is not set")

        self._on_focus = on_focus
        self._socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self._socket.connect(path)
        self._socket.sendall(_pack(_MSG_SUBSCRIBE, json.dumps(["window"]).encode()))
        self._socket.setblocking(False)

        self._watch_id = GLib.io_add_watch(
            self._socket.fileno(),
            GLib.IO_IN | GLib.IO_HUP | GLib.IO_ERR,
            self._on_readable,
        )
        logger.debug("Sway focus backend listening on %s", path)

    def _on_readable(self, _fd: int, condition: int) -> bool:
        if condition & (GLib.IO_HUP | GLib.IO_ERR):
            logger.error("Sway IPC socket closed")
            return False

        assert self._socket is not None
        try:
            data = self._socket.recv(65536)
        except BlockingIOError:
            return True
        except OSError as error:
            logger.error("Sway IPC read failed: %s", error)
            return False

        if not data:
            return False

        self._buffer += data
        self._consume_buffer()
        return True

    def _consume_buffer(self) -> None:
        while len(self._buffer) >= _HEADER.size:
            _, length, message_type = _HEADER.unpack(self._buffer[: _HEADER.size])
            total = _HEADER.size + length
            if len(self._buffer) < total:
                break

            payload = self._buffer[_HEADER.size : total]
            self._buffer = self._buffer[total:]

            if message_type & _EVENT_MASK:
                self._handle_event(payload)

    def _handle_event(self, payload: bytes) -> None:
        try:
            data = json.loads(payload.decode())
        except (json.JSONDecodeError, UnicodeDecodeError):
            return

        change = data.get("change")
        if change not in ("focus", "title", "new"):
            return

        event = _container_to_event(data.get("container") or {})
        if event is not None and self._on_focus is not None:
            self._on_focus(event)

    def stop(self) -> None:
        if self._watch_id is not None:
            GLib.source_remove(self._watch_id)
            self._watch_id = None
        if self._socket is not None:
            self._socket.close()
            self._socket = None
        self._buffer = b""

    def get_current(self) -> Optional[FocusEvent]:
        path = self._socket_path()
        if not path:
            return None

        try:
            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
                sock.connect(path)
                sock.sendall(_pack(_MSG_GET_TREE))
                while True:
                    message_type, payload = _recv_message(sock)
                    if message_type == _MSG_GET_TREE:
                        break
        except OSError as error:
            logger.error("Failed to query Sway tree: %s", error)
            return None

        try:
            tree = json.loads(payload.decode())
        except (json.JSONDecodeError, UnicodeDecodeError):
            return None

        focused = _find_focused(tree)
        if focused is None:
            return None
        return _container_to_event(focused)