psdiff

#!/usr/bin/env python
# psdiff (part of ossobv/vcutil) // wdoekes/2016-2020,2021,2023,2025
#   // Public Domain
#
# Generic (coarse) monitoring of daemon processes. Use in conjunction
# with a monitoring suite like Zabbix.
#
import argparse
import collections
import os
import re
import subprocess
import sys
import textwrap
import time
import warnings


DBNAME = '/var/lib/psdiff.db'
DBLINE = u'{indent}{cmdline}  {{user={user}{extra}}}'


def udiff(alines, blines):
    """
    Replacement for difflib.ndiff that doesn't take exponential time.

    The fancy ndiff is nice, but on some machines it takes almost a
    minute for psdiff dumps of 200 lines. That's not acceptable. We'll
    have to settle for something less pretty but faster.

    See: https://bugs.python.org/issue6931

    We'd prefer the vcutils::udiff tool over difflib.unified_diff
    because it syncs on the next line that matches instead of grouping
    removals and additions together.

    vcutil::udiff::

         begin
        +4
        -1
        +5
        -2
        +6
         end

    difflib.unified_diff::

         begin
        -1
        -2
        +4
        +5
        +6
         end

    However, loading udiff optionally makes the output differ based on
    the (non)existence of said tool, and loading it from python is
    rather ugly. We'll settle for ugly-diff for now.
    """
    # # py2/py3 (importlib.util refused to load non-.py files)
    # import imp
    # try:
    #     udifflib = imp.load_source('udiff', './udiff')
    #     udifflib.filediff  # test existence
    # except (AttributeError, IOError, ImportError):
    import difflib

    def diff_func(a, b):
        return difflib.unified_diff(a, b, lineterm='')

    iter_ = iter(diff_func(alines, blines))
    for line in iter_:
        if line.startswith('@@ '):
            break  # skip +++/---
    for line in iter_:
        if not line.startswith('@@ '):
            yield '{} {}'.format(line[0], line[1:])


class Process(object):
    split = re.compile(r'\s+')

    @classmethod
    def from_line(cls, line, root):
        args = cls.split.split(line, 3)

        user = args[0]
        pid = int(args[1])
        ppid = int(args[2])
        exe = args[3][0:8].rstrip()
        assert args[3][8] == ' '
        cmdline = args[3][9:]

        return cls(ppid, pid, user, exe, cmdline, root=root)

    def __init__(self, parent, pid, user, exe, cmdline, root=None):
        self.parent = parent
        self.pid = pid
        self.user = user
        self.exe = exe
        self.cmdline = cmdline
        self.root = root or self
        self.include = True

        if not root:
            self.process_map = {}

        self.root.process_map[pid] = self
        self.children = set()

        # Enriched info
        self._listens = set()

    def skip(self):
        self.include = False

        # Move the _listens to the parent.
        if self._listens:
            obj = self
            while obj:
                obj = obj.parent
                if obj.include:
                    break
                assert not obj._listens, (self, obj)
            obj.add_listens(self._listens)
            self._listens = None  # make it empty/unusable

    def has_parent(self, include_self=False,
                   cmdline__startswith=None, pid=None):
        obj = self
        if not include_self:
            obj = obj.parent

        while obj:
            if (cmdline__startswith is not None and
                    obj.cmdline.startswith(cmdline__startswith)):
                return True
            if pid is not None and obj.pid == pid:
                return True
            obj = obj.parent

        return False

    def fix_links(self):
        if self.parent is not None:
            # Convert ppid to parent.
            self.parent = self.root.process_map[self.parent]
            # Add us as child of the parent.
            self.parent.children.add(self)

    def get_process(self, pid):
        if not pid:
            return None
        return self.root.process_map[pid]

    def to_string(self, indent=0):
        if self._listens:
            # FIXME: TODO: Version/human sorting of these listening addresses.
            listens = list(sorted(self._listens))
            if len(listens) > 1:
                extra = ', listen=({})'.format(','.join(listens))
            else:
                extra = ', listen={}'.format(listens[0])
        else:
            extra = ''
        return DBLINE.format(
            indent=(indent * u'  '), cmdline=self.cmdline.rstrip(),
            user=self.user, extra=extra)

    def sort(self):
        # Sort the children and convert the set into a list.
        for child in self.children:
            child.sort()
        self.children = list(sorted(self.children))

    def add_listens(self, listens):
        self._listens |= set(listens)

    def __hash__(self):
        # Needs to be reimplemented because Python3 drops the
        # auto-generated one when __eq__ is defined.
        return id(self)

    def __eq__(self, other):
        # Only identity comparison yields same.
        return (id(self) == id(other))

    def __lt__(self, other):
        # Quick, check identity:
        if id(self) == id(other):
            return False

        # Lazy comparison.
        if self.cmdline != other.cmdline:
            return (self.cmdline < other.cmdline)
        if self.user != other.user:
            return (self.user < other.user)
        if len(self.children) != len(other.children):
            return (len(self.children) < len(other.children))
        assert isinstance(self.children, list), self.children
        assert isinstance(other.children, list), other.children
        return (self.children < other.children)

    def __str__(self):
        return self.to_string()


