rguard

#!/bin/env python3

# Copyright 2015, 2017 Ryan Sawhill Aroha <rsaw@redhat.com>
# Copyright 2020. Matthias Homann <palto@mailbox.org>
#
#    This program 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.
#
#    This program 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 <gnu.org/licenses/gpl.html> for more details.

import argparse
import time
import subprocess
import os
import logging
import signal
import sys


prog = "rguard"
vers = {
    "version": "2.0.2",
    "date": "2020/14/08",
}


class RebootGuard:
    """Object for enabling/disabling blocks that prevent shutdown, reboot, etc.

    Keyword arguments:
        - blockMethod -- The method by which shutdown is blocked (has default)

    After initialization, call guard(enforce=True) or guard(enforce=False)
    """

    def __init__(self, blockMethod="systemd_refusemanualstart"):
        """Initialize reboot-blocking method choice.

        rsaw says: Please feel free to add other methods.
        """
        possibleBlockMethods = {
            "systemd_refusemanualstart": (self.reboot_guard_systemd_refusemanualstart),
            # FIXME: Wouldn't be all that useful right now,
            # but it's worth implementing:
            # 'systemd_inhibitorlock': self.reboot_guard_systemd_inhibitorlock,
        }
        if blockMethod in possibleBlockMethods:
            self.blockMethod = possibleBlockMethods[blockMethod]
        else:
            raise ValueError(
                f"Invalid blockMethod ('{blockMethod}'); "
                f"valid choices:\n  {', '.join(possibleBlockMethods.keys())}"
            )

    def guard(self, enforce=True):
        """Enable or disable reboot-guard."""
        self.blockMethod(enforce)

    def reboot_guard_systemd_refusemanualstart(self, enforce=True, targets=None):
        """
        Block/unblock entering targets by way of
        RefuseManualStart=yes/no.
        """
        if not targets:
            targets = [
                "poweroff",
                "reboot",
                "halt",
                # "hibernate",
                # "hybrid-sleep",
                # "suspend",
            ]
        reloadDaemon = False
        for t in targets:
            t += ".target"
            if enforce:
                reloadDaemon |= self.add_systemd_unit_start_block(t)
            else:
                reloadDaemon |= self.del_systemd_unit_start_block(t)
        if reloadDaemon:
            logging.debug("Reloading systemd via: systemctl daemon-reload")
            if subprocess.run(["systemctl", "daemon-reload"]).returncode != 0:
                logging.error("Unexpected error reloading systemd")

    def systemd_unit_is_blocked(self, unit):
        """
        Return True if systemd *unit* is configured with
        RefuseManualStart=yes
        """
        cmd = ["systemctl", "show", unit, "-p", "RefuseManualStart"]
        try:
            out = subprocess.check_output(cmd)
        except subprocess.SubprocessError:
            print("ERR")
            logging.error(f"Unexpected error running: {' '.join(cmd)}")
            return False
        if out.strip() == b"RefuseManualStart=yes":
            return True
        return False

    def add_systemd_unit_start_block(self, unit, cfgFileName="reboot-guard.conf"):
        """
        Return True after successfully adding a block
        to prevent start of *unit*.
        """
        if self.systemd_unit_is_blocked(unit):
            logging.debug(f"Unit {unit} already configured with RefuseManualStart=yes")
        else:
            dir_name = f"/run/systemd/system/{unit}.d"
            cfg = f"{dir_name}/{cfgFileName}"
            if self.mkdir_p(dir_name):
                logging.debug(f"Adding RefuseManualStart=yes to {cfg}")
                try:
                    with open(cfg, "w") as f:
                        f.write("[Unit]\nRefuseManualStart=yes\n")
                except (IOError, PermissionError):
                    logging.error(f"Unexpected error writing to '{cfg}'")
                else:
                    logging.warning(f"☹  Blocked {unit}")
                    return True
        return False

    def del_systemd_unit_start_block(self, unit, cfgFileName="reboot-guard.conf"):
        """Return True after successfully unblocking manual start of *unit*."""
        if not self.systemd_unit_is_blocked(unit):
            logging.debug(f"Unit {unit} already configured with RefuseManualStart=no")
        else:
            dir_name = f"/run/systemd/system/{unit}.d"
            cfg = f"{dir_name}/{cfgFileName}"
            logging.debug(f"Removing {cfg}")
            if self.rm_rf(cfg):
                logging.warning(f"☻  Unblocked {unit}")
                return True
        return False

    def mkdir_p(self, path):
        """Make directory *path*."""
        if subprocess.run(["mkdir", "-p", path]).returncode == 0:
            return True
        else:
            logging.error(f"Unexpected error making directory '{path}'")
            return False

    def rm_rf(self, path):
        """Recursively remove *path*."""
        if subprocess.run(["rm", "-rf", path]).returncode == 0:
            return True
        else:
            logging.error(f"Unexpected error deleting path '{path}'")
            return False


