Switch to asyncio with dbus-next

antidrift/auth.py

@@ -0,0 +1,33 @@
+from dbus_next.auth import Authenticator, _AuthResponse
+
+
+class AuthExternal(Authenticator):
+    """An authenticator class for the external auth protocol for use with the
+    :class:`MessageBus <dbus_next.message_bus.BaseMessageBus>`.
+    :sealso: https://dbus.freedesktop.org/doc/dbus-specification.html#auth-protocol
+    """
+    def __init__(self, user_uid):
+        self.user_uid = user_uid
+        self.negotiate_unix_fd = False
+        self.negotiating_fds = False
+
+    def _authentication_start(self, negotiate_unix_fd=False) -> str:
+        self.negotiate_unix_fd = negotiate_unix_fd
+        hex_uid = str(self.user_uid).encode().hex()
+        return f'AUTH EXTERNAL {hex_uid}'
+
+    def _receive_line(self, line: str):
+        response, args = _AuthResponse.parse(line)
+
+        if response is _AuthResponse.OK:
+            if self.negotiate_unix_fd:
+                self.negotiating_fds = True
+                return "NEGOTIATE_UNIX_FD"
+            else:
+                return "BEGIN"
+
+        if response is _AuthResponse.AGREE_UNIX_FD:
+            return "BEGIN"
+
+        raise AuthError(f'authentication failed: {response.value}: {args}')
+

antidrift/client.py

@@ -1,19 +1,43 @@
-import time
+import asyncio
-import antidrift.daemon
+from dbus_next.aio import MessageBus
+from dbus_next import BusType
+from dbus_next.errors import DBusError
 from antidrift.config import Config
 from argparse import Namespace
 from rich import print
+from antidrift.daemon import IFACE, OPATH, BUS_NAME
 
 
-def antidrift_daemon_is_running() -> bool:
+async def get_dbus_interface():
-    """Check if AntiDrift is running via the DBUS"""
+    try:
-    interface = antidrift.daemon.get_dbus_interface()
+        bus = await MessageBus(bus_type=BusType.SESSION).connect()
+        introspection = await bus.introspect(BUS_NAME, OPATH)
+        proxy_obj = bus.get_proxy_object(BUS_NAME, OPATH, introspection)
+        return proxy_obj.get_interface(IFACE)
+    except DBusError:
+        return None
+
+
+async def run(args: Namespace, config: Config):
+    interface = await get_dbus_interface()
+    reply = "🟡 ad daemon active but no command"
     if interface is None:
-        return False
+        reply = "🔴 ad inactive"
-    reply = interface.status()
+    elif args.start:
-    if reply:
+        reply = await interface.call_start(args.start)
-        return True
+    elif args.stop:
-    return False
+        reply = await interface.call_stop()
+    elif args.pause:
+        reply = await interface.call_pause()
+    elif args.unpause:
+        reply = await interface.call_unpause()
+    elif args.schedule:
+        reply = await interface.call_schedule(args.schedule)
+    elif args.tailf:
+        tailf(config)
+    elif args.status:
+        reply = await interface.call_status()
+    print(reply)
 
 
 def tailf(config):
@@ -25,21 +49,3 @@ def tailf(config):
                 time.sleep(0.1)
             else:
                 print(line.strip())
-
-
-def run(args: Namespace, config: Config):
-    interface = antidrift.daemon.get_dbus_interface()
-    reply = "🟡 ad daemon active but no command"
-    if interface is None:
-        reply = "🔴 ad inactive"
-    elif args.start:
-        reply = interface.start(args.start)
-    elif args.stop:
-        reply = interface.stop()
-    elif args.schedule:
-        reply = interface.schedule(args.schedule)
-    elif args.tailf:
-        tailf(config)
-    elif args.status:
-        reply = interface.status()
-    print(reply)

antidrift/config.py

@@ -22,7 +22,7 @@ class Config(BaseModel):
     daemon_log_file: Path = Path()
     client_log_file: Path = Path()
     config_file: Path = Path()
-    polling_cycle_ms: int = 500
+    polling_cycle_ms: int = 2000
     enforce_delay_ms: int = 5000
 
     class Config:
@@ -46,6 +46,7 @@ class State(BaseModel):
     active_blackblocks: List[Block] = []
     active_whiteblocks: List[Block] = []
     inactive_blackblocks: List[Block] = []
+    pause: bool = False
 
     class Config:
         extra = "forbid"