class ProcessFormatter(object):
    def __init__(self, root):
        self.root = root

        # Add self.adjust hook to alter process traits before sort.
        self.visit(self.adjust)
        # Skip processes.
        self.visit((lambda process: (
            None if self.include(process) else process.skip())))
        # Sort processes.
        self.visit((lambda process: process.sort()))

    def visit(self, callable_):
        "Visit all processes with callable."
        for process in self.root.process_map.values():
            callable_(process)

    def to_strings(self, process, indent=0):
        "Return a list of stringified children with indentation."
        ret = []
        if process.include:
            ret.append(self.to_string(process, indent))
        for child in process.children:  # has been sorted already
            ret.extend(self.to_strings(child, indent + 1))
        return ret

    def __str__(self):
        return u'\n'.join(self.to_strings(self.root)) + '\n'

    def adjust(self, process):
        """
        The possibility to adjust cmdline and other process traits.

        This is called before sort, so you'll want to use this to alter
        cmdline.
        """
        pass

    def include(self, process):
        "The possibility to exclude processes from the listing."
        return True

    def to_string(self, process, indent=0):
        "The old hook to alter cmdline appearance."
        return process.to_string(indent)


class FilteredProcessFormatter(ProcessFormatter):
    def __init__(self, *args, **kwargs):
        self._include_once = set()
        super(FilteredProcessFormatter, self).__init__(*args, **kwargs)

    def adjust(self, process):
        super(FilteredProcessFormatter, self).adjust(process)

        if process.cmdline.startswith((
                'astcanary',  # astcanary /var/run/asterisk/... <pid>
                '/usr/sbin/amavisd-new ')):
            # These processes have fluctuating arguments. Drop them.
            process.cmdline = process.cmdline.split(' ', 1)[0]

        elif process.cmdline.startswith((
                '/usr/sbin/zabbix_proxy: ',
                '/usr/sbin/zabbix_server: ')):
            # zabbix_proxy and zabbix_server add " [info]" which changes.
            # Drop it.
            process.cmdline = process.cmdline.split(' [', 1)[0]

        elif process.cmdline.startswith((
                '/usr/bin/containerd-shim-runc-v2 ',
                'containerd-shim ',
                'docker-containerd-shim ')):
            # Docker 18+ instances have fluctuating arguments:
            #   /usr/bin/containerd-shim-runc-v2 -namespace moby -id <ID> ...
            #   containerd-shim ... -workdir /var/...containerd/<ID> ...
            args = process.cmdline.split()
            if '-id' in args:
                # -id [ID]
                pos = args.index('-id') + 1
                if pos < len(args):
                    args[pos] = '<ID>'
            if '-workdir' in args:
                # -workdir [PATH]
                # /var/lib/containerd/io.containerd.runtime.v1.linux/moby/<ID>
                pos = args.index('-workdir') + 1
                if pos < len(args):
                    args[pos] = args[pos].rsplit('/', 1)[0] + '/<ID>'
            if len(args) == 4 and args[3] == 'docker-runc':
                # Docker 17- instances have fluctuating arguments:
                #   docker-containerd-shim <ID> /var/...containerd/<ID> \
                #     docker-runc
                args[1] = '<ID>'
                args[2] = args[2].rsplit('/', 1)[0] + '/<ID>'

            process.cmdline = ' '.join(args)

        elif process.exe == 'haproxy':
            # haproxy might listen to a random UDP port because of syslog:
            # > log 10.20.30.40:514 format rfc5424 local0 notice
            # Turn it into a single UDP socket.
            process._listens = set(
                ('udp:0.0.0.0:*' if lst.startswith('udp:0.0.0.0:') else lst)
                for lst in process._listens)

    def include(self, process):
        # Ignore kernel threads.
        if process.has_parent(include_self=True, pid=2):
            return False

        # Systemd renames itself after an update. We can't rename it
        # back to /sbin/init because it may have been called differently
        # (/sbin/init splash or whatever) in the first place.
        elif process.pid == 1:
            # /sbin/init [splash]
            # /lib/systemd/systemd --system --deserialize 19
            process.cmdline = 'INIT'

        # Children of these commands are generally not daemons, skip
        # them:
        elif process.has_parent(include_self=True, cmdline__startswith=(
                'CRON', 'SCREEN', '-tmux',
                '/USR/SBIN/CRON',  # older cron
                '/usr/sbin/CRON',  # newer cron
                # Is a daemon, but spawns children of init for extra work.
                '/usr/bin/python /usr/bin/salt-minion',
                # Comes and goes.
                '/usr/libexec/fwupd/fwupd')):
            return False

        # We want to monitor these daemons, but not their
        # (grand)children, as they come and go:
        elif process.has_parent(include_self=False, cmdline__startswith=(
                'sshd: /usr/sbin/sshd ',
                'gocollect',                        # ubuntu (upstart)
                '/lib/systemd/systemd-udevd',
                '/usr/lib/postfix/master',          # debian/ubuntu
                '/usr/lib/postfix/sbin/master',     # ubuntu16.04+
                '/usr/libexec/postfix/master',      # redhat
                '/usr/lib/postgresql/',
                '/usr/sbin/dovecot',
                '/usr/sbin/gocollect',              # sysv/systemd
                '/usr/sbin/sshd ',
                '/usr/sbin/vsftpd',
                '/usr/sbin/zabbix_agent2',
                '/usr/sbin/zabbix_agentd')):
            return False

        # These children may come and go, but we expect at least one:
        # - multiprocess apache creates at least N processes but may add/remove
        #   based on demand
        elif process.cmdline.startswith((
                '/usr/sbin/apache2 ',       # debian/ubuntu
                '/usr/sbin/httpd ',         # redhat
                'php-fpm: ')):
            key = (process.parent.pid, process.user, process.cmdline)
            if key in self._include_once:
                return False
            else:
                self._include_once.add(key)

        # These ones may come and go. Don't care if they exist.
        # - uuidd gets spawned at will
        # - zfs get can be slow; it's used by various processes (like kubelet)
        elif ((process.user == 'uuidd' and
                process.cmdline == '/usr/sbin/uuidd --socket-activation') or
                (process.user == 'root' and
                 process.cmdline.startswith(('zfs get ', 'zfs list')))):
            return False

        # Special case since Ubuntu noble (systemd 255), we expect this one in
        # combination with sd-pam, but no other children.
        elif process.cmdline in (
                '/lib/systemd/systemd --user',
                '/usr/lib/systemd/systemd --user') and all(
                    child.exe == '(sd-pam)' and not child.children
                    for child in process.children):
            return False

        # (sd-pam) comes and goes. Generally as child of systemd --user.
        elif process.cmdline == process.exe == '(sd-pam)':
            return False

        return super(FilteredProcessFormatter, self).include(process)


