Clean up spawn as root and CLI interface

antidrift.desktop

@@ -0,0 +1,6 @@
+[Desktop Entry]
+Name=AntiDrift
+Exec=antidrift
+Terminal=false
+Type=Application
+StartupNotify=false

antidrift.py

@@ -1,73 +0,0 @@
-import logging
-import shutil
-import sys
-import os
-import signal
-import subprocess
-from pathlib import Path
-import antidrift.client as client
-from antidrift.config import Config
-from antidrift.daemon import AntiDriftDaemon
-from dbus.mainloop.glib import DBusGMainLoop
-DBusGMainLoop(set_as_default=True)
-signal.signal(signal.SIGINT, signal.SIG_DFL)
-
-
-def init_logging(log_file: Path):
-    class DuplicateFilter(logging.Filter):
-        def filter(self, record) -> bool:
-            current_log = (record.module, record.levelno, record.msg)
-            if current_log != getattr(self, "last_log", None):
-                self.last_log = current_log
-                return True
-            return False
-    format_str = '[bold pale_green3]%(asctime)s[/bold pale_green3] | ' \
-                 '[light_steel_blue]%(levelname)-8s[/light_steel_blue] | ' \
-                 '%(message)s'
-    logging.basicConfig(filename=log_file,
-                        format=format_str,
-                        datefmt='%a %H:%M:%S',
-                        encoding='utf-8',
-                        level=logging.DEBUG)
-    logger = logging.getLogger()
-    logger.addFilter(DuplicateFilter())
-
-
-def check_for_xdotool():
-    """ Check if xdotool is in path and exit if not """
-    result = shutil.which("xdotool")
-    if not result:
-        logging.critical("Please install xdotool")
-        sys.exit(1)
-
-
-def main_daemon(config):
-    init_logging(config.daemon_log_file)
-    add = AntiDriftDaemon(config)
-    add.run()
-
-
-def main() -> None:
-    """ Main routine that dispatches to client or daemon """
-    config = Config.load(os.path.expanduser("~/.config/antidrift.yaml"))
-    check_for_xdotool()
-
-    if client.antidrift_daemon_is_running():
-        init_logging(config.client_log_file)
-        client.client_mode(config)
-    elif len(sys.argv) == 1:
-        if os.geteuid() == 0:
-            newpid = os.fork()
-            if newpid == 0:
-                main_daemon(config)
-        else:
-            cmd = ["sudo", "antidrift"]
-            subprocess.Popen(cmd)
-    elif len(sys.argv) == 2 and sys.argv[1] == '--daemon_user':
-        main_daemon(config)
-    else:
-        print("ad not running")
-
-
-if __name__ == "__main__":
-    main()

antidrift/client.py

@@ -1,25 +1,13 @@
-from antidrift.config import Config
-from typing import Optional
-from rich import print
-import argparse
 import time
-import dbus
+import antidrift.daemon
-import dbus.service
+from antidrift.config import Config
-
+from argparse import Namespace
-
+from rich import print
-def get_dbus_interface() -> Optional[dbus.Interface]:
-    try:
-        bus = dbus.SessionBus()
-        bus_object = bus.get_object("com.antidrift", "/com/antidrift")
-        interface = dbus.Interface(bus_object, "com.antidrift")
-        return interface
-    except dbus.exceptions.DBusException:
-        return None
 
 
 def antidrift_daemon_is_running() -> bool:
-    """ Check if AntiDrift is running via the DBUS """
+    """Check if AntiDrift is running via the DBUS"""
-    interface = get_dbus_interface()
+    interface = antidrift.daemon.get_dbus_interface()
     if interface is None:
         return False
     reply = interface.status()
@@ -29,7 +17,7 @@ def antidrift_daemon_is_running() -> bool:
 
 
 def tailf(config):
-    with open(config.daemon_log_file, 'r') as f:
+    with open(config.daemon_log_file, "r") as f:
         f.seek(0, 2)
         while True:
             line = f.readline()
@@ -39,21 +27,12 @@ def tailf(config):
                 print(line.strip())
 
 
-def client_mode(config: Config):
+def run(args: Namespace, config: Config):
-    parser = argparse.ArgumentParser(description='AntiDrift CLI.')
+    interface = antidrift.daemon.get_dbus_interface()
-    parser.add_argument('--start', metavar='whiteblock', nargs='+',
+    reply = "🟡 ad daemon active but no command"
-                        help='start whiteblocks')
+    if interface is None:
-    parser.add_argument('--stop', action='store_true',
+        reply = "🔴 ad inactive"
-                        help='stop session')
+    elif args.start:
-    parser.add_argument('--schedule', metavar='blackblock',
-                        help='schedule blackblock')
-    parser.add_argument('--status', action='store_true',
-                        help='get status from daemon')
-    parser.add_argument('--tailf', action='store_true',
-                        help='tail -f log file')
-    args = parser.parse_args()
-    interface = get_dbus_interface()
-    if args.start:
         reply = interface.start(args.start)
     elif args.stop:
         reply = interface.stop()