antidrift/daemon.py

@@ -1,59 +1,82 @@
 from datetime import datetime
 import csv
-import dbus
-import dbus.service
 import logging
 import os
 import pwd
 import re
 import sys
 import time
+import asyncio
 
 import antidrift.xwindow as xwindow
 from antidrift.xwindow import XWindow
 from antidrift.config import Config, State, Block
-from gi.repository import GLib, Gio
+from antidrift.auth import AuthExternal
 from typing import List, Optional
 
+from dbus_next.aio import MessageBus
+from dbus_next.service import ServiceInterface, method
+from dbus_next import Variant, BusType
+
 BUS_NAME = "com.antidrift"
 IFACE = "com.antidrift"
 OPATH = "/com/antidrift"
 
 
-def reload_callback(m, f, o, event):
+class AntiDriftDaemon(ServiceInterface):
-    filename = f.get_basename()
-    m = f"[dark_orange3]Restart after change in '{filename}'.[/dark_orange3]"
-    logging.warning(m)
-    os.execv(sys.executable, ["python3"] + sys.argv)
-
-
-def get_dbus_interface() -> Optional[dbus.Interface]:
-    try:
-        bus = dbus.SessionBus()
-        bus_object = bus.get_object(BUS_NAME, OPATH)
-        interface = dbus.Interface(bus_object, IFACE)
-        return interface
-    except dbus.exceptions.DBusException:
-        return None
-
-
-class AntiDriftDaemon(dbus.service.Object):
     def __init__(self, config: Config):
-        user_name = os.environ.get("SUDO_USER", pwd.getpwuid(os.getuid()).pw_name)
+        super().__init__(IFACE)
-        user_uid = pwd.getpwnam(user_name).pw_uid
-        euid = os.geteuid()
-        os.seteuid(user_uid)
-        bus = dbus.bus.BusConnection(f"unix:path=/run/user/{user_uid}/bus")
-        os.seteuid(euid)
-
-        bus.request_name(BUS_NAME)
-        bus_name = dbus.service.BusName(BUS_NAME, bus=bus)
-        dbus.service.Object.__init__(self, bus_name, OPATH)
         self.config = config
         self.reset_block_state()
         self.enforce_count = 0
         self.enforce_value = int(config.enforce_delay_ms / config.polling_cycle_ms)
 
+    async def init_bus(self):
+        """
+        We are switching the effective UID to the target user's UID in order to
+        connect to the D-Bus session bus with the correct permissions and
+        authentication. Once the D-Bus connection is established, we restore
+        the original effective UID to maintain the appropriate privilege
+        levels.
+        """
+        user_name = os.environ.get("SUDO_USER", pwd.getpwuid(os.getuid()).pw_name)
+        user_uid = pwd.getpwnam(user_name).pw_uid
+        euid = os.geteuid()
+        os.seteuid(user_uid)
+        auth = AuthExternal(user_uid)
+        bus_address = f"unix:path=/run/user/{user_uid}/bus"
+        bus = MessageBus(bus_address=bus_address, bus_type=BusType.SESSION, auth=auth)
+        await bus.connect()
+        bus.export(OPATH, self)
+        await bus.request_name(BUS_NAME)
+        os.seteuid(euid)
+        return bus
+
+    async def run(self, debug: bool = False):
+        bus = await self.init_bus()
+
+        async def _enforce():
+            while True:
+                if self.state.pause is True:
+                    await self.enforce_pause()
+                else:
+                    await self.enforce()
+                await asyncio.sleep(self.config.polling_cycle_ms / 1000)
+
+        async def _log():
+            while True:
+                if self.state.pause is False:
+                    self.log_window()
+                await asyncio.sleep(60)  # Sleep for 60 seconds
+
+        # Start _enforce and _log as background tasks
+        asyncio.create_task(_enforce())
+        asyncio.create_task(_log())
+
+        xwindow.notify(f"AntiDrift running.")
+        stop = asyncio.Event()
+        await stop.wait()
+
     def reset_block_state(self):
         self.state = State(
             active_blackblocks=self.config.blackblocks,
@@ -61,8 +84,8 @@ class AntiDriftDaemon(dbus.service.Object):
             inactive_blackblocks=[],
         )
 