def diff(a, b):
    a = a.rstrip().split('\n')  # drop trailing LF
    b = b.rstrip().split('\n')  # drop trailing LF
    if len(a) == 1 and not a[0]:
        a = []
    if len(b) == 1 and not b[0]:
        b = []

    changes = []
    remap = {' ': 0, '-': -1, '+': 1}
    for change in udiff(a, b):
        if change[0] != '?':
            changes.append((remap[change[0]], change[1:]))

    return changes


def ps_faxu(with_network=False):
    cmd = ['ps', 'ax', '-o', 'user,pid,ppid,fname,args']
    try:
        output = subprocess.check_output
    except AttributeError:
        # Blegh. Python 2.6. (You did already `pip install argparse`, yes?)
        proc = subprocess.Popen(cmd, bufsize=-1, stdout=subprocess.PIPE)
        output = proc.communicate()[0]
        proc.wait()
    else:
        output = subprocess.check_output(cmd, bufsize=-1)

    # Get socket data as quickly as possible.
    sockets_by_pid = None
    if with_network:
        try:
            # Get listen addresses per pid in a dictionary:
            # {1234: ['tcp:0.0.0.0:80', 'tcp:0.0.0.0:443']}
            sockets_by_pid = netstat()
        except (subprocess.CalledProcessError, OSError):
            # If there is no 'ss' (super-netstat) or we have too few perms
            # to use it fruitfully, accept it.
            # - CalledProcessError = ss returned non-zero
            # - FileNotFoundError(OSError) = there is no ss
            # - PermissionDeniedError(OSError) = we're not allowed to call ss
            warnings.warn(
                '--net info requested but failed (no perms? no ss(1)?)')

    output = output.decode('ascii', 'replace')

    root = Process(None, 0, 'root', 'root', 'root')
    for i, line in enumerate(output.split('\n')):
        if i == 0 or not line:
            pass
        else:
            Process.from_line(line, root)

    # Update processes with proper links. This must be done last because
    # the process output is unordered and we may not have the parent
    # process info yet earlier.
    for process in root.process_map.values():
        process.fix_links()

    # We have pids and netstat has pids. Amend the processes with socket info.
    if sockets_by_pid is not None:
        process_pids = set(root.process_map.keys())
        socket_pids = set(sockets_by_pid.keys())
        both_pids = process_pids & socket_pids
        for pid in both_pids:
            root.process_map[pid].add_listens(sockets_by_pid[pid])
        for pid in (socket_pids - both_pids):
            root.process_map[0].add_listens(sockets_by_pid[pid])

    return root


def netstat():
    # -H no-header
    # -n numeric
    # -l listen
    # -p with-process
    # -t include-tcp
    # -u include-udp
    # -w include-raw
    # -S include-sctp
    # -0 include-packet
    cmd = ['ss', '-HnlptuwS0']
    try:
        output = subprocess.check_output
    except AttributeError:
        # Blegh. Python 2.6. (You did already `pip install argparse`, yes?)
        proc = subprocess.Popen(cmd, bufsize=-1, stdout=subprocess.PIPE)
        output = proc.communicate()[0]
        proc.wait()
    else:
        output = subprocess.check_output(cmd, bufsize=-1)
    output = output.decode('ascii', 'replace')

    output = output.strip().split('\n')
    lines = [line.split(None, 7) for line in output]
    # Netid  State  Recv-Q  Send-Q  LocalAddr:Port  PeerAddr:Port  Proc
    proc_listens = collections.defaultdict(list)
    for line in lines:
        addr = '{}:{}'.format(line[0], line[4])
        if len(line) < 7:
            # We need root to get the process info. Bail immediately.
            raise OSError(1, 'Operation not permitted')  # EPERM
        users = line[6]
        assert users.startswith('users:('), line
        assert users.endswith(')'), line
        idx = 7
        state = 0
        while True:
            # ("zabbix_agentd",pid=1744582,fd=6)
            if state == 0:
                assert users[idx] == '('
                state = 1
            elif state == 1:
                assert users[idx] == '"'
                state = 2
            elif state == 2:
                if users[idx] == '"':
                    state = 3
            elif state == 3:
                assert users[idx] == ','
                nidx = users.index(')', idx + 1)
                pidinfo = users[idx + 1:nidx].split(',', 1)[0]
                assert pidinfo.startswith('pid='), users
                pid = int(pidinfo[4:])
                proc_listens[pid].append(addr)
                state = 4
                idx = nidx
            elif state == 4:
                if users[idx] == ')':
                    state = 5
                    break
                elif users[idx] == ',':
                    state = 0
            idx += 1

    new_listens = {}
    for key, value in proc_listens.items():
        new_listens[key] = value

    # {fd:   [addr1, addr2]}
    # {1047: ['icmp6:*:58', 'tcp:127.0.0.1:2601']}
    return new_listens


