#!/usr/bin/env python2.6

"""Wrapper on rpm for verification purposes"""


from __future__ import print_function

import argparse
import logging
import logging.handlers
import os
import psutil
import re
import select
import subprocess
import sys
from time import sleep
import threading


try:
    import rpm
    from rpmUtils.miscutils import getSigInfo
except:
    print("This tool can only run on redhat-derivated hosts")
    sys.exit(-1)


MISSING_RE = re.compile("^missing")
MATCH_RE = re.compile(r'^(?P<size>\.|S)'
                      r'(?P<mode>\.|M)'
                      r'(?P<md5>[\?\.5])'
                      r'(?P<device>\.|D)'
                      r'(?P<link>\.|L)'
                      r'(?P<user>\.|U)'
                      r'(?P<group>\.|G)'
                      r'(?P<mtime>\.|T)'
                      r'(?P<capability>\.|P)'
                      r'  (?P<attribute>[cdglr ]?) '
                      r'(?P<filename>\/.*)$')

TIMEOUT = 3600 * 22
GRACE = 60
SIG_THREASHOLD = 10


def renice(process_id, nice=18, ioclass=psutil.IOPRIO_CLASS_BE, iovalue=6):
    """Change the nice priority and ionice class/priority of a pid"""
    process = psutil.Process(process_id)
    process.set_nice(nice)
    process.set_ionice(ioclass, value=iovalue)


def timeout(process, name):
    """Kill process when runed"""
    if process.poll() is None:
        logging.error('%s timeout', name)
        process.terminate()
    sleep(GRACE)
    if process.poll() is None:
        logging.error('%s timeout HARD', name)
        process.terminate()


class Timed(object):
    """Kill process after delay (plus GRACE)"""
    def __init__(self, delay, process, name):
        self._timer = threading.Timer(delay, timeout, [process, name])
    def __enter__(self):
        self._timer.start()
    def __exit__(self, ex_type, _ex_value, _tb):
        self._timer.cancel()


def process_stdin(failedchecksumfiles, failedlink, failedcapability, line):
    """Process the stdout of rpm --verify --all"""
    if MISSING_RE.match(line):
        #skip missing files
        return
    match = MATCH_RE.match(line)
    if match is None:
        if line == '':
            return
        logging.info("Error parsing ''%s''", line)
        return
    if match.group('attribute') == 'c':
        #skip config files
        return
    if match.group('md5') != '.':
        failedchecksumfiles.append(match.group('filename').strip())
    if match.group('link') != '.':
        failedlink.append(match.group('filename').strip())
    if match.group('capability') != '.':
        failedcapability.append(match.group('filename').strip())


def get_broken_packages(nice=19, ioclass=psutil.IOPRIO_CLASS_BE, iovalue=7):
    """Use rpm --verify --all to check installed packages"""
    failedchecksumfiles = []
    failedlink = []
    failedcapability = []

    process = subprocess.Popen(["rpm", "-V", "--all"],
                               stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)
    renice(process.pid, nice=nice, ioclass=ioclass, iovalue=iovalue)

    with Timed(1.0 * TIMEOUT, process, "rpm -V --all"):
        reads = [process.stdout.fileno(), process.stderr.fileno()]
        while True:
            ret = select.select(reads, [], [])
            for fdesc in ret[0]:
                if fdesc == process.stdout.fileno():
                    process_stdin(failedchecksumfiles, failedlink,
                                  failedcapability,
                                  process.stdout.readline().strip())
                if fdesc == process.stderr.fileno():
                    line = process.stdout.readline().strip()
                    if line != '':
                        logging.error(line)
            if process.poll() != None:
                break
        for line in iter(process.stdout.readline, ''):
            process_stdin(failedchecksumfiles, failedlink,
                          failedcapability, process.stdout.readline().strip())
        for line in iter(process.stderr.readline, ''):
            logging.error(process.stdout.readline().strip())

    return (failedchecksumfiles, failedlink, failedcapability)


def format_package(header):
    """Format a package name: name-version-release"""
    return '{0}-{1}-{2}'.format(header[rpm.RPMTAG_NAME],
                                header[rpm.RPMTAG_VERSION],
                                header[rpm.RPMTAG_RELEASE])


def log_change(desc, filename, header):
    """Log the modification of a package"""
    if header is None:
        package = 'UNKNOWN package'
    else:
        package = 'package {0}'.format(format_package(header))
    logging.info('%s %s from %s has been modified', desc, filename, package)


def analyse_broken(failedchecksumfiles, failedlink, failedcapability):
    """Analyse the output of rpm --verify --all"""
    try:
        transactionset = rpm.TransactionSet()
        headers = transactionset.dbMatch()

        for header in headers:
            for filename in header[rpm.RPMTAG_FILENAMES]:
                if filename in failedchecksumfiles:
                    log_change('File', filename, header)
                    failedchecksumfiles.remove(filename)
                if filename in failedlink:
                    log_change('Link', filename, header)
                    failedlink.remove(filename)
                if filename in failedcapability:
                    log_change('Capacity on file', filename, header)
                    failedcapability.remove(filename)

        for filename in failedchecksumfiles:
            log_change('File', filename, None)
        for filename in failedlink:
            log_change('Link', filename, None)
        for filename in failedcapability:
            log_change('Capacity on file', filename, None)
    except rpm.error:
        logging.exception("Error in the RPM code")