-    @dbus.service.method(dbus_interface=IFACE, in_signature="as", out_signature="s")
+    @method()
-    def start(self, whiteblocks: List[str]) -> str:
+    def start(self, whiteblocks: 'as') -> 's':
         self.reset_block_state()
         all_whiteblocks = {wb.name: wb for wb in self.config.whiteblocks}
         success_wbs, fail_blocks = [], []
@@ -83,8 +106,8 @@ class AntiDriftDaemon(dbus.service.Object):
             logging.warning(m)
         return r
 
-    @dbus.service.method(dbus_interface=IFACE, in_signature="s", out_signature="s")
+    @method()
-    def schedule(self, blackblock_name: str) -> str:
+    def schedule(self, blackblock_name: 's') -> 's':
         """Schedule blackblock based if it has a non-zero timeout value."""
         all_blackblocks = {bb.name: bb for bb in self.config.blackblocks}
         if blackblock_name not in all_blackblocks:
@@ -101,27 +124,46 @@ class AntiDriftDaemon(dbus.service.Object):
         def allow():
             self.allow_blackblock(blackblock)
 
-        delay_ms = blackblock.delay * 1000 * 60
+            delay_sec = blackblock.delay * 60
-        GLib.timeout_add(delay_ms, allow)
+
-        m = f"Scheduled [sky_blue3]{blackblock_name}[/sky_blue3] in {blackblock.delay} minutes."
+        delay_sec = blackblock.delay * 60
+        loop = asyncio.get_event_loop()
+        loop.call_later(delay_sec, allow)
+        m = f"Scheduled [sky_blue3]{blackblock.name}[/sky_blue3] in {blackblock.delay} minutes."
         logging.info(m)
         return m
 
-    @dbus.service.method(dbus_interface=IFACE, in_signature="", out_signature="s")
+    @method()
-    def stop(self) -> str:
+    def stop(self) -> 's':
         self.reset_block_state()
         m = "Blacklist only mode."
         logging.info(m)
         return m
 
-    @dbus.service.method(dbus_interface=IFACE, in_signature="", out_signature="s")
+    @method()
-    def status(self) -> str:
+    def pause(self) -> 's':
+        self.state.pause = True
+        m = "Antidrift paused."
+        logging.info(m)
+        return m
+
+    @method()
+    def unpause(self) -> 's':
+        self.state.pause = False
+        m = "Antidrift unpaused."
+        logging.info(m)
+        return m
+
+    @method()
+    def status(self) -> 's':
         white_active = bool(self.state.active_whiteblocks)
         black_active = bool(self.state.active_blackblocks)
         m = "🟢 ad "
         inactive_bbs = " ".join(
             map(lambda b: "-" + b.name, self.state.inactive_blackblocks)
         )
+        if self.state.pause is True:
+            return "🟡 ad paused"
         match (white_active, black_active):
             case (True, _):
                 m += "wb: "
@@ -150,15 +192,6 @@ class AntiDriftDaemon(dbus.service.Object):
         s = " ".join(map(lambda b: b.name, self.state.active_whiteblocks))
         return f"intention is {s} work" if s else "no intention"
 
-    def log_window(self):
-        self.config.window_log_file.parent.mkdir(parents=True, exist_ok=True)
-        window = XWindow()
-        ts = int(time.time())
-        intention = self.get_intention()
-        log_line = f"{ts}, {window.name}, {window.cls}, {intention}\n"
-        with self.config.window_log_file.open('a') as f:
-            f.write(log_line)
-
     def log_window(self):
         window = XWindow()
         utc_timestamp = datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
@@ -174,54 +207,31 @@ class AntiDriftDaemon(dbus.service.Object):
             writer = csv.writer(f)
             writer.writerow(log_line)
 
-    def run(self, debug: bool = False):
+    async def enforce_pause(self):
-        def _enforce():
+        xwindow.notify("Goint to minimize window because of pause...")
-            self.enforce()
+        for _ in range(8):
-            GLib.timeout_add(self.config.polling_cycle_ms, _enforce)
+            await asyncio.sleep(1)
+            if not XWindow().is_active() or self.state.pause is False:
+                return
 