def eval_psdiff_d(dirname):
    sources = [
        os.path.join(dirname, f) for f in os.listdir(dirname)
        if not f.startswith('.') and f.endswith('.py')]

    sources.sort()
    mixins = []
    for source_file in sources:
        try:
            with open(source_file, 'r') as fh:
                source = fh.read()
            # You really don't need FilteredProcessFormatter or
            # ProcessFormatter or eval_psdiff_conf or any of that.
            # If you're going to do really fancy stuff, use custom
            # psdiff.conf and go from there.
            io = {}
            exec(source, io)
            mixins.append(io['ProcessFormatterMixin'])  # one Mixin per file
        except Exception:
            import traceback
            raise ValueError(
                'exec error reading {!r}\n\n    {}'.format(
                    source_file,
                    '\n    '.join(traceback.format_exc().split('\n'))))

    return {'LocalProcessFormatterMixins': mixins}


def eval_psdiff_conf(filename):
    with open(filename, 'r') as fh:
        source = fh.read()
    # Ooohh.. eval/exec. Supply FilteredProcessFormatter and
    # ProcessFormatter so they can be used as superclass.
    # And pass eval_psdiff_d so you can use a custom
    # psdiff.conf and _also_ load psdiff.d files.
    io = {
        'FilteredProcessFormatter': FilteredProcessFormatter,
        'ProcessFormatter': ProcessFormatter,
        'eval_psdiff_conf': eval_psdiff_conf,
        'eval_psdiff_d': eval_psdiff_d,
    }
    exec(source, io)
    return {
        'LocalFilteredProcessFormatter': io['LocalFilteredProcessFormatter']}


def get_formatter_class():
    for path in ('/usr/local/etc/psdiff.conf', '/etc/psdiff.conf'):
        # First check, and then open without exception handling. That way we
        # see if anything is wrong with permissions and such.
        if os.path.exists(path):
            return eval_psdiff_conf(path)['LocalFilteredProcessFormatter']

    # No psdiff.conf? Check for psdiff.d.
    for path in ('/usr/local/etc/psdiff.d', '/etc/psdiff.d'):
        if os.path.exists(path):
            # Mere existence of the path is enough for us to use that: if you
            # create an empty /usr/local/etc/psdiff.d, then /etc/psdiff.d will
            # NOT be used.
            # Don't forget the 'object' in ProcessFormatterMixin(object) for
            # python2.
            mixins = eval_psdiff_d(path)['LocalProcessFormatterMixins']
            class_ = type(
                'LocalFilteredProcessFormatter',
                tuple(mixins + [FilteredProcessFormatter]),
                {})
            return class_

    # Nothing found? Return the plain unaltered version.
    return FilteredProcessFormatter


def get_new_output(formatter_class, with_network=False):
    root = ps_faxu(with_network=with_network)
    formatter = formatter_class(root)
    return formatter.__str__()  # returns unicode(!) on py2


def show_diff(changes, missing=True, extra=True):
    has_diff = False
    for which, line in changes:
        if which and isinstance('', bytes):  # py2
            line = line.encode('utf-8', 'replace')
        if which < 0 and missing:
            show_diff_line('-', line)
            has_diff = True
        elif which > 0 and extra:
            show_diff_line('+', line)
            has_diff = True
    return has_diff


def show_diff_line(prefix, line):
    """
    This is a temporary fix. Instead of parsing the line and
    reformatting, we should parse it earlier and do a diff on the parsed
    line instead.

    Line looks like:

        /usr/sbin/zabbix_agentd --foreground  {user=zabbix}

    Or maybe:

        /usr/sbin/zabbix_agentd --foreground  {user=zabbix, listen=(
            tcp:0.0.0.0:10051)}

    Or maybe with additional properties just before the trailing '}'.

    If it looks like listen=() we list those entries on separate lines.
    """
    line_no_indent = line.lstrip()
    indent = line[0:len(line) - len(line_no_indent)]

    info_idx = line_no_indent.rindex('{')
    paren_idx = line_no_indent.find('(', info_idx)
    if paren_idx != -1:
        lines = [line_no_indent[0:(paren_idx + 1)]]
        lines.extend(
            '{},'.format(line)
            for line in line_no_indent[paren_idx + 1:].split(','))
        lines[-1] = lines[-1][:-1]  # remove excess ','
    else:
        lines = [line_no_indent]

    # Print first line as is, print the rest with extra indentation.
    print('{}{}{}'.format(prefix, indent, lines[0]))
    for line in lines[1:]:
        print('{}{}    {}'.format(prefix, indent, line))


