#!/usr/bin/env python3
#
# Copyright 2017 SUNET. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification, are
# permitted provided that the following conditions are met:
#
#    1. Redistributions of source code must retain the above copyright notice, this list of
#       conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright notice, this list
#       of conditions and the following disclaimer in the documentation and/or other materials
#       provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY SUNET ``AS IS'' AND ANY EXPRESS OR IMPLIED
# WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
# FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SUNET OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# The views and conclusions contained in the software and documentation are those of the
# authors and should not be interpreted as representing official policies, either expressed
# or implied, of SUNET.
#
# Author : Fredrik Thulin <fredrik@thulin.net>
#

import os
import grp
import pwd
import sys
import argparse

import logging
import logging.handlers

_defaults = {'debug': False,
             'syslog': False,
             'unless': False,
             'basedir': '/etc/ssl',
             'owner': 'root',
             'group': 'root',
             'mode': 0o444,
             }


class ArgumentError(Exception):
    pass


class LoadError(Exception):
    pass


def parse_args(defaults):
    """
    Parse the command line arguments
    """
    parser = argparse.ArgumentParser(description = 'Tool to create almost arbitrary cert bundles',
                                     add_help = True,
                                     formatter_class = argparse.ArgumentDefaultsHelpFormatter,
                                     )

    parser.add_argument('--debug',
                        dest = 'debug',
                        action = 'store_true', default = defaults['debug'],
                        help = 'Enable debug operation',
                        )
    parser.add_argument('--syslog',
                        dest = 'syslog',
                        action = 'store_true', default = defaults['syslog'],
                        help = 'Enable syslog output',
                        )
    parser.add_argument('--unless',
                        dest = 'unless',
                        action = 'store_true', default = defaults['unless'],
                        help = 'Puppet exec-condition "unless" mode',
                        )
    parser.add_argument('--basedir',
                        dest = 'basedir',
                        default = defaults['basedir'],
                        help = 'Base directory',
                        metavar = 'PATH',
                        )
    parser.add_argument('--owner',
                        dest = 'owner',
                        default = defaults['owner'],
                        help = 'Output file owner',
                        metavar = 'USER',
                        )
    parser.add_argument('--group',
                        dest = 'group',
                        default = defaults['group'],
                        help = 'Output file group',
                        metavar = 'GROUP',
                        )
    parser.add_argument('--mode',
                        dest = 'mode',
                        default = defaults['mode'],
                        help = 'Output file mode',
                        metavar = 'OCTAL',
                        )
    parser.add_argument('--outfile',
                        dest = 'outfile',
                        default = None,
                        help = 'Output filename',
                        metavar = 'FILENAME',
                        )
    # Positional arguments
    parser.add_argument(dest = 'script',
                        help = 'Script to create output from',
                        metavar = 'SCRIPT',
                        type = str,
                        nargs = '+',
                        )
    res = parser.parse_args()
    # Further parse script into a list of (key, filename) tuples
    script = []
    outfile = None
    for this in res.script:
        if '=' not in this:
            raise ArgumentError('Bad script part {!r}'.format(this))
        key, value = this.split('=')
        filename = _expand_filename(key, value, res.basedir)
        if key == 'out':
            outfile = filename
        else:
            script += [(key, filename)]
    res.script = script
    if not res.outfile and not outfile:
        raise ArgumentError('Outfile need to be given either as --outfile or script out=path')
    res.outfile = outfile
    try:
        res.owner = pwd.getpwnam(res.owner).pw_uid
    except KeyError:
        raise ArgumentError('No such user: {}'.format(res.owner))
    try:
        res.group = grp.getgrnam(res.group).gr_gid
    except KeyError:
        raise ArgumentError('No such group: {}'.format(res.group))
    return res


def main(myname = 'cert-bundler', args = None, logger = None, defaults = _defaults):
    """
    :param myname: String, used for logging
    :param args: Command line arguments
    :param logger: logging logger
    :param defaults: Default command line arguments

    :type myname: string
    :type args: None or argsparse object
    :type logger: logging.logger
    :type defaults: dict
    """
    if not args:
        args = parse_args(defaults)

    # initialize various components
    if not logger:
        level = logging.INFO
        if args.debug:
            level = logging.DEBUG
        logging.basicConfig(level = level, stream = sys.stderr,
                            format = '%(asctime)s: %(threadName)s %(levelname)s %(message)s')
        logger = logging.getLogger(myname)
    if args.debug:
        logger.setLevel(logging.DEBUG)
        # log to stderr when debugging
        formatter = logging.Formatter('%(asctime)s %(name)s: %(levelname)s %(message)s')
        stream_h = logging.StreamHandler(sys.stderr)
        stream_h.setFormatter(formatter)
        logger.addHandler(stream_h)
    if args.syslog:
        syslog_h = logging.handlers.SysLogHandler(address = '/dev/log')
        formatter = logging.Formatter('%(name)s: %(levelname)s %(message)s')
        syslog_h.setFormatter(formatter)
        logger.addHandler(syslog_h)

    logger.debug('Processing script {}'.format(args.script))
    try:
        return _process_script(args, logger)
    except (LoadError, ArgumentError) as exc:
        if args.unless:
            # Puppet won't display errors from --unless commands, so we have to lie slightly here
            # so that Puppet actually tries to run the command without --unless in order to get
            # the failure shown in the Puppet output :(
            logger.info('Got a LoadError in unless mode - letting Puppet (re-)write output file : {!s}'.format(exc))
            return 1
        else:
            logger.error('{}'.format(exc))
        return False


