#!/usr/bin/python3
# -*- coding: utf-8 -*-

import argparse
import collections
import pickle
import errno
import glob
import os
import re
import subprocess
import tempfile


class HelperError(subprocess.CalledProcessError):
    def __init__(self, returncode, cmd, output=None, erroutput=None):
        subprocess.CalledProcessError.__init__(self, returncode, cmd, output)
        self.erroutput = erroutput

    def __str__(self):
        return "Command '%s' exited with status %d and stderr %r" % (
            self.cmd,
            self.returncode,
            self.erroutput,
        )


class DBusError(HelperError):
    pass


def mapcache(container):
    def f(functocache):
        def actualfunc(*a, **kw):
            cachekey = pickle.dumps((a, kw))
            if cachekey in container:
                res = container[cachekey]
                if isinstance(res, Exception):
                    raise res
                return res
            try:
                v = functocache(*a, **kw)
                container[cachekey] = v
                return v
            except Exception as e:
                container[cachekey] = e
                raise

        return actualfunc

    return f


@mapcache({})
def helper(cmd):
    """Return standard output of command, or raise HelperError."""
    env = dict(os.environ)
    env["LANG"] = "C"

    errfile = tempfile.TemporaryFile()
    try:
        out = subprocess.check_output(
            cmd,
            stderr=errfile,
            env=env,
            universal_newlines=True,
        )
    except subprocess.CalledProcessError as e:
        errfile.seek(0, 0)
        erroutput = errfile.read().decode("utf-8")
        raise HelperError(e.returncode, e.cmd, e.output, erroutput)
    finally:
        errfile.close()
    return out


def package(f):
    """Get the name of the package containing the specified file.

    Returns None when the file is not known to RPM."""
    cmd = ["rpm", "-qf", "--", f]
    try:
        out = helper(cmd)
    except HelperError as e:
        enoent = "No such file or directory\n"
        if e.returncode == 1 and e.erroutput.endswith(enoent):
            return None
        raise
    if out.endswith("is not owned by any package\n"):
        return None
    return out[:-1]


def unit_name(pid):
    """Get the systemd unit name associated with the PID.  Returns None
    if no unit associated, or no Python D-Bus support on this system."""
    cmd = [
        "dbus-send",
        "--print-reply=literal",
        "--system",
        "--dest=org.freedesktop.systemd1",
        "/org/freedesktop/systemd1",
        "org.freedesktop.systemd1.Manager.GetUnitByPID",
        "uint32:%d" % pid,
    ]
    try:
        out = helper(cmd)
    except HelperError as e:
        loadedunit = "does not belong to any loaded unit"
        if e.returncode == 1 and loadedunit in e.erroutput:
            return None
        raise DBusError(e.returncode, e.cmd, e.output, e.erroutput)

    unit_path = out.strip()

    cmd = [
        "dbus-send",
        "--print-reply=literal",
        "--system",
        "--dest=org.freedesktop.systemd1",
        unit_path,
        "org.freedesktop.DBus.Properties.Get",
        "string:org.freedesktop.systemd1.Unit",
        "string:Id",
    ]
    try:
        out = helper(cmd)
    except HelperError as e:
        raise DBusError(e.returncode, e.cmd, e.output, e.erroutput)

    return re.sub(" +variant +", "", out.rstrip())


def has_changed(filename, old_inode):
    """Given a path name, and an old inode, will return (False, old_inode)
    if the file has not changed on disk, or (True, new_inode) if it has
    or (True, None) if the file was deleted."""
    try:
        stres = os.stat(filename)
    except OSError as e:
        if e.errno == errno.ENOENT:
            return True, None
        raise
    if stres.st_ino != old_inode:
        return True, stres.st_ino
    return False, old_inode


def get_process_info(pid):
    contents = open("/proc/%d/cmdline" % pid).read().split("\0")
    if contents[-1] == "":
        contents.pop()
    return " ".join(contents).replace("\n", " ")


def memoize(f, cache):
    def g(*args):
        if args in cache:
            return cache[args]
        r = f(*args)
        cache[args] = r
        return r

    return g


def splitinodepath(buf):
    inode = []
    buf = iter(buf)
    for character in buf:
        if character == " ":
            break
        inode.append(character)
    inode = int("".join(inode))
    for character in buf:
        if character != " ":
            break
    if character != "":
        path = [character]
        for character in buf:
            path.append(character)
        path = "".join(path[:-1])
    else:
        path = ""
    return inode, path