def db_has_network(db):
    """
    Whether the DB file has with_network output in it.

    Useful for backwards compatibility: only show with_network output if
    the DB had it.
    """
    # This is a very cheap/shady check for the newer DB format. It
    # should generally work though. I don't expect that string in normal
    # arguments.
    return bool(', listen=' in db)


def write(output, has_old=False):
    if isinstance('', bytes):  # py2
        output = output.encode('utf-8', 'replace')

    new_name = '{}.new'.format(DBNAME.rsplit('.', 1)[0])
    old_name = '{}.old'.format(DBNAME.rsplit('.', 1)[0])
    with open(new_name, 'w') as fh:
        fh.write(output)
    if has_old:
        os.rename(DBNAME, old_name)
    os.rename(new_name, DBNAME)


def get_argv():
    "Fix up argv"
    argv = sys.argv[:]

    # show-missing or show-extra: convert to show --missing|--extra
    if 'show-missing' in argv[1:]:
        warnings.warn(
            'Replacing show-missing with show --missing', DeprecationWarning)
        idx = argv.index('show-missing', 1)
        if argv[idx] == 'show-missing':
            argv[idx:idx + 1] = ['show', '--missing']
    elif 'show-extra' in argv[1:]:
        warnings.warn(
            'Replacing show-extra with show --extra', DeprecationWarning)
        idx = argv.index('show-extra', 1)
        if argv[idx] == 'show-extra':
            argv[idx:idx + 1] = ['show', '--extra']

    # No args? Needed on python2 where ArgumentParser does not proceed after
    # not finding an action.
    if all(arg.startswith('-') for arg in argv[1:]):
        argv.append('show')

    return argv