def _process_script(args, logger):
    data, private = load_files(args, logger)
    if args.unless:
        return _check_puppet_unless(args, args.outfile, logger, data)

    data, private = load_files(args, logger)
    perms = 0o444
    if private or 'private/' in args.outfile:
        perms = perms & 0o770
        logger.debug('Using non-world accessible permissions {} for output file {!r}'.format(args.outfile, oct(perms)))

    logger.debug('Writing {} bytes to file {} (perms {})'.format(len(data), args.outfile, oct(perms)))

    fd = os.open(args.outfile + '.tmp', os.O_CREAT | os.O_WRONLY, perms)
    uid = args.owner if args.owner else -1
    gid = args.group if args.group else -1
    os.fchown(fd, uid, gid)
    os.write(fd, data)
    os.close(fd)
    os.rename(args.outfile + '.tmp', args.outfile)

    logger.debug('Wrote {} bytes to file {!r}: {}'.format(len(data), args.outfile, args.script))
    parts = [x for x, _ in args.script]
    logger.info('Wrote {} bytes to file {!r} ({})'.format(len(data), args.outfile, '+'.join(parts)))

    return True


def _check_puppet_unless(args, outfile, logger, data):
    """
    Puppet unless-mode means 'do the exec for real unless the --unless invocation exits zero'.
    :param args:
    :param outfile:
    :param logger:
    :param data:
    :return:
    """
    try:
        contents, _private = _load_one_file('out', outfile, args, logger)
    except LoadError:
        logger.debug('Got LoadError for outfile {!r} - letting Puppet (re-)write it'.format(outfile))
        return 1
    if data == contents:
        logger.info('File {!r} has not be changed'.format(outfile))
        return 0
    logger.info('File {!r} ({} bytes) needs to be updated ({} bytes)'.format(outfile, len(contents), len(data)))
    return 1


def load_files(args, logger):
    """
    Load all the input file contents.

    :param args:
    :param logger:
    :return: All the files contents, in order.
    """
    res = b''
    private = 0
    for (key, filename) in args.script:
        data, this_p = _load_one_file(key, filename, args, logger)
        res += data
        private += this_p
    return res, private


def _load_one_file(key, filename, args, logger):
    """
    Load one files contents.

    The return value `private' is 1 if the contents are believed to be private key material.

    :return: Data, private

    :rtype: bytes, int
    """
    private = 0
    logger.debug('Loading {!r} from file {!r}'.format(key, filename))
    try:
        stat = os.stat(filename)
    except OSError as exc:
        raise LoadError('File {!r} access error: {!s}'.format(filename, exc))
    if not stat.st_size:
        raise LoadError('File {!r} has zero size'.format(filename))
    if key == 'key':
        private = 1
        world = stat.st_mode & 0o007
        if world:
            logger.warning('Key file {!r} has world r/w/x permissions: {}'.format(filename, oct(stat.st_mode)))
    if stat.st_mode & 0o007 == 0:
        # input file has restrictive permissions, that should propagate to outfile
        private = 1
    fd = os.open(filename, os.O_RDONLY)
    res = os.read(fd, 1024 * 1024)
    logger.debug('Read {} bytes from file {}'.format(len(res), filename))
    _private_key = res.find(b' PRIVATE KEY') >= 0
    if not _private_key and res.find(b'-----BEGIN CERTIFICATE-----') < 0:
        raise LoadError('File {!r} does not appear to be a certificate or a private key'.format(filename))
    if _private_key:
        private = 1
    return res, private


def _expand_filename(key, value, basedir):
    if value.startswith('/'):
        if os.path.dirname(value) == '/':
            # A value like /foo.pem turns into /etc/ssl/foo.pem
            return os.path.join(basedir, value)
        # Any value with more than just one leading slash is returned as is
        return value
    if '/' in value:
        # A path like 'certs/ca.crt'
        return os.path.join(basedir, value)
    # For values not containing slashes, guess the directory
    if key == 'key' or value.endswith('.key'):
        return os.path.join(basedir, 'private', value)
    if value.endswith('.crt'):
        return os.path.join(basedir, 'certs', value)
    if value.endswith('.pem'):
        return os.path.join(basedir, value)
    raise ArgumentError('Could not deduce directory for {}={}'.format(key, value))


if __name__ == '__main__':
    try:
        progname = os.path.basename(sys.argv[0])
        res = main(progname)
        if isinstance(res, int):
            sys.exit(res)
        if res:
            sys.exit(0)
        sys.exit(1)
    except KeyboardInterrupt:
        sys.exit(0)