class ConditionChecker:
    """An object to handle arbitrary condition checks.

    Keyword arguments:
        - forbidFiles -- Forbidden files for which to check
        - requireFiles -- Required files for which to check
        - units -- Active systemd units for which to check
        - cmds -- Process cmd names for which to check
        - cmdArgs -- Process cmd name+args for which to check
        - runCmds -- Arbitrary commands to execute (return code is checked)

    Every argument expects a list.

    After initialization, call check_conditions() which will return True only
    if all checks pass.
    """

    def __init__(
        self,
        forbidFiles=None,
        requireFiles=None,
        units=None,
        cmds=None,
        cmdArgs=None,
        runCmds=None,
    ):
        """Initialize condition checks."""
        self.forbidFiles = forbidFiles
        self.requireFiles = requireFiles
        self.units = units
        self.cmds = cmds
        self.cmdArgs = cmdArgs
        self.runCmds = runCmds
        self.lastStatus = None

    def check_conditions(self):
        """Run through all condition checks and return True if all pass."""
        if (
            self.test_forbidden_files()
            or self.test_required_files()
            or self.test_active_units()
            or self.test_running_cmds()
            or self.test_running_cmd_args()
            or self.test_execute_commands()
        ):
            if self.lastStatus or self.lastStatus is None:
                logging.warning("☹  Failed a condition-check")
            else:
                logging.debug("Still failing a condition-check")
            self.lastStatus = False
        else:
            if not self.lastStatus:
                logging.warning("☻  Passed all condition-checks")
            else:
                logging.debug("Still passing all condition-checks")
            self.lastStatus = True
        return self.lastStatus

    def test_forbidden_files(self):
        if self.forbidFiles:
            for f in self.forbidFiles:
                if os.path.exists(f):
                    logging.info(f"✘ Fail: Forbid-file '{f}' shouldn't exist")
                    return True
            logging.info("✔ Pass: No forbid-files exist")
        elif self.lastStatus is None:
            logging.debug("No forbid-file checks defined")
        return False

    def test_required_files(self):
        if self.requireFiles:
            for f in self.requireFiles:
                if not os.path.exists(f):
                    logging.info(f"✘ Fail: Require-file '{f}' should exist")
                    return True
            logging.info("✔ Pass: All require-files exist")
        elif self.lastStatus is None:
            logging.debug("No require-file checks defined")
        return False

    def test_active_units(self):
        if self.units:
            for u in self.units:
                if (
                    subprocess.run(
                        ["systemctl", "is-active", u], capture_output=True
                    ).returncode
                    == 0
                ):
                    logging.info(f"✘ Fail: Unit '{u}' shouldn't be active")
                    return True
            logging.info("✔ Pass: No active units match")
        elif self.lastStatus is None:
            logging.debug("No systemd unit checks defined")
        return False

    def test_running_cmds(self):
        if self.cmds:
            for c in self.cmds:
                if (
                    subprocess.run(
                        ["pgrep", "--exact", c], capture_output=True
                    ).returncode
                    == 0
                ):
                    logging.info(f"✘ Fail: Process named '{c}' shouldn't exist")
                    return True
            logging.info("✔ Pass: No running processes match commands")
        elif self.lastStatus is None:
            logging.debug("No command checks defined")
        return False

    def test_running_cmd_args(self):
        if self.cmdArgs:
            for a in self.cmdArgs:
                if (
                    subprocess.run(
                        ["pgrep", "--full", "--exact", a], capture_output=True
                    ).returncode
                    == 0
                ):
                    logging.info(
                        f"✘ Fail: Process name+args matching '{a}' " f"shouldn't exist"
                    )
                    return True
            logging.info("✔ Pass: No running processes match command args")
        elif self.lastStatus is None:
            logging.debug("No command-args checks defined")
        return False

    def test_execute_commands(self):
        z = {True: "zero", False: "non-zero"}
        if self.runCmds:
            for c in self.runCmds:
                desiredOutcome = True
                if "!" in c[0]:
                    desiredOutcome = False
                    c = c[1:]
                if "@" in c[0]:
                    shell = True
                    c = c[1:]
                    logging.debug(f"Running shell cmd: {c}")
                else:
                    shell = False
                    logging.debug(f"Running cmd: {c}")
                    c = c.split()
                if (
                    subprocess.run(c, capture_output=True, shell=shell).returncode == 0
                ) != desiredOutcome:
                    logging.info(
                        f"✘ Fail: Exec-command {{{c}}} needs to return "
                        f"{z[desiredOutcome]}"
                    )
                    return True
            logging.info("✔ Pass: All exec-command tests return as required")
        elif self.lastStatus is None:
            logging.debug("No exec-command checks defined")
        return False