def main():
    parser = argparse.ArgumentParser(
        prog='psdiff',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=textwrap.dedent('''\
            Monitor differences between the list of expected running processes
            and the actual running processes.
            '''),
        epilog=textwrap.dedent('''\
            Expected usage
            --------------

            - set up server with various processes;
            - run `psdiff write' to store a dump in /var/lib/psdiff.db;
            - have zabbix (or your favorite monitoring tool) call
              `psdiff show --missing' and `psdiff show --extra';
            - have the monitoring tool show errors if there is output for any
              of the commands.

            This is just a STARTING POINT, it is NOT a replacement for DETAILED
            process monitoring. You will still want to add daemon-specific
            monitoring through other means.


            Adjustments through psdiff.conf
            -------------------------------

            On startup, an attempt is made to import /usr/local/etc/psdiff.conf
            or /etc/psdiff.conf (a python file) where it looks for a class
            called `LocalFilteredProcessFormatter', which will be used as
            formatter class instead of the builtin FilteredProcessFormatter.

            For example:

                class LocalFilteredProcessFormatter(
                        FilteredProcessFormatter):
                    def adjust(self, process):
                        super(LocalFilteredProcessFormatter, self).adjust(
                            process)

                        # haproxy(1) sometimes adds " -sf PIDLIST" at the tail
                        if process.cmdline.startswith('/usr/sbin/haproxy'):
                            process.cmdline = (
                                process.cmdline.split(' -sf ', 1)[0])

                        # Java processes get unordered arguments...
                        if process.cmdline.startswith((
                                '/usr/bin/java', 'java')):
                            args = process.cmdline.split(' ')
                            process.cmdline = ' '.join(
                                [args[0]] + sorted(args[1:]))

                    def include(self, process):
                        # atop(1) has fluctuating arguments. I don't care
                        # whether it runs.
                        if process.cmdline.startswith('/usr/bin/atop '):
                            return False

                        # Don't monitor children of LXD/LXC containers.
                        if process.has_parent(
                                #include_self=True,  # hide the monitor too?
                                cmdline__startswith='[lxc monitor] '):
                            return False

                        return (
                            super(LocalFilteredProcessFormatter, self)
                            .include(process))

                # vim: set syn=python:


            Adjustments through psdiff.d
            ----------------------------

            If psdiff.conf is not found, but a psdiff.d directory is, then all
            .py files will be loaded for ProcessFilteredMixin classes and used
            to create a big LocalFilteredProcessFormatter class through
            multiple inheritance.

            For example /etc/psdiff.d/kvm.py:

                class ProcessFormatterMixin(object):
                    def adjust(self, process):
                        super(ProcessFormatterMixin, self).adjust(process)

                        # For kvm, show only "/usr/bin/kvm -id X -name Y".
                        if process.cmdline.startswith('/usr/bin/kvm '):
                            p = process.cmdline.split()
                            new = [p[0]]
                            [new.extend(p[i:i+2])
                             for i, j in enumerate(p)
                             if j.startswith(('-id', '-name'))]
                            process.cmdline = ' '.join(new)

            Method resolution order (MRO) is alphabetical, with the
            FilteredProcessFormatter base last.


            Zabbix example
            --------------

                UserParameter=psdiff.missing,psdiff show --missing --retry 2>&1
                UserParameter=psdiff.extra,psdiff show --extra --retry 2>&1

            With triggers like this:

                {Template Role Daemons:psdiff.missing.strlen()}<>0 or
                {Template Role Daemons:psdiff.missing.nodata(30m)}=1
            '''))

    if hasattr(argparse, 'BooleanOptionalAction'):
        parser.add_argument(
            '--net', action=argparse.BooleanOptionalAction,
            dest='with_network', help='with network info')
    else:
        parser.add_argument(
            '--net', action='store_true', dest='with_network',
            help='with network info')
        parser.add_argument(
            '--no-net', action='store_false', dest='with_network',
            help='without network info')

    actions = parser.add_subparsers(dest='action', help='subcommand help')
    show_action = actions.add_parser('show', help='show current diff')
    show_action.add_argument(
        '--missing', action='store_true', help='show missing (only)')
    show_action.add_argument(
        '--extra', action='store_true', help='show extra (only)')
    show_action.add_argument(
        '--retry', action='store_true',
        help=("retry up to 2 seconds; avoids false positives caused "
              "by restarts and short lived children"))
    dump_action = actions.add_parser('dump', help='show current state')
    dump_action.add_argument(
        '--verbose', '-v', action='store_true',
        help='show diff format instead of db format')
    write_action = actions.add_parser('write', help='save current state')
    write_action.add_argument(
        '--verbose', '-v', action='store_true',
        help='show which changes are written')
    manual_action = actions.add_parser('manual', help='do manual writes')
    manual_action.add_argument(
        'manual', choices=('add', 'remove'), help='manually add/remove')
    manual_action.add_argument('user', help='process owner')
    manual_action.add_argument('cmdline', help='process cmdline')

    args = parser.parse_args(args=get_argv()[1:])

    # If no action was chosen, default to 'show' and (unfortunately) add
    # the subparser specific arguments.
    if not args.action:
        args.action = 'show'
        args.missing = args.extra = args.retry = False

    # First load up config.
    formatter_class = get_formatter_class()

    # Then load up old db.
    try:
        with open(DBNAME, 'r') as fh:
            old_output = fh.read()
        if isinstance('', bytes):  # py2
            old_output = old_output.decode('utf-8', 'replace')
    except IOError as e:
        if e.errno != 2:  # no such file
            raise
        old_output = u''

    # Python2 compatibility.
    if not hasattr(argparse, 'BooleanOptionalAction'):
        if '--net' not in sys.argv[1:] and '--no-net' not in sys.argv[1:]:
            # Otherwise it will be set to the negation of the first defined
            # argument.
            args.with_network = None

    process(args, old_output, formatter_class)