def dict_append(dic, key, val):
    """Add a value to a dict of lists"""
    try:
        dic[key].append(val)
    except KeyError:
        dic[key] = [val]


def extract_sig(sig):
    """Extract the signature ID from signature header"""
    raw = sig[1][2]
    if raw.startswith(' Key ID '):
        return raw[7:].strip()
    return raw.strip()


def get_signing_keys():
    """Get the key used to sign every installed package"""
    sigs = {}
    wrong_sigs = {}
    try:
        transactionset = rpm.TransactionSet()
        headers = transactionset.dbMatch()

        for header in headers:
            if not (header[rpm.RPMTAG_DSAHEADER] or
                    header[rpm.RPMTAG_RSAHEADER]):
                dict_append(wrong_sigs, 'UNSIGNED!', format_package(header))
                continue
            sig = getSigInfo(header)
            if sig[0] == 0:
                dict_append(sigs, extract_sig(sig), format_package(header))
            else:
                dict_append(wrong_sigs, extract_sig(sig),
                            format_package(header))
    except rpm.error:
        logging.exception("Error in the RPM code")
    return (sigs, wrong_sigs)


def process_wrong_sigs(wrong_sigs, ignore_unsigned):
    """Process the packages without signature or broken"""
    try:
        unsigneds = wrong_sigs['UNSIGNED!']
        del wrong_sigs['UNSIGNED!']
    except KeyError:
        unsigneds = []
    if not ignore_unsigned:
        for package in unsigneds:
            logging.info('Unsigned package %s', package)
    for key in wrong_sigs:
        for package in wrong_sigs[key]:
            logging.info('Wrong signature with key %s on package %s',
                         key, package)


def process_sigs(sigs, known_keys=None):
    """Process correctly signed packages' keys: do we trust them?"""
    if known_keys is not None:
        for key in known_keys:
            try:
                del sigs[key]
            except KeyError:
                pass
    for key in sigs:
        if known_keys is None:
            if len(sigs[key]) < SIG_THREASHOLD:
                logging.info('Suspicious key %s, only signing %s packages:',
                             key, len(sigs[key]))
                for package in sigs[key]:
                    logging.info('Suspicious key %s used on package %s', key,
                                 package)
        else:
            for package in sigs[key]:
                logging.info('Unknown key %s used on package %s', key, package)


def get_conf():
    """Parse sys.argv"""
    parser = argparse.ArgumentParser(description='Verify installed RPMS')
    parser.add_argument('--nice', default=19, type=int,
                        help='Set the scheduling priority (default: %(default)s)')
    parser.add_argument('--ioclass', default=psutil.IOPRIO_CLASS_BE, type=int,
                        help='Set the io scheduling class (default: %(default)s)')
    parser.add_argument('--ionice', default=7, type=int,
                        help='Set the io scheduling priority (default: %(default)s)')
    parser.add_argument('--priority-diff', default=1, type=int,
                        help='Priority difference between parsing and hashing')
    parser.add_argument('-f', action='store_true', dest='here',
                        help='Print logs on the screen instead of using syslog')
    parser.add_argument('--disable-hash', action='store_true',
                        help='Disable rpm verification (hashing files)')
    parser.add_argument('--disable-signatures', action='store_true',
                        help='Disable signature verification')
    parser.add_argument('--trusted-keys',
                        help='Coma-separated list of trusted signing keys')
    parser.add_argument('--ignore-unsigned', action='store_true',
                        help='Ignore unsigned rpms (dangerous!)')
    return parser.parse_args()


def do_it(conf):
    """Do all the checks required"""
    renice(os.getpid(), nice=(conf.nice - conf.priority_diff),
           ioclass=conf.ioclass, iovalue=(conf.ionice - conf.priority_diff))
    if not conf.disable_hash:
        (failedchecksumfiles, failedlink, failedcapability) = get_broken_packages(nice=conf.nice, ioclass=conf.ioclass, iovalue=conf.ionice)
        analyse_broken(failedchecksumfiles, failedlink, failedcapability)
    if not conf.disable_signatures:
        (sigs, wrong_sigs) = get_signing_keys()
        process_wrong_sigs(wrong_sigs, conf.ignore_unsigned)
        if conf.trusted_keys:
            process_sigs(sigs, known_keys=conf.trusted_keys.split(','))
        else:
            process_sigs(sigs)


def main():
    """Wrapper arround do_it: conf & log"""
    conf = get_conf()
    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)
    if conf.here:
        handler = logging.StreamHandler()
    else:
        handler = logging.handlers.SysLogHandler(address='/dev/log')
    formatter = logging.Formatter('cern-rpmverify: %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logging.info('started')
    try:
        do_it(conf)
        logging.info('finished')
    except KeyboardInterrupt:
        logging.error('killed')
    except Exception:
        logging.exception('failed')
        raise


if __name__ == '__main__':
    main()
