psdiff

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


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


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

        if not root:
            self.process_map = {}

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

    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):
        return DBLINE.format(
            indent=(indent * u'  '), cmdline=self.cmdline.rstrip(),
            user=self.user)

    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 __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)
        # 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 self.include(process):
            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):
        super(FilteredProcessFormatter, self).__init__(*args, **kwargs)
        self._include_once = set()

    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)

    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',
                # User systemd comes and goes as it pleases with (sd-pam).
                '/lib/systemd/systemd --user')):
            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

        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():
    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)
    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()

    return root


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):
    root = ps_faxu()
    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:
            print('-{0}'.format(line))
            has_diff = True
        elif which > 0 and extra:
            print('+{0}'.format(line))
            has_diff = True
    return has_diff


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[:]

    # No argument: default to show
    if len(argv) == 1:
        argv.append('show')

    # show-missing or show-extra: convert to show --missing|--extra
    if argv[1] == 'show-missing':
        argv[1:2] = ['show', '--missing']
    elif argv[1] == 'show-extra':
        argv[1:2] = ['show', '--extra']

    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
            '''))

    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"))
    actions.add_parser('dump', help='show current state')
    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:])

    # 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''

    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':
        new_output = get_new_output(formatter_class)
        if isinstance('', bytes):  # py2
            new_output = new_output.encode('utf-8', 'replace')
        print(new_output[0:-1])  # without trailing LF

    elif args.action == 'write':
        new_output = get_new_output(formatter_class)
        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

        # 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)

            # 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


if __name__ == '__main__':
    main()
# vim: set ts=8 sw=4 sts=4 et ai: