dron 0.1.20241008__py3-none-any.whl

dron/launchd.py

@@ -0,0 +1,415 @@
+from __future__ import annotations
+
+import itertools
+import json
+import os
+import re
+import shlex
+import sys
+import textwrap
+from datetime import timedelta
+from pathlib import Path
+from subprocess import PIPE, Popen, check_call, check_output
+from tempfile import TemporaryDirectory
+from typing import Any, Iterator, Sequence
+
+from .api import (
+    OnCalendar,
+    OnFailureAction,
+    When,
+)
+from .common import (
+    ALWAYS,
+    MANAGED_MARKER,
+    Command,
+    LaunchdUnitState,
+    MonitorEntry,
+    MonitorParams,
+    State,
+    Unit,
+    UnitFile,
+    logger,
+    unwrap,
+)
+
+# TODO custom launchd domain?? maybe instead could do dron/ or something?
+_LAUNCHD_DOMAIN = f'gui/{os.getuid()}'
+
+
+# in principle not necessary...
+# but makes it much easier to filter out logs & lobs from launchctl dump
+DRON_PREFIX = 'dron.'
+
+
+def _launchctl(*args: Path | str) -> list[Path | str]:
+    return ['launchctl', *args]
+
+
+def _launch_agent(path: str) -> Path:
+    # symlink for autostart
+    assert path.endswith('.plist'), path  # meh
+    assert not Path(path).is_absolute(), path
+
+    LA = Path('~/Library/LaunchAgents').expanduser()
+    link = LA / path
+    return link
+
+
+def fqn(name: Unit) -> str:
+    return _LAUNCHD_DOMAIN + '/' + DRON_PREFIX + name
+
+
+def launchctl_load(*, unit_file: UnitFile) -> None:
+    # bootstrap is nicer than load
+    # load is super defensive, returns code 0 on errors
+    check_call(_launchctl('bootstrap', _LAUNCHD_DOMAIN, unit_file))
+    _launch_agent(unit_file.name).symlink_to(unit_file)
+
+
+def launchctl_unload(*, unit: Unit) -> None:
+    # bootout is more verbose than unload
+    # in addition unload is super defensive, returns code 0 on errors
+    check_call(_launchctl('bootout', fqn(unit)))
+    _launch_agent(unit + '.plist').unlink()
+
+
+def launchctl_kickstart(*, unit: Unit) -> None:
+    check_call(_launchctl('kickstart', fqn(unit)))
+
+
+def launchctl_reload(*, unit: Unit, unit_file: UnitFile) -> None:
+    # don't think there is a better way?
+    launchctl_unload(unit=unit)
+    launchctl_load(unit_file=unit_file)
+
+
+def launchd_wrapper(*, job: str, on_failure: list[str]) -> list[str]:
+    # fmt: off
+    return [
+        sys.executable,
+        '-m',
+        'dron.launchd_wrapper',
+        *itertools.chain.from_iterable(('--notify', n) for n in on_failure),
+        '--job', job,
+        '--',
+    ]
+    # fmt: on
+
+
+def remove_launchd_wrapper(cmd: str) -> str:
+    if ' dron.launchd_wrapper ' not in cmd:
+        return cmd
+    # uhh... not super reliable, but this is only used for monitor so hopefully fine
+    [_, cmd] = cmd.split(' -- ', maxsplit=1)
+    return cmd
+
+
+def plist(
+    *,
+    unit_name: str,
+    command: Command,
+    on_failure: Sequence[OnFailureAction],
+    when: When | None=None,
+) -> str:
+    # TODO hmm, kinda mirrors 'escape' method, not sure
+    cmd: Sequence[str]
+    if isinstance(command, (list, tuple)):
+        cmd = tuple(map(str, command))
+    elif isinstance(command, Path):
+        cmd = [str(command)]
+    elif isinstance(command, str) and ' ' not in command:
+        cmd = [command]
+    else:
+        # unquoting and splitting is way trickier than quoting and joining...
+        # not sure how to implement it p
+        # maybe we just want bash -c in this case, dunno how to implement properly
+        raise RuntimeError(command)
+    del command
+
+    mschedule = ''
+    if when is None:
+        # support later
+        raise RuntimeError(unit_name)
+
+    if when == ALWAYS:
+        mschedule = '<key>KeepAlive</key>\n<true/>'
+    else:
+        assert isinstance(when, OnCalendar), when
+        # https://www.freedesktop.org/software/systemd/man/systemd.time.html#
+        seconds = {
+            'minutely': 60,
+            'hourly'  : 60 * 60,
+            'daily'   : 60 * 60 * 24,
+        }.get(when)
+        if seconds is None:
+            # ok, try systemd-like spec..
+            specs = [
+                (re.escape('*:0/')   + r'(\d+)', 60),
+                (re.escape('*:*:0/') + r'(\d+)', 1 ),
+            ]
+            for rgx, mult in specs:
+                m = re.fullmatch(rgx, when)
+                if m is not None:
+                    num = m.group(1)
+                    seconds = int(num) * mult
+                    break
+        if seconds is None:
+            # try to parse as hh:mm at least
+            m = re.fullmatch(r'(\d\d):(\d\d)', when)
+            assert m is not None, when
+            hh = m.group(1)
+            mm = m.group(2)
+            mschedule = '\n'.join([
+                '<key>StartCalendarInterval</key>',
+                '<dict>',
+                '<key>Hour</key>'  , f'<integer>{int(hh)}</integer>',
+                '<key>Minute</key>', f'<integer>{int(mm)}</integer>',
+                '</dict>',
+            ])
+        else:
+            mschedule = '\n'.join(('<key>StartInterval</key>', f'<integer>{seconds}</integer>'))
+
+    assert mschedule != '', unit_name
+
+    # meh.. not sure how to reconcile it better with systemd
+    on_failure = [
+        x.replace('--job %n', f'--job {unit_name}') + ' --stdin'
+        for x in on_failure
+    ]
+
+    # attempt to set argv[0] properly
+    # hmm I was hoping it would make desktop notifications ('background service added' nicer)
+    # but even after that it still only shows executable script name. ugh
+    # program_argv = (unit_name, *cmd[1:])
+    program_argv = (
+        *launchd_wrapper(job=unit_name, on_failure=on_failure),
+        *cmd,
+    )
+    del cmd
+    program_argvs = '\n'.join(f'<string>{c}</string>' for c in program_argv)
+
+    # TODO add log file, although mailer is already capturing stdout
+    # TODO hmm maybe use the same log file for all dron jobs? would make it easier to rotate?
+    res = f'''
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+
+    <key>Label</key>
+    <string>{DRON_PREFIX}{unit_name}</string>
+    <key>ProgramArguments</key>
+    <array>
+{textwrap.indent(program_argvs, " " * 8)}
+    </array>
+
+    <key>RunAtLoad</key>
+    <true/>
+
+{textwrap.indent(mschedule, " " * 8)}
+
+    <key>Comment</key>
+    <string>{MANAGED_MARKER}</string>
+</dict>
+</plist>
+'''.lstrip()
+    return res
+
+
+from .common import LaunchdUnitState
+
+
+def launchd_state(*, with_body: bool) -> Iterator[LaunchdUnitState]:
+    # sadly doesn't look like it has json interface??
+    dump = check_output(['launchctl', 'dumpstate']).decode('utf8')
+
+    name: str | None = None
+    extras: dict[str, Any] = {}
+    arguments: list[str] | None = None
+    all_props: str | None = None
+    fields = [
+        'path',
+        'last exit code',
+        'pid',
+        'run interval',
+    ]
+    for line in dump.splitlines():
+        if name is None:
+            # start of job description group
+            name = line.removesuffix(' = {')
+            all_props = ''
+            continue
+        elif line == '}':
+            # end of job description group
+            path: str | None = extras.get('path')
+            if path is not None and 'dron' in path:
+                # otherwsie likely some sort of system unit
+                unit_file = Path(path)
+                body = unit_file.read_text() if with_body else None
+
+                # TODO extract 'state'??
+
+                periodic_schedule = extras.get('run interval')
+                calendal_schedule = 'com.apple.launchd.calendarinterval' in unwrap(all_props)
+
+                schedule: str | None = None
+                if periodic_schedule is not None:
+                    schedule = 'every ' + periodic_schedule
+                elif calendal_schedule:
+                    # TODO parse properly
+                    schedule = 'calendar'
+                else:
+                    # NOTE: seems like keepalive attribute isn't present in launcd dumpstate output
+                    schedule = 'always'
+
+                yield LaunchdUnitState(
+                    unit_file=Path(path),
+                    body=body,
+                    cmdline=tuple(extras['arguments']),
+                    # might not be present when we killed process manually?
+                    last_exit_code=extras.get('last exit code'),
+                    # pid might not be present (presumably when it's not running)
+                    pid=extras.get('pid'),
+                    schedule=schedule,
+                )
+            name = None
+            all_props = None
+            extras = {}
+            continue
+
+        all_props = unwrap(all_props) + line + '\n'
+
+        if arguments is not None:
+            if line == '\t}':
+                extras['arguments'] = arguments
+                arguments = None
+            else:
+                arg = line.removeprefix('\t\t')
+                arguments.append(arg)
+        else:
+            xx = line.removeprefix('\t')
+            for f in fields:
+                zz = f'{f} = '
+                if xx.startswith(zz):
+                    extras[f] = xx.removeprefix(zz)
+                    break
+            # special handling..
+            if xx.startswith('arguments = '):
+                arguments = []
+
+
+def verify_unit(*, unit_name: str, body: str) -> None:
+    with TemporaryDirectory() as tdir:
+        tfile = Path(tdir) / unit_name
+        tfile.write_text(body)
+        check_call([
+            'plutil', '-lint',
+            '-s',  # silent on success
+            tfile,
+        ])
+
+
+def cmd_past(unit: Unit) -> None:
+    sub = fqn('dron.' + unit)
+    cmd = [
+        # todo maybe use 'stream'??
+        'log', 'show', '--info',
+        # '--last', '24h',
+        # hmm vvv that doesn't work, if we pass pid, predicate is ignored?
+        # '--process', '1',
+        # hmm, oddly enough "&&" massively slows the predicate??
+        #'--predicate', f'processIdentifier=1 && (subsystem contains "gui/501/dron.{unit}")',
+        '--predicate', f'subsystem contains "{sub}"',
+        '--style', 'ndjson',
+        '--color', 'always',
+    ]
+    with Popen(cmd, stdout=PIPE, encoding='utf8') as p:
+        out = p.stdout; assert out is not None
+        for line in out:
+            j = json.loads(line)
+            if j.get('finished') == 1:
+                # last event at the very end
+                continue
+            subsystem = j['subsystem']
+            # sometimes subsystem contains pid at the end, need to chop it off
+            # also that's wjy we can't use "subsystem = " predicate :(
+            subsystem = subsystem.split(' ')[0]
+            if sub != subsystem:
+                continue
+            msg = j['eventMessage']
+
+            interesting = re.search(' spawned .* because', msg) or 'exited ' in msg
+            if not interesting:
+                continue
+            ts = j['timestamp']
+            print(ts, sub, msg)
+
+
+def cmd_run(*, unit: Unit, do_exec: bool) -> None:
+    if not do_exec:
+        launchctl_kickstart(unit=unit)
+        return
+
+    states = []
+    for s in launchd_state(with_body=False):
+        if s.unit_file.stem == unit:
+            states.append(s)
+    [state] = states
+    cmdline = state.cmdline
+    assert cmdline is not None, unit
+
+    ## cut off launchd wrapper
+    sep_i = cmdline.index('--')
+    cmdline = cmdline[sep_i + 1 :]
+    ##
+
+    cmds = ' '.join(map(shlex.quote, cmdline))
+    logger.info(f'running: {cmds}')
+    os.execvp(
+        cmdline[0],
+        list(cmdline),
+    )
+
+
+def get_entries_for_monitor(managed: State, *, params: MonitorParams) -> list[MonitorEntry]:
+    # for now kinda copy pasted from systemd
+
+    entries: list[MonitorEntry] = []
+    for s in managed:
+        assert isinstance(s, LaunchdUnitState), s
+
+        unit_file = s.unit_file
+        name = unit_file.name.removesuffix('.plist')
+
+        is_seconds = re.fullmatch(r'every (\d+) seconds', s.schedule or '')
+        if is_seconds is not None:
+            delta = timedelta(seconds=int(is_seconds.group(1)))
+            # meh, but works for now
+            ss = f'every {delta}'
+        else:
+            ss = str(s.schedule)
+
+        schedule = ss
+        command = None
+        if params.with_command:
+            cmdline = s.cmdline
+            assert cmdline is not None, name  # not None for launchd units
+            command = ' '.join(map(shlex.quote, cmdline))
+            command = remove_launchd_wrapper(command)
+
+        status_ok = s.last_exit_code == '0'
+        status = 'success' if status_ok else f'exitcode {s.last_exit_code}'
+
+        pid = s.pid
+
+        entries.append(MonitorEntry(
+            unit=name,
+            status=status,
+            left='n/a',
+            next='n/a',
+            schedule=schedule,
+            command=command,
+            pid=pid,
+            status_ok=status_ok,
+        ))
+    return entries

dron/launchd_wrapper.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+import argparse
+import shlex
+import sys
+from pathlib import Path
+from subprocess import PIPE, STDOUT, Popen
+from typing import Iterator, NoReturn
+
+from loguru import logger
+
+LOG_DIR = Path('~/Library/Logs/dron').expanduser()
+
+
+def main() -> NoReturn:
+    p = argparse.ArgumentParser()
+    p.add_argument('--notify', action='append')
+    p.add_argument('--job', required=True)
+    # hmm, this doesn't work with keyword args??
+    # p.add_argument('cmd', nargs=argparse.REMAINDER)
+    args, rest = p.parse_known_args()
+
+    assert rest[0] == '--', rest
+    cmd = rest[1:]
+
+    notify_cmds = [] if args.notify is None else args.notify
+    job = args.job
+
+    LOG_DIR.mkdir(parents=True, exist_ok=True)
+    log_file = LOG_DIR / f'{job}.log'
+
+    logger.add(log_file, rotation='100 MB')  # todo configurable? or rely on osx rotation?
+
+    # hmm, a bit crap transforming everything to stdout? but not much we can do?
+    captured_log = []
+    try:
+        with Popen(cmd, stdout=PIPE, stderr=STDOUT) as po:
+            out = po.stdout
+            assert out is not None
+            for line in out:
+                captured_log.append(line)
+                sys.stdout.buffer.write(line)
+        rc = po.poll()
+
+        if rc == 0:
+            # short circuit
+            sys.exit(0)
+    except Exception as e:
+        # Popen istelf still fail due to permission denied or something
+        logger.exception(e)
+        captured_log.append(str(e).encode('utf8'))
+        rc = 123
+
+
+    def payload() -> Iterator[bytes]:
+        yield f"exit code: {rc}\n".encode()
+        yield b'command: \n'
+        yield (' '.join(map(shlex.quote, cmd)) + '\n').encode('utf8')
+        yield f'log file: {log_file}\n'.encode()
+        yield b'\n'
+        yield b'output (stdout + stderr):\n\n'
+        # TODO shit -- if multiple notifications, can't use generator for captured_log
+        # unless we notify simultaneously?
+        yield from captured_log
+
+    for line in payload():
+        logger.info(line.decode('utf8').rstrip('\n'))  # meh
+
+    for notify_cmd in notify_cmds:
+        logger.info(f'notifying: {notify_cmd}')
+        try:
+            with Popen(notify_cmd, shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE) as po:
+                sin = po.stdin
+                assert sin is not None
+                for line in payload():
+                    sin.write(line)
+                (sout, serr) = po.communicate()
+                for l in sout.decode('utf8').splitlines():
+                    logger.debug(l)
+                for l in serr.decode('utf8').splitlines():
+                    logger.debug(l)
+            assert po.poll() == 0, notify_cmd
+        except Exception as e:
+            logger.error(f'notificaiton failed: {notify_cmd}')
+            logger.exception(e)
+
+    sys.exit(rc)
+
+
+if __name__ == '__main__':
+    main()