class GracefulDeath:
    """Catch signals to allow graceful shutdown."""

    def __init__(self, rguard):
        self.rguard = rguard
        catchSignals = [
            signal.SIGHUP,
            signal.SIGINT,
            signal.SIGQUIT,
            signal.SIGUSR1,
            signal.SIGUSR2,
            signal.SIGTERM,
        ]
        for signum in catchSignals:
            signal.signal(signum, self.handler)

    def handler(self, signum, frame):
        self.lastSignal = signum
        self.receivedSignal = True
        if signum in [
            signal.SIGINT,
            signal.SIGQUIT,
            signal.SIGTERM,
        ]:
            logging.warning(f"Gracefully exiting due to receipt of signal {signum}")
            self.rguard.guard(enforce=False)
            sys.exit()
        else:
            logging.warning(f"Ignoring signal {signum}")


class CustomFormatter(argparse.RawDescriptionHelpFormatter):
    """This custom formatter eliminates the duplicate metavar in help lines."""

    def _format_action_invocation(self, action):
        if not action.option_strings:
            (metavar,) = self._metavar_formatter(action, action.dest)(1)
            return metavar
        else:
            parts = []
            if action.nargs == 0:
                parts.extend(action.option_strings)
            else:
                default = action.dest.upper()
                args_string = self._format_args(action, default)
                for option_string in action.option_strings:
                    parts.append("%s" % option_string)
                parts[-1] += " %s" % args_string
            return ", ".join(parts)


helpExamples = f"""EXAMPLES:
  As mentioned above, all of the condition-checking options can be used
  multiple times. All specified checks must pass to allow shutdown.

  {prog} --forbid-file /shutdown/prevented/if/this/file/exists

  {prog} --require-file /shutdown/prevented/until/this/file/exists

  {prog} --cmd some-cmd-name-which-when-running-prevents-shutdown
           * e.g., 'bash' or 'firefox'

  {prog} --args 'syndaemon -i 1.0 -t -K -R'
           * Prevent shutdown if this exact cmd + args are found running

  {prog} --run 'ping -c2 -w1 -W1 10.0.0.1'
           * Allow shutdown only if single command (ping in this case) succeeds

  {prog} --run '!ping -c2 -w1 -W1 10.0.0.1'
           * Allow shutdown only if single command FAILS

  {prog} --run '!findmnt /mnt'
           * Allow shutdown if nothing mounted at mountpoint

  {prog} --run '@some_cmd | some_other_cmd; another_cmd'
           * Allow shutdown only if last cmd in shell succeeds
           * When using '@', full shell syntax is supported
           * e.g.: '|', '&&', '||', ';', '>', '>>', '<', etc

  {prog} --run '!@lsof -i:ssh | grep -q ESTABLISHED'
           * Allow shutdown only if shell commands FAIL
           * In this example, only if there are no established ssh connections

"""


def parse_args():
    """Parse argv into usable input."""
    # Setup argument parser
    description = (
        "Block systemd-initiated shutdown until configurable condition checks pass"
    )
    version = f"{prog} v{vers['version']} last mod {vers['date']}"
    epilog = (
        f"{helpExamples}"
        f"VERSION:\n"
        f"  {version}\n"
        f"  See <https://github.com/palto42/reboot-guard to report bugs or RFEs"
    )

    def fmt(prog):
        return CustomFormatter(prog)

    p = argparse.ArgumentParser(
        prog=prog,
        description=description,
        add_help=False,
        epilog=epilog,
        formatter_class=fmt,
    )
    # Add args
    g0 = p.add_argument_group(
        "SET AND QUIT", description="Execute a single action with no condition checks."
    )
    gg0 = g0.add_mutually_exclusive_group()
    gg0.add_argument(
        "-1",
        "--install-guard",
        action="store_true",
        help="Install reboot-guard and immediately exit",
    )
    gg0.add_argument(
        "-0",
        "--remove-guard",
        action="store_true",
        help="Remove reboot-guard (if present) and immediately exit",
    )
    g1 = p.add_argument_group(
        "CONFIGURE CONDITIONS",
        description=(
            "Establish what condition checks must pass to allow shutdown.\n"
            "Each option may be specified multiple times."
        ),
    )
    g1.add_argument(
        "-f",
        "--forbid-file",
        action="append",
        metavar="FILE",
        help="Prevent shutdown while FILE exists",
    )
    g1.add_argument(
        "-F",
        "--require-file",
        action="append",
        metavar="FILE",
        help="Prevent shutdown until FILE exists",
    )
    g1.add_argument(
        "-u",
        "--unit",
        action="append",
        help="Prevent shutdown while systemd UNIT is active",
    )
    g1.add_argument(
        "-c",
        "--cmd",
        action="append",
        metavar="CMD",
        help="Prevent shutdown while CMD exactly matches at least 1 running process",
    )
    g1.add_argument(
        "-a",
        "--args",
        action="append",
        metavar="ARGS",
        help="Prevent shutdown while ARGS exactly matches the full cmd+args of at least 1 running process",
    )
    g1.add_argument(
        "-r",
        "--run",
        action="append",
        metavar="COMMAND",
        help=(
            "Prevent shutdown while execution of COMMAND fails; prefix COMMAND with '!' to prevent shutdown"
            " while execution succeeds (MORE: Prefix COMMAND with '@' to execute it as a shell command,"
            " e.g., to use pipes '|' or other logic; examples: '@cmd|cmd' or '!@cmd|cmd')"
        ),
    )
    g2 = p.add_argument_group("GENERAL OPTIONS", description="")
    g2.add_argument(
        "-h",
        "--help",
        dest="showHelp",
        action="store_true",
        help="Show this help message and exit",
    )
    g2.add_argument(
        "-i",
        "--interval",
        type=float,
        metavar="SEC",
        default=60,
        help="Modify the sleep interval between condition checks (default: 60 seconds)",
    )
    g2.add_argument(
        "-n",
        "--ignore-signals",
        action="store_true",
        help=(
            "Ignore the most common signals (HUP, INT, QUIT, USR1, USR2, TERM),"
            " in which case a graceful exit requires the next option or else SIGKILL"
        ),
    )
    g2.add_argument(
        "-x",
        "--exit-on-pass",
        action="store_true",
        help="Exit (and remove reboot-guard) the first time all condition checks pass",
    )
    g2.add_argument(
        "-b",
        "--start-blocked",
        action="store_true",
        help="Install reboot-guard at start, regardless of initial conditions check",
    )
    g2.add_argument(
        "-v",
        "--loglevel",
        choices=["debug", "info", "warning", "error"],
        default="warning",
        help="Specify minimum message type to print (default: warning)",
    )
    g2.add_argument(
        "-t",
        "--timestamps",
        action="store_true",
        help="Enable timestamps in message output (not necessary if running as systemd service)",
    )
    opts = p.parse_args()
    if opts.showHelp:
        p.print_help()
        sys.exit()
    return opts


def main():
    """Parse args, check conditions, install reboot-guard, setup infinite loop."""
    opts = parse_args()
    if opts.timestamps:
        logging.basicConfig(
            format="%(asctime)s: %(levelname)s: %(message)s",
            level=opts.loglevel.upper(),
            datefmt="%m/%d %H:%M:%S",
        )
    else:
        logging.basicConfig(
            format="%(levelname)s: %(message)s", level=opts.loglevel.upper()
        )
    if os.geteuid() > 0:
        logging.critical(
            "Preemptive permission denied -- exiting because EUID not root"
        )
        sys.exit(1)
    # Handle install/remove-and-exit options
    if opts.install_guard or opts.remove_guard:
        r = RebootGuard()
        if opts.install_guard:
            r.guard(enforce=True)
        else:
            r.guard(enforce=False)
        logging.info("Exiting due to -1 or -0 option")
        sys.exit()
    c = ConditionChecker(
        forbidFiles=opts.forbid_file,
        requireFiles=opts.require_file,
        units=opts.unit,
        cmds=opts.cmd,
        cmdArgs=opts.args,
        runCmds=opts.run,
    )
    # Initialize reboot-guard
    r = RebootGuard()
    # Initialize signal handler
    _ = GracefulDeath(r)
    # Run initial check
    if c.check_conditions() and not opts.start_blocked:
        if opts.exit_on_pass:
            logging.warning("Exiting due to passed condition checks")
            sys.exit()
    else:
        if opts.start_blocked:
            logging.debug("Forced to start with blocked targets.")
        r.guard(enforce=True)
    # Run forever-loop
    while True:
        time.sleep(opts.interval)
        if c.check_conditions():
            r.guard(enforce=False)
            if opts.exit_on_pass:
                logging.warning("Exiting due to passed condition checks")
                sys.exit()
        else:
            r.guard(enforce=True)


if __name__ == "__main__":
    # Execute main loop if run directly
    main()