-        def _log():
+        window = XWindow()
-            self.log_window()
+        if window.is_active():
-            ONE_MINUTE_IN_MS = 60 * 1000
-            GLib.timeout_add(ONE_MINUTE_IN_MS, _log)
-
-        # autorestart on file change for development
-        monitors = []
-        files = [
-            "antidrift.py",
-            "antidrift/daemon.py",
-            "antidrift/client.py",
-            "antidrift/config.py",
-        ]
-        if debug:
-            logging.warning("[red]Running in debug mode.[/red]")
-            for filename in files:
-                gio_file = Gio.File.new_for_path(filename)
-                monitor = gio_file.monitor_file(Gio.FileMonitorFlags.NONE, None)
-                monitor.connect("changed", reload_callback)
-                monitors.append(monitor)
-
-        self.config.window_log_file.parent.mkdir(parents=True, exist_ok=True)
-        _enforce()
-        _log()
-        mainloop = GLib.MainLoop()
-        mainloop.run()
-        xwindow.notify(f"AntiDrift running.")
-
-    def enforce(self):
-        if self.enforce_count >= self.enforce_value:
-            window = XWindow()
-            xwindow.notify(f"Minimize {window.name[:30]}.")
             window.minimize()
-            self.enforce_count = 0
+
-        elif self.enforce_count > 0 and window_is_blocked(self.state, True):
+    async def enforce(self):
-            self.enforce_count += 1
+        if not window_is_blocked(self.state):
-        elif self.enforce_count == 0 and window_is_blocked(self.state):
+            return
-            self.enforce_count += 1
+
-            delay = int(self.config.enforce_delay_ms / 1000)
+        delay = int(self.config.enforce_delay_ms / 1000)
-            xwindow.notify(f"AntiDrift will minimize in {delay}s.")
+        for i in range(delay, 0, -1):
-        elif self.enforce_count > 0:
+            await asyncio.sleep(1)
-            xwindow.notify("We are gucci again.")
+            xwindow.notify(f"AntiDrift will minimize in {i}s.")
-            self.enforce_count = 0
+            if not window_is_blocked(self.state, silent=True):
+                xwindow.notify("We are gucci again.")
+                return
+        window = XWindow()
+        xwindow.notify(f"Minimize {window.name[:30]}.")
+        window.minimize()
 
 
 def window_is_blocked(state: State, silent: bool = False) -> bool:

antidrift/xwindow.py

@@ -37,6 +37,9 @@ class XWindow:
     def kill(self):
         self._run(["windowkill", self.window])
 
+    def is_active(self):
+        return True if self.name else False
+
 
 def notify(message: str) -> None:
     """Notify user via the Xorg notify-send command."""

main.py

@@ -7,6 +7,7 @@ import subprocess
 import argparse
 import psutil
 import rich
+import asyncio
 from rich.logging import RichHandler
 from pathlib import Path
 
@@ -25,10 +26,10 @@ def get_args():
     parser.add_argument("--daemon", action="store_true", help="run daemon")
     parser.add_argument("--status", action="store_true", help="get status from daemon")
     parser.add_argument("--tailf", action="store_true", help="tail -f log file")
-    parser.add_argument(
+    parser.add_argument("--start", metavar="whiteblock", nargs="+", help="start whiteblocks")
-        "--start", metavar="whiteblock", nargs="+", help="start whiteblocks"
-    )
     parser.add_argument("--stop", action="store_true", help="stop session")
+    parser.add_argument("--pause", action="store_true", help="pause antidrift")
+    parser.add_argument("--unpause", action="store_true", help="unpause antidrift")
     parser.add_argument("--schedule", metavar="blackblock", help="schedule blackblock")
     args = parser.parse_args()
     return args
@@ -97,7 +98,8 @@ def main_daemon():
         if newpid == 0:
             config = Config.load(os.path.expanduser("~/.config/antidrift.yaml"))
             init_logging(config.daemon_log_file)
-            AntiDriftDaemon(config).run()
+            daemon = AntiDriftDaemon(config)
+            asyncio.run(daemon.run())
     else:
         if sys.argv[0] == "antidrift":
             kill_existing_antidrift()
@@ -106,7 +108,8 @@ def main_daemon():
         else:
             config = Config.load(os.path.expanduser("~/.config/antidrift.yaml"))
             init_logging(config.daemon_log_file, dev_mode=True)
-            AntiDriftDaemon(config).run(debug=True)
+            daemon = AntiDriftDaemon(config)
+            asyncio.run(daemon.run(debug=True))
 
 
 def main() -> None:
@@ -117,7 +120,7 @@ def main() -> None:
         main_daemon()
     else:
         config = Config.load(os.path.expanduser("~/.config/antidrift.yaml"))
-        antidrift.client.run(args, config)
+        asyncio.run(antidrift.client.run(args, config))
 
 
 if __name__ == "__main__":