dron/monitor.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+from dataclasses import asdict, fields
+from functools import lru_cache
+from typing import Any
+
+from textual import events, work
+from textual.app import App, ComposeResult
+from textual.widgets import DataTable, Log
+
+from .common import MonitorEntry, MonitorParams, unwrap
+from .dron import get_entries_for_monitor, managed_units
+
+
+# todo try RichLog? https://textual.textualize.io/guide/input
+class LogHeader(Log):
+    DEFAULT_CSS = """
+    LogHeader {
+        height: 5;
+        dock: top;
+    }
+    """
+
+
+MonitorEntries = dict[str, MonitorEntry]
+
+
+@lru_cache(None)
+def get_columns() -> list[str]:
+    cols: list[str] = []
+    for f in fields(MonitorEntry):
+        cols.append(f.name)
+
+    # TODO how to check it statically? MonitorEntry.pid isn't giving anything?
+    cols.remove('pid')  # probs don't need?
+    cols.remove('status_ok')
+    return cols
+
+
+def as_row(entry: MonitorEntry) -> dict[str, Any]:
+    cols = get_columns()
+    res = {k: v for k, v in asdict(entry).items() if k in cols}
+
+    v = res['status']
+    color = 'green' if entry.status_ok else 'red'
+    res['status'] = f'[{color}]' + res['status'] + f'[/{color}]'
+
+    # meh. workaround for next/left being max datetime
+    if res['next'].startswith('9999-'):
+        res['left'] = '--'
+        res['next'] = '[yellow]' + 'never' + '[/yellow]'
+
+    if entry.pid is not None:
+        res['left'] = '--'
+        res['next'] = '[yellow]' + 'running' + '[/yellow]'
+    return res
+
+
+class MonitorApp(App):
+
+    def __init__(self, *, monitor_params: MonitorParams, refresh_every: int) -> None:
+        super().__init__()
+        self.monitor_params = monitor_params
+        self.refresh_every = refresh_every
+
+    def compose(self) -> ComposeResult:
+        # TODO self.log is already defined?? what is it?
+        self.logx = LogHeader(auto_scroll=True, highlight=True)
+
+        # TODO input field to filter out jobs?
+
+        # doesn't do anything??
+        # self.log("HELLOOO")
+
+        self.table: DataTable = DataTable(
+            cursor_type='row',
+            zebra_stripes=True,  # alternating colours
+        )
+
+        # NOTE: useful for debugging
+        # yield self.logx
+        yield self.table
+
+    def on_mount(self) -> None:
+        table = self.table
+
+        for col in get_columns():
+            table.add_column(label=col, key=col)
+
+        # todo how to do async here as well?
+        entries = self.get_entries()
+        self.update(entries)
+
+        self.set_focus(table)
+
+        # TODO run and update it continuously? not sure
+        # what if it isn't computed within interval?
+        self.set_interval(interval=self.refresh_every, callback=self.update_in_background)
+
+    def get_entries(self) -> MonitorEntries:
+        managed = list(managed_units(with_body=False))  # body slows down this call quite a bit
+        entries = get_entries_for_monitor(managed=managed, params=self.monitor_params)
+        return {e.unit: e for e in entries}
+
+    def update(self, entries: MonitorEntries) -> None:
+        table = self.table
+
+        # self.logx.write_line(f"HI {datetime.now().isoformat()}")
+
+        current_rows: set[str] = {unwrap(x.value) for x in table.rows}
+        to_remove = {x for x in current_rows if x not in entries}
+        to_add = {x: entry for x, entry in entries.items() if x not in current_rows}
+        to_update = {x: entry for x, entry in entries.items() if x in current_rows}
+
+        for row_key in to_remove:
+            table.remove_row(row_key=row_key)
+
+        for row_key, entry in to_add.items():
+            vals = as_row(entry).values()
+            table.add_row(*vals, key=row_key)
+
+        for row_key, entry in to_update.items():
+            for col, value in as_row(entry).items():
+                table.update_cell(row_key=row_key, column_key=col, value=value, update_width=True)
+
+        columns = get_columns()
+
+        def sort_key(row):
+            data = dict(zip(columns, row))
+            is_running = 'running' in data['next']
+            failed = 'exit-code' in data['status']
+            return (not is_running, not failed, data['unit'])
+
+        # TODO hmm kinda annoying, doesn't look like it preserves cursor position
+        #  if the item pops on top of the list when a service is running?
+        #  but I guess not a huge deal now
+        table.sort(*columns, key=sort_key)
+
+    @work(exclusive=True, thread=True)
+    def update_in_background(self) -> None:
+        entries = self.get_entries()
+        self.call_from_thread(self.update, entries)
+
+    def on_key(self, event: events.Key) -> None:
+        actions = {
+            'j': self.table.action_cursor_down,
+            'k': self.table.action_cursor_up,
+            'h': self.table.action_cursor_left,
+            'l': self.table.action_cursor_right,
+        }
+        action = actions.get(event.key)
+        if action is not None:
+            action()
+            return
+        if event.key == 'q':
+            self.exit()
+        # TODO would be nice to display log on enter press or something?