def get_needs_restarting(ignore_prefixes=None):
    """Gets list of processes that need restarting, each with its reasons.

    Args:
        ignore_prefixes: iterable of path prefixes to ignore in the analysis
    """
    if ignore_prefixes is None:
        ignore_prefixes = []
    else:
        ignore_prefixes = list(ignore_prefixes)

    cache = dict()
    hc = memoize(has_changed, cache)
    to_restart = collections.OrderedDict()
    for mapfile in glob.glob("/proc/*/maps"):
        try:
            pid = int(mapfile.split("/")[2])
        except ValueError:
            continue

        def skipenoentopen(f):
            try:
                return open(f)
            except IOError as e:
                if e.errno != errno.ENOENT:
                    raise
            return []

        maps = (
            (x[1], splitinodepath(x[4]))
            for x in (x.split(" ", 4) for x in skipenoentopen(mapfile) if x)
            if "x" in x[1]
        )
        map_to_inode = collections.OrderedDict(
            (x[1][1], x[1][0])
            for x in maps
            if int(x[1][0]) > 0
            and not x[1][1].startswith("[")
            and not any(x[1][1].startswith(p) for p in ignore_prefixes)
        )
        for name, inode in list(map_to_inode.items()):
            if name.endswith(" (deleted)") and not os.path.exists(name):
                name = name[:-10]
            if pid in to_restart and name in to_restart[pid]:
                continue
            changed, new_inode = hc(name, inode)
            if changed:
                if pid not in to_restart:
                    to_restart[pid] = collections.OrderedDict()
                if new_inode is None:
                    to_restart[pid][name] = "was deleted"
                else:
                    to_restart[pid][name] = "changed inode from %s to %s" % (
                        inode,
                        new_inode,
                    )
    return to_restart


def add_package_info(needs_restarting):
    """Take the result of needs_restarting and substitute files
    for package names when possible."""
    new_needs_restarting = collections.OrderedDict()
    for pid, filesandreasons in list(needs_restarting.items()):
        new_needs_restarting[pid] = collections.OrderedDict()
        for name, reason in list(filesandreasons.items()):
            already = set()
            inpkg = False
            pkg = package(name)
            if pkg:
                inpkg = True
                if pkg not in already:
                    u = "was installed as an update"
                    new_needs_restarting[pid][pkg] = u
                    already.add(pkg)
            if not inpkg:
                new_needs_restarting[pid][name] = reason
    return new_needs_restarting


def aggregate_by_systemd_service(needs_restarting):
    by_service = collections.OrderedDict()
    for pid, files in list(needs_restarting.items()):
        unitname = unit_name(pid)
        if unitname not in by_service:
            by_service[unitname] = []
        by_service[unitname].append((pid, files))
    return by_service


def get_parser():
    parser = argparse.ArgumentParser(
        description=(
            "Show services in need of restart and processes "
            "keeping files open that were deleted or updated "
            "since they were opened by those processes."
        )
    )
    parser.add_argument(
        "-b",
        help="bare form – only list the units that need restarting",
        action="store_true",
        dest="bare",
        default=False,
    )
    parser.add_argument(
        "-p",
        help="show updated packages instead of updated paths whenever "
        "possible; has no effect when -b is specified",
        action="store_true",
        dest="pkgs",
        default=False,
    )
    parser.add_argument(
        "-s",
        help="only list system units – ignore user and session ones",
        action="store_true",
        dest="system_only",
        default=False,
    )
    parser.add_argument(
        "-u",
        help="ignore the systemd unit containing needs-restart",
        action="store_true",
        dest="ignore_self_and_parents",
        default=False,
    )
    parser.add_argument(
        "-i",
        metavar="PATH_PREFIX",
        help=(
            "ignore deleted / modified mmap()ed path(s) starting "
            "with this prefix; you can specify this parameter "
            "multiple times"
        ),
        action="append",
        dest="ignore_prefixes",
    )
    return parser


def get_column_width():
    try:
        rows, columns = subprocess.check_output(
            ["stty", "size"], stderr=open(os.devnull, "w")
        ).split()
        columns = int(columns)
    except subprocess.CalledProcessError:
        columns = None
    return columns


def main():
    args = get_parser().parse_args()
    columns = get_column_width()

    needs_restarting = get_needs_restarting(args.ignore_prefixes)
    if args.pkgs:
        needs_restarting = add_package_info(needs_restarting)
    by_service = aggregate_by_systemd_service(needs_restarting)
    my_unit_name = unit_name(os.getpid())
    for service, processes in list(by_service.items()):
        if (
            args.system_only
            and service
            and (service.startswith("user@") or service.startswith("session-"))
        ):
            continue
        if args.ignore_self_and_parents and my_unit_name == service:
            continue
        if args.bare:
            if service is not None:
                print(service)
        else:
            if service is None:
                service = "(no unit)"
            print("* unit: %s" % service)
            for pid, files in processes:
                print("  * PID:", pid)
                cmd = get_process_info(pid)
                line = "    command: %s" % cmd
                if columns is not None and len(line) > columns:
                    line = line[: columns - 3] + "..."
                print(line)
                for f, reason in list(files.items()):
                    print("    *", f, reason)


if __name__ == "__main__":
    main()