@@ -63,6 +42,4 @@ def client_mode(config: Config):
         tailf(config)
     elif args.status:
         reply = interface.status()
-    else:
-        reply = '[red]no command[/red]'
     print(reply)

antidrift/config.py

@@ -12,7 +12,7 @@ class Block(BaseModel):
     delay: int = 0
 
     class Config:
-        extra = 'forbid'
+        extra = "forbid"
 
 
 class Config(BaseModel):
@@ -25,7 +25,7 @@ class Config(BaseModel):
     enforce_delay_ms: int = 5000
 
     class Config:
-        extra = 'forbid'
+        extra = "forbid"
 
     @classmethod
     def load(cls, config_file: str) -> Config:
@@ -43,4 +43,4 @@ class State(BaseModel):
     inactive_blackblocks: List[Block] = []
 
     class Config:
-        extra = 'forbid'
+        extra = "forbid"

antidrift/daemon.py

@@ -4,12 +4,13 @@ import pwd
 import re
 import sys
 import antidrift.xwindow as xwindow
-from antidrift.config import Config, State, Block
-from gi.repository import GLib, Gio
-from typing import List
 import dbus
 import dbus.service
 
+from antidrift.config import Config, State, Block
+from gi.repository import GLib, Gio
+from typing import List, Optional
+
 BUS_NAME = "com.antidrift"
 IFACE = "com.antidrift"
 OPATH = "/com/antidrift"
@@ -19,12 +20,21 @@ def reload_callback(m, f, o, event):
     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)
+    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)
         user_uid = pwd.getpwnam(user_name).pw_uid
         euid = os.geteuid()
@@ -44,10 +54,10 @@ class AntiDriftDaemon(dbus.service.Object):
         self.state = State(
             active_blackblocks=self.config.blackblocks,
             active_whiteblocks=[],
-                inactive_blackblocks=[])
+            inactive_blackblocks=[],
+        )
 
-    @dbus.service.method(dbus_interface=IFACE,
+    @dbus.service.method(dbus_interface=IFACE, in_signature="as", out_signature="s")
-                         in_signature="as", out_signature="s")
     def start(self, whiteblocks: List[str]) -> str:
         self.reset_block_state()
         all_whiteblocks = {wb.name: wb for wb in self.config.whiteblocks}
@@ -59,7 +69,7 @@ class AntiDriftDaemon(dbus.service.Object):
             else:
                 fail_blocks.append(block_name)
         if success_wbs:
-            wbs = ', '.join(success_wbs)
+            wbs = ", ".join(success_wbs)
             r = f"Start whiteblocks [sky_blue3]{wbs}[/sky_blue3]."
             logging.info(r)
         else:
@@ -69,10 +79,9 @@ class AntiDriftDaemon(dbus.service.Object):
             logging.warning(m)
         return r
 
-    @dbus.service.method(dbus_interface=IFACE,
+    @dbus.service.method(dbus_interface=IFACE, in_signature="s", out_signature="s")
-                         in_signature="s", out_signature="s")
     def schedule(self, blackblock_name: str) -> str:
-        """ Schedule blackblock based if it has a non-zero timeout value. """
+        """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:
             m = f"No blackblock [red3]{blackblock_name}[/red3]."
@@ -87,41 +96,42 @@ class AntiDriftDaemon(dbus.service.Object):
 
         def allow():
             self.allow_blackblock(blackblock)
+
         delay_ms = blackblock.delay * 1000 * 60
         GLib.timeout_add(delay_ms, 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,
+    @dbus.service.method(dbus_interface=IFACE, in_signature="", out_signature="s")
-                         in_signature="", out_signature="s")
     def stop(self) -> str:
         self.reset_block_state()
-        m = 'Blacklist only mode.'
+        m = "Blacklist only mode."
         logging.info(m)
         return m
 
-    @dbus.service.method(dbus_interface=IFACE,
+    @dbus.service.method(dbus_interface=IFACE, in_signature="", out_signature="s")
-                         in_signature="", out_signature="s")
     def status(self) -> str:
         white_active = bool(self.state.active_whiteblocks)
         black_active = bool(self.state.active_blackblocks)
-        m = 'ad '
+        m = "🟢 ad "
-        inactive_bbs = ' '.join(map(lambda b: "-" + b.name, self.state.inactive_blackblocks))
+        inactive_bbs = " ".join(
+            map(lambda b: "-" + b.name, self.state.inactive_blackblocks)
+        )
         match (white_active, black_active):
             case (True, _):
-                m += 'wb: '
+                m += "wb: "
-                m += ' '.join(map(lambda b: b.name, self.state.active_whiteblocks))
+                m += " ".join(map(lambda b: b.name, self.state.active_whiteblocks))
                 if inactive_bbs:
-                    m += ' '
+                    m += " "
                     m += inactive_bbs
             case (False, True):
-                m += 'bb'
+                m += "bb"
                 if inactive_bbs:
-                    m += ': '
+                    m += ": "
                     m += inactive_bbs
             case _:
-                m = 'inactive'
+                m = "inactive"
         return m
 
     def allow_blackblock(self, blackblock: Block):
@@ -132,15 +142,20 @@ class AntiDriftDaemon(dbus.service.Object):
         m = f"Blackblock [sky_blue3]{blackblock.name}[/sky_blue3] is now allowed."
         logging.info(m)
 
-    def run(self):
+    def run(self, debug: bool = False):
         def _enforce():
             self.enforce()
             GLib.timeout_add(self.config.polling_cycle_ms, _enforce)
 
         # autorestart on file change for development
         monitors = []
-        files = ["antidrift.py", "antidrift/daemon.py", "antidrift/client.py",
+        files = [
-                 "antidrift/config.py"]
+            "antidrift.py",
+            "antidrift/daemon.py",
+            "antidrift/client.py",
+            "antidrift/config.py",
+        ]
+        if debug:
             for filename in files:
                 gio_file = Gio.File.new_for_path(filename)
                 monitor = gio_file.monitor_file(Gio.FileMonitorFlags.NONE, None)
@@ -178,7 +193,7 @@ def window_is_blocked(state: State, silent: bool = False) -> bool:
         return False
 
     def keyword_matches_window(keyword: str, window: xwindow.XWindow):
-        if keyword.startswith('/') and keyword.endswith('/'):
+        if keyword.startswith("/") and keyword.endswith("/"):
             try:
                 r = re.compile(keyword[1:-1], re.IGNORECASE)
                 if r.findall(window.name):
@@ -186,7 +201,7 @@ def window_is_blocked(state: State, silent: bool = False) -> bool:
                 else:
                     return False
             except re.error:
-                m = f'Invalid regex [red3]{keyword}[/red3].'
+                m = f"Invalid regex [red3]{keyword}[/red3]."
                 logging.warning(m)
                 return False
         else:
@@ -204,8 +219,10 @@ def window_is_blocked(state: State, silent: bool = False) -> bool:
             elif keyword_matches_window(k, window):
                 if not silent:
                     xwindow.notify(f"{window.name[:30]} blocked by {b.name}.")
-                    logging.warning(f"[red]{window.name[:50]}[/red] "
+                    logging.warning(
-                                    f"blocked by [red]{b.name}[/red].")
+                        f"[red]{window.name[:50]}[/red] "
+                        f"blocked by [red]{b.name}[/red]."
+                    )
                 return True
     if not whiteblocks:
         if not silent:
@@ -215,9 +232,17 @@ def window_is_blocked(state: State, silent: bool = False) -> bool:
         for k in w.keywords:
             if keyword_matches_window(k, window):
                 if not silent:
-                    logging.debug(f"[pale_green3]{window.name[:30]}[/pale_green3] "
+                    logging.debug(
-                                  f"allowed by [sky_blue3]{w.name}[/sky_blue3].")
+                        f"[pale_green3]{window.name[:30]}[/pale_green3] "
+                        f"allowed by [sky_blue3]{w.name}[/sky_blue3]."
+                    )
                 return False
     if not silent:
         xwindow.notify(f"'{window.name[:30]}…' not on any whiteblock.")
     return True
+
+
+def run(config: Config):
+    add = AntiDriftDaemon(config)
+    xwindow.notify(f"AntiDrift.run()")
+    add.run()

antidrift/obsolete.py

@@ -1,4 +1,3 @@
-
 def is_window_blocked(window_name: str, blocked: List[re.Pattern]) -> bool:
     for b in blocked:
         if b.findall(window_name):
@@ -23,11 +22,8 @@ def kill_sequence(blocked: List[re.Pattern]) -> None:
 
 
 def notify(message: str) -> None:
-    """ Notify user via the Xorg notify-send command. """
+    """Notify user via the Xorg notify-send command."""
-    env = {
+    env = {**os.environ, "DBUS_SESSION_BUS_ADDRESS": "unix:path=/run/user/1000/bus"}
-        **os.environ,
-        "DBUS_SESSION_BUS_ADDRESS": "unix:path=/run/user/1000/bus"
-    }
     user = env.get("SUDO_USER", None)
     if user is None:
         cmd = ["notify-send", message]
@@ -141,7 +137,7 @@ def write_window_names(config: Config, window_names: Set[str]) -> None:
 
 
 def main() -> None:
-    """ Run main_root as root except while debugging. """
+    """Run main_root as root except while debugging."""
     config_path = "~/.config/aw-focus/config.json"
     config = Config.load_config(config_path)
 
@@ -167,5 +163,3 @@ def main_root(config: Config) -> None:
         kill_window_if_blocked(blocked)
         if config.enforce_aw_commit:
             enforce_aw_commit()
-
-

antidrift/xwindow.py

@@ -36,7 +36,7 @@ class XWindow:
 
 
 def notify(message: str) -> None:
-    """ Notify user via the Xorg notify-send command. """
+    """Notify user via the Xorg notify-send command."""
     logging.debug(f"{message} - [grey]notify[/grey]")
     env = dict(os.environ)
     user = env.get("SUDO_USER", None)

main.py

@@ -0,0 +1,107 @@
+import logging
+import shutil
+import sys
+import os
+import signal
+import subprocess
+import argparse
+import rich
+from rich.logging import RichHandler
+from pathlib import Path
+
+import antidrift.client
+import antidrift.daemon
+from antidrift.config import Config
+
+from dbus.mainloop.glib import DBusGMainLoop
+
+DBusGMainLoop(set_as_default=True)
+signal.signal(signal.SIGINT, signal.SIG_DFL)
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="AntiDrift CLI.")
+    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(
+        "--start", metavar="whiteblock", nargs="+", help="start whiteblocks"
+    )
+    parser.add_argument("--stop", action="store_true", help="stop session")
+    parser.add_argument("--schedule", metavar="blackblock", help="schedule blackblock")
+    args = parser.parse_args()
+    return args
+
+
+def init_logging(log_file: Path, dev_mode: bool = False):
+    class DuplicateFilter(logging.Filter):
+        def filter(self, record) -> bool:
+            current_log = (record.module, record.levelno, record.msg)
+            if current_log != getattr(self, "last_log", None):
+                self.last_log = current_log
+                return True
+            return False
+
+    if dev_mode:
+        format_str = "%(message)s"  # RichHandler will handle the formatting
+        logging.basicConfig(
+            level=logging.DEBUG,
+            format=format_str,
+            datefmt="%a %H:%M:%S",
+            handlers=[RichHandler(rich_tracebacks=True, markup=True)],
+        )
+    else:
+        format_str = (
+            "[bold pale_green3]%(asctime)s[/bold pale_green3] | "
+            "[light_steel_blue]%(levelname)-8s[/light_steel_blue] | "
+            "%(message)s"
+        )
+        logging.basicConfig(
+            filename=log_file,
+            format=format_str,
+            datefmt="%a %H:%M:%S",
+            encoding="utf-8",
+            level=logging.DEBUG,
+        )
+
+    logger = logging.getLogger()
+    logger.addFilter(DuplicateFilter())
+
+
+def check_for_xdotool():
+    """Check if xdotool is in path and exit if not"""
+    result = shutil.which("xdotool")
+    if not result:
+        logging.critical("Please install xdotool")
+        sys.exit(1)
+
+
+def main_daemon(config):
+    if os.geteuid() == 0:
+        newpid = os.fork()
+        if newpid == 0:
+            init_logging(config.daemon_log_file)
+            antidrift.daemon.run(config)
+    else:
+        if sys.argv[0] == "antidrift":
+            cmd = ["sudo", "antidrift", "--daemon"]
+            subprocess.Popen(cmd)
+        else:
+            init_logging(config.daemon_log_file, dev_mode=True)
+            logging.warning("[red]Running in development mode.[/red]")
+            antidrift.daemon.run(config)
+
+
+def main() -> None:
+    """Main routine that dispatches to client or daemon"""
+    config = Config.load(os.path.expanduser("~/.config/antidrift.yaml"))
+    check_for_xdotool()
+    args = get_args()
+    if args.daemon:
+        main_daemon(config)
+    else:
+        antidrift.client.run(args, config)
+
+
+if __name__ == "__main__":
+    main()