def process(args, old_output, formatter_class):
    # NOTE: We never print() with u'' below, because in py2 it would "guess"
    # the encoding of the recipient (tty) instead of choosing utf-8.

    if args.action == 'dump':
        if args.with_network is None and db_has_network(old_output):
            args.with_network = True

        new_output = get_new_output(
            formatter_class, with_network=args.with_network)
        if isinstance('', bytes):  # py2
            new_output = new_output.encode('utf-8', 'replace')
        if args.verbose:
            for line in new_output.rstrip().split('\n'):
                show_diff_line('', line)
        else:
            print(new_output[0:-1])  # without trailing LF

    elif args.action == 'write':
        if args.with_network is None and db_has_network(old_output):
            args.with_network = True

        new_output = get_new_output(
            formatter_class, with_network=args.with_network)

        if old_output != new_output:
            write(new_output, has_old=bool(old_output))
            print('Wrote to {}'.format(DBNAME))
            if args.verbose:
                changes = diff(old_output, new_output)
                show_diff(changes)
        else:
            print('No changes to {}'.format(DBNAME))

    elif args.action == 'manual':
        if not old_output:
            raise ValueError('add/remove only works for existing db')

        # For now, only accept children of init (i.e. depth 4)
        indent = u'    '  # root -> INIT -> (process)
        diff_line = DBLINE.format(
            indent=indent, cmdline=args.cmdline, user=args.user)

        if args.manual == 'add':
            new_lines = []
            for line in old_output.split(u'\n'):
                if (diff_line and (line == '' or (
                        line.startswith(indent)
                        and line[len(indent)] != ' '
                        and diff_line < line))):
                    new_lines.append(diff_line)
                    diff_line = None
                new_lines.append(line)
            new_output = u'\n'.join(new_lines)
        else:  # args.manual == 'remove'
            assert args.manual == 'remove', args
            new_output = u'\n'.join(
                i for i in old_output.split(u'\n') if i != diff_line)

        if old_output == new_output:
            raise ValueError('no change!')
        write(new_output, has_old=True)
        print('Wrote to {}'.format(DBNAME))
        changes = diff(old_output, new_output)
        show_diff(changes)

    else:
        assert args.action.startswith('show'), args

        # Check old output format.
        if db_has_network(old_output):
            if args.with_network is None:
                if os.getuid() == 0:
                    # DB format has network, use it.
                    args.with_network = True
                else:
                    # DB format has network, but we cannot use it.
                    warnings.warn(
                        'Not listing --net info without root powers')

            if not args.with_network:
                # Look at the output without the network info.
                old_output = re.sub(r', listen=[^}]*}', '}', old_output)

        # If neither missing nor extra is set, then we want to see both.
        if not args.missing and not args.extra:
            args.missing = args.extra = True

        # If args.retry, then try fetching a changeset for 5 times
        # before concluding that something really has changed.
        for sleeptime in (0.1, 0.3, 0.6, 1.0, 0):
            new_output = get_new_output(
                formatter_class, with_network=args.with_network)

            # Quick optimization.
            if old_output == new_output:
                changes = ()
                break

            changes = diff(old_output, new_output)

            # There are changes. If we're not retrying, start showing
            # the results.
            if not args.retry or not sleeptime:
                break

            # Do the changes apply to us?
            if args.missing and not any(
                    which < 0 for which, line in changes):
                break
            if args.extra and not any(
                    which > 0 for which, line in changes):
                break

            # Sleep a while.
            time.sleep(sleeptime)

        # Show changes, if any.
        has_diff = show_diff(changes, missing=args.missing, extra=args.extra)
        status = (0, 1)[has_diff]

        sys.exit(status)  # possibly non-zero exit


def formatwarning(message, category, filename, lineno, line=None):
    """
    Override default Warning layout, from:

        /PATH/TO/argv0:326: UserWarning:
            [Errno 2] No such file or directory: '/0.d/05.d'
          warnings.warn(str(e))

    To:

        argv0:330: UserWarning:
            [Errno 2] No such file or directory: '/0.d/05.d'
    """
    return '{basename}:{lineno}: {category}: {message}\n'.format(
        basename=os.path.basename(filename), lineno=lineno,
        category=category.__name__, message=message)
warnings.formatwarning = formatwarning  # noqa


if __name__ == '__main__':
    if not sys.stderr.isatty():
        warnings.filterwarnings(
            'ignore', 'Not listing --net info without root powers')
        warnings.simplefilter('ignore', DeprecationWarning)

    main()

# vim: set ts=8 sw=4 sts=4 et ai: