#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
#   Bacula(R) - The Network Backup Solution
#
#
# This script is a simple key-manager for the Volume Encryption done by the
# Storage Daemon.
#
# One key is automatically generated at LABEL time for every VOLUME and
# is stored in the "KEY_DIR" directory.
#
# The keys are a random sequence of bytes as generated by /dev/urandom.
#
# Two encryption methods are available: AES_128_XTS & AES_256_XTS
#
# A third encryption method, called "NULL", exists for testing purposes only.
#
# The main purpose of this script is to provide an example and illustrate the
# protocol between the key-manager and the Storage Daemon. It may be used in
# a production environment.
#
# Use --help to get help the command line parameters of this script.
# If you modify this script, rename it to avoid the possibility of this script
# being overwritten during an upgrade.
#
# The script gets its input from environment variables and returns its
# output via STDOUT.
#
# The Storage Daemon passes the following environment variables:
#
# - OPERATION: This is can "LABEL" when the volume is labeled. In this case
#    the script should generate a new key for the volume. This variable can
#    also be "READ" when the volume has already been labeled and the Storage
#    Daemon needs the already existing key to read or append data to the volume.
#
# - VOLUME_NAME: This is the name of the volume.
#
# Some variables already exist to support a "Master Key" in the future.
# This feature is not yet supported, but will come later:
#
# - ENC_CIPHER_KEY: This is a base64 encoded version of the key encrypted by
#    the "master key"
#
# - MASTER_KEYID: This is a base64 encoded version of the Key Id of
#    the "master key" that was used to encrypt the ENC_CIPHER_KEY value above.
#
# The Storage Daemon expects some values in return via STDOUT:
#
# - volumename: This is a repetition of the name of the volume that is
#    given to the script. This field is optional and ignored by Bacula.
#
# - cipher: This is the cipher that the Storage Daemon must use.
#    The Storage Daemon knows the following ciphers: AES_128_XTS and AES_256_XTS.
#    Of course the key lengths vary with the cipher.
#
# - cipher_key: This is the symmetric key in base64 format.
#
# - comment: This is a single line of text that is optional and ignored
#    by the SD.
#
# - error: This is a single line error message.
#    This is optional, but when provided, the SD considers that the script
#    returned an error and will display this error in the job log.
#
# The Storage Daemon expects an exit code of 0. If the script exits with a
# different error code, any output is ignored and the Storage Daemon will
# display a generic message with the exit code in the job log.
#
# To return an error to the Storage Daemon, the script must set the "error"
# variable string and return an error code of 0.
#
# Here are some input/output samples to illustrate the script's funtion:
#
#   $ OPERATION=LABEL VOLUME_NAME=Volume0001 ./key-manager.py getkey --cipher AES_128_XTS --key-dir tmp/keys 
#   cipher: AES_128_XTS
#   cipher_key: G6HksAYDnNGr67AAx2Lb/vecTVjZoYAqSLZ7lGMyDVE=
#   volume_name: Volume0001
#
#   $ OPERATION=READ VOLUME_NAME=Volume0001 ./key-manager.py getkey --cipher AES_128_XTS --key-dir tmp/keys 
#   cipher: AES_128_XTS
#   cipher_key: G6HksAYDnNGr67AAx2Lb/vecTVjZoYAqSLZ7lGMyDVE=
#   volume_name: Volume0001
#
#   $ cat tmp/keys/Volume0001 
#   cipher: AES_128_XTS
#   cipher_key: G6HksAYDnNGr67AAx2Lb/vecTVjZoYAqSLZ7lGMyDVE=
#   volume_name: Volume0001
#
#   $ OPERATION=READ VOLUME_NAME=MissingVol ./key-manager.py getkey --cipher AES_128_XTS --key-dir tmp/keys 2>/dev/null
#   error: no key information for volume "MissingVol"
#   $ echo $?
#   0
#
#   $ OPERATION=BAD_CMD VOLUME_NAME=Volume0002 ./key-manager.py getkey --cipher AES_128_XTS --key-dir tmp/keys 2>/dev/null
#   error: environment variable OPERATION invalid "BAD_CMD" for volume "Volume0002"
#   $ echo $?
#   0
#
# ------------
# BEGIN SCRIPT
# ------------
import sys
import logging
import argparse
import re
import os
import base64
import codecs
import random
import tempfile

if sys.version_info[0] < 3:
    # python 2.7
    import ConfigParser as configparser
else:
    # python >= 3.X
    import configparser

# logging.raiseExceptions=False

LOG_FILE="@working_dir@/key-manager.log"
KEY_DIR="@sysconfdir@/keydir"
CONFIG_FILE="@sysconfdir@/key-manager.conf"
GNUPGHOME="@sysconfdir@/gnupg"

# trick to use the .in as a python script
if LOG_FILE.startswith('@'):
    LOG_FILE=os.path.join(tempfile.gettempdir(), 'key-manager.log')
if KEY_DIR.startswith('@'):
    KEY_DIR=os.path.join(tempfile.gettempdir(), 'keydir')
if CONFIG_FILE.startswith('@'):
    CONFIG_FILE=os.path.join(tempfile.gettempdir(), 'key-manager.conf')
if GNUPGHOME.startswith('@'):
    GNUPGHOME=os.path.join(tempfile.gettempdir(), 'gnupg')

MASTER_KEYID_SIZE=20
want_to_have_all_the_same_keys=False
#want_to_have_all_the_same_keys=True
CIPHERS=[ 'NULL', 'AES_128_XTS', 'AES_256_XTS' ]
DEFAULT_CIPHER=CIPHERS[1]
MAX_NAME_LENGTH=128
volume_re=re.compile('[A-Za-z0-9:.-_]{1,128}')

# the config from the configuration file if any
config=None

class CryptoCtx:
    master_key_id=None
    cipher=DEFAULT_CIPHER
    stealth=False
    passphrase=None

# raiseExceptions=False
class MyFileHandler(logging.FileHandler):
    """raise an exception when the format don't match the parameters
       instead of printing an error on stderr
    """
    def emit(self, record):
        """dont use try/except and dont call handleError"""
        try:
            msg = self.format(record)
            stream = self.stream
            stream.write(msg)
            stream.write(self.terminator)
            self.flush()
        except Exception:
            if False:
                self.handleError(record)
            else:
                raise

def escape_volume_name(name):
    escapechar='='
    replace_esc='{}0x{:02x}'.format(escapechar, ord(escapechar))
    replace_colon='{}0x{:02x}'.format(escapechar, ord(':'))
    newname=name.replace(escapechar, replace_esc)
    newname=newname.replace(':', replace_colon)
    return newname

def add_console_logger():
    console=logging.StreamHandler()
    console.setFormatter(logging.Formatter('%(levelname)-3.3s %(filename)s:%(lineno)d %(message)s', '%H:%M:%S'))
    console.setLevel(logging.INFO) # must be INFO for prod
    logging.getLogger().addHandler(console)
    return console

def add_file_logger(filename):
    filelog=logging.FileHandler(filename)
    # %(asctime)s  '%Y-%m-%d %H:%M:%S'
    filelog.setFormatter(logging.Formatter('%(asctime)s %(levelname)-3.3s %(filename)s:%(lineno)d %(message)s', '%H:%M:%S'))
    filelog.setLevel(logging.INFO)
    logging.getLogger().addHandler(filelog)
    return filelog

def volume_regex_type(arg_value):
    if not volume_re.match(arg_value):
        raise argparse.ArgumentTypeError
    return arg_value

def setup_logging(debug, verbose, logfile):
    level=logging.WARNING
    if debug:
        level=logging.DEBUG
    elif verbose:
        level=logging.INFO

    logging.getLogger().setLevel(level)

    if logfile:
        filelog=MyFileHandler(logfile)
        # %(asctime)s  '%Y-%m-%d %H:%M:%S'
        filelog.setFormatter(logging.Formatter('%(asctime)s %(levelname)-3.3s %(filename)s:%(lineno)d %(message)s', '%Y-%m-%d %H:%M:%S'))
        filelog.setLevel(level)
        logging.getLogger().addHandler(filelog)
    else:
        console=logging.StreamHandler()
        console.setFormatter(logging.Formatter('%(levelname)-3.3s %(filename)s:%(lineno)d %(message)s', '%H:%M:%S'))
        console.setLevel(level)
        logging.getLogger().addHandler(console)

def test(args):
    assert 'hello123'==escape_volume_name('hello123')
    assert 'vol=0x3dname=0x3a.-_end'==escape_volume_name('vol=name:.-_end')

def bytes_xor(data, key):
    """trivial encode and decode function"""
    enc=[]
    for i, ch in enumerate(data):
        enc.append(ch ^ key[i%len(key)])
    return bytes(enc)

def check_force_cipher_env(cipher):
    return os.getenv('FORCE_CIPHER', cipher)

def get_crypto_ctx_from_config(args, volume_name, master_keyid=None):
    """ retrieve the master-key defined in the config file or a default CTX
     return
            None : for error
            MasterKey object : the master-key or a default context if no config
    """

    if args.config:
        config=configparser.ConfigParser()
        try:
            config.read(args.config)
        except configparser.ParsingError as exc:
            logging.error("parsing configuration file \"%s\": %s", args.config, str(exc))
            print('error: parsing configuration file \"{}\"\n'.format(args.config))
            return None
        the_section=None
        if master_keyid:
            if config.has_section(master_keyid):
                the_section=master_keyid
            else:
                logging.error("configuration file \"%s\" has no master-key \"%s\"", args.config, master_keyid)
                print('error: configuration file \"{}\" has no master-key \"{}\"\n'.format(args.config, master_keyid))
                return None
        else:
            # search for the section matching the volume
            for section in config.sections():
                try:
                    volume_regex=config.get(section, 'volume_regex')
                except configparser.NoOptionError:
                    logging.debug("ignore section \"%s\"", section)
                    continue
                try:
                    match=re.match(volume_regex, volume_name)
                except re.error:
                    logging.error("regular expression error in configuration file \"%s\" in section \"%s\" : %s", args.config, section, str(exc))
                    print("error: regular expression error in configuration file \"{}\" in section \"{}\" : {}".format(args.config, section, str(exc)))
                    return None
                if match:
                    the_section=section
                    break
            if not the_section:
                logging.debug("no master-key defined for volume \"%s\"", volume_name)

        crypto_ctx=CryptoCtx()
        if the_section==None:
            # no master key
            crypto_ctx.master_key_id=None
            crypto_ctx.cipher=args.cipher
        else:
            crypto_ctx.master_key_id=the_section
            try:
                crypto_ctx.gnupghome=config.get(the_section, 'gnupghome')
                if crypto_ctx.gnupghome.startswith('"') and crypto_ctx.gnupghome.endswith('"'):
                    crypto_ctx.gnupghome=crypto_ctx.gnupghome[1:-1]
            except configparser.NoOptionError:
                crypto_ctx.gnupghome=GNUPGHOME
            try:
                crypto_ctx.cipher=config.get(the_section, 'cipher')
            except configparser.NoOptionError:
                crypto_ctx.cipher=args.cipher
            try:
                crypto_ctx.stealth=config.getboolean(the_section, 'stealth')
            except configparser.NoOptionError:
                pass
            try:
                crypto_ctx.passphrase=config.get(the_section, 'passphrase')
            except configparser.NoOptionError:
                pass
            logging.info("use masterkey %r and cipher \"%s\" for volume \"%s\"", crypto_ctx.master_key_id, crypto_ctx.cipher, volume_name)
    else:
        crypto_ctx=CryptoCtx()
        crypto_ctx.cipher=args.cipher

    return crypto_ctx

def generate_key(crypto_ctx, volume_name):
    if crypto_ctx.cipher=='AES_128_XTS':
        key_size=32
    elif crypto_ctx.cipher=='AES_256_XTS':
        key_size=64
    elif crypto_ctx.cipher=='NULL':
        key_size=16
    else:
        logging.error('unknown cipher %s', crypto_ctx.cipher)
        print('error: unknown cipher %s'.format(crypto_ctx.cipher))
        return None # unknown cipher
    urandom=open('/dev/urandom', 'rb')
    key=urandom.read(key_size)
    if want_to_have_all_the_same_keys:
        key=b'A'*key_size
    key_base64=codecs.decode(base64.b64encode(key))
    r=dict()
    r['cipher']=crypto_ctx.cipher
    r['cipher_key']=key_base64
    r['volume_name']=volume_name
    if crypto_ctx.master_key_id:
        try:
            import gnupg
            gnupg.GPG   # check that we have the module and not the GnuPG directory
        except (ImportError, AttributeError):
            logging.error('module gnupg is not installed')
            print('error: python module gnupg is not installed')
            return None
        gpg=gnupg.GPG(gnupghome=crypto_ctx.gnupghome)
        master_keyid_base64=codecs.decode(base64.b64encode(codecs.encode(crypto_ctx.master_key_id)))
        r['master_keyid']=master_keyid_base64
        enc_key=gpg.encrypt(key, crypto_ctx.master_key_id, armor=False)
        enc_key_base64=codecs.decode(base64.b64encode(enc_key.data))
        r['enc_cipher_key']=enc_key_base64
    return r

def decrypt_key(crypto_ctx, volume_name, enc_cipher_key):
    try:
        import gnupg
        gnupg.GPG   # check that we have the module and not the GnuPG directory
    except (ImportError, AttributeError):
        logging.error('module gnupg is not installed')
        print('error: python module gnupg is not installed')
        return None
    r=dict()
    r['cipher']=crypto_ctx.cipher
    gpg=gnupg.GPG(gnupghome=crypto_ctx.gnupghome)
    master_keyid_base64=codecs.decode(base64.b64encode(codecs.encode(crypto_ctx.master_key_id)))
    r['master_keyid']=master_keyid_base64
    passphrase=crypto_ctx.passphrase
    cipher_key=gpg.decrypt(enc_cipher_key, passphrase=passphrase)
    if cipher_key.ok==False:
        logging.error('decryption error for volume "{}":'.format(volume_name, cipher_key.status))
        print('error: decryption error for volume "{}":'.format(volume_name, cipher_key.status))
        return None
    cipher_key_base64=codecs.decode(base64.b64encode(cipher_key.data))
    r['cipher_key']=cipher_key_base64
    r['volume_name']=volume_name
    return r

def decode_data(data):
    d=dict()
    for line in data.split('\n'):
        if line:
            k, v=line.split(':', 1)
            d[k.strip()]=v.strip()
    return d

def encode_data(dct, exclude=None):
    lines=[]
    for key, value in dct.items():
        if not exclude or not key in exclude:
            lines.append('{}: {}'.format(key, value))
    lines.append('')
    return '\n'.join(lines)

def getkey0(args):
    operation=os.getenv('OPERATION')
    volume_name=os.getenv('VOLUME_NAME')
    if not volume_name:
        logging.error("environment variable VOLUME_NAME missing or empty")
        print('error: environment variable VOLUME_NAME missing or empty\n')
        return 0
    if not operation:
        logging.error("environment variable OPERATION missing or empty")
        print('error: environment variable OPERATION missing or empty\n')
        return 0
    if not operation in [ 'LABEL', 'READ']:
        logging.error("environment variable OPERATION invalid \"%s\" for volume \"%s\"", operation, volume_name)
        print("error: environment variable OPERATION invalid \"{}\" for volume \"{}\"\n".format(operation, volume_name))
        return 0

    enc_cipher_key=os.getenv('ENC_CIPHER_KEY')
    master_keyid=os.getenv('MASTER_KEYID')
    logging.info('getkey OPERATION="%s" VOLUME_NAME="%s"%s%s', operation, volume_name, ' ENC_CIPHER_KEY="{}"'.format(enc_cipher_key) if enc_cipher_key else "", ' MASTER_KEYID="{}"'.format(master_keyid) if master_keyid else "")
    key_filename=os.path.join(args.key_dir, escape_volume_name(volume_name))
    if operation=='LABEL':
        crypto_ctx=get_crypto_ctx_from_config(args, volume_name)
        if crypto_ctx==None:
            return 0 # error reading the config file
        crypto_ctx.cipher=check_force_cipher_env(crypto_ctx.cipher)
        if os.path.isfile(key_filename):
            logging.info("delete old keyfile for volume \"%s\" : %s", volume_name, key_filename)
            os.unlink(key_filename)
        ctx=generate_key(crypto_ctx, volume_name)
        if ctx==None:
            return 0 # error while generating the key (wrong cipher or gnupg not installed)
        logging.info("generate key volume=%s cipher=%s enckey=%s masterkey=%s", ctx['volume_name'], ctx['cipher'], ctx.get('enc_cipher_key', ''), ctx.get('master_keyid', ''))
        if crypto_ctx.stealth:
            # don't keep an un-encrypted version of the cipher_key
            # use the masterkey id to decrypte the enckey
            exclude=set(['cipher_key'])
        else:
            exclude=set()
        data=encode_data(ctx, exclude=exclude)
        f=open(key_filename, 'wt')
        f.write(data)
        f.close()
        output=encode_data(ctx) # including the 'cipher_key'
    elif operation=='READ':
        ctx=dict()
        if os.path.isfile(key_filename):
            # use data in the key file
            data=open(key_filename, 'rt').read()
            ctx=decode_data(data)
        if 'cipher_key' in ctx:
            logging.info("read key volume=%s cipher=%s", ctx['volume_name'], ctx['cipher'])
            output=encode_data(ctx)
        elif not enc_cipher_key:
            logging.error("no cipher key nor encrypted cipher key for volume \"%s\"", volume_name)
            print('error: no cipher key nor encrypted cipher key for volume "{}"'.format(volume_name))
            return 0
        else:
            enc_cipher_key_raw=base64.b64decode(codecs.encode(enc_cipher_key))
            master_keyid_raw=base64.b64decode(codecs.encode(master_keyid))
            master_keyid_ascii=codecs.decode(master_keyid_raw)
            # maybe we can retrieve the passphrase for the master-key
            crypto_ctx=get_crypto_ctx_from_config(args, volume_name, master_keyid_ascii)
            if crypto_ctx==None:
                return 0 # error no master-key
                # maybe the master_keyid from the volume could have done the job
                # if gnupg still remember this master-key despit it has been
                # removed from the key-manager config file
            crypto_ctx.cipher=check_force_cipher_env(crypto_ctx.cipher)
            # use the master-key to decrypt the enc_cipher_key
            ctx=decrypt_key(crypto_ctx, volume_name, enc_cipher_key_raw)
            if ctx==None:
                return 0 # error decrypting key
            logging.info("read key volume=%s cipher=%s cipher_key=%s masterkey=%s", ctx['volume_name'], ctx['cipher'], ctx['cipher_key'], ctx['master_keyid'])
            output=encode_data(ctx)
    else:
        output='error: unknown operation \"%r\"'.format(operation)
    print(output)
    return 0

def getkey(args):
    try:
        getkey0(args)
    except:
        logging.exception("unhandled exception in getkey0")
        sys.exit(1)

mainparser=argparse.ArgumentParser(description='Bacula Storage Daemon key manager ')
subparsers=mainparser.add_subparsers(dest='command', metavar='', title='valid commands')

common_parser=argparse.ArgumentParser(add_help=False)
common_parser.add_argument('--key-dir', '-k', metavar='DIRECTORY', type=str, default=KEY_DIR, help='the directory where to store the keys')
common_parser.add_argument('--config', '-C', metavar='CONFIG', type=str, help='the configuration file')
common_parser.add_argument('--log', metavar='LOGFILE', type=str, default=LOG_FILE, help='setup the logfile')
common_parser.add_argument('--debug', '-d', action='store_true', help='enable debugging')
common_parser.add_argument('--verbose', '-v', action='store_true', help='be verbose')

parser=subparsers.add_parser('getkey', description="Retrieve a key", parents=[common_parser, ],
    help="retrieve a key or generate one if don't exist yet")
parser.add_argument('--cipher', '-c', metavar='CIPHER', choices=CIPHERS, default=DEFAULT_CIPHER, help='set the default cipher in {}'.format(', '.join(CIPHERS)))
parser.set_defaults(func=getkey)

parser=subparsers.add_parser('test', description="Run some internal test of the code")
parser.set_defaults(func=test)

args=mainparser.parse_args()
args._parser=mainparser

setup_logging(getattr(args, 'debug', None), getattr(args, 'verbose', None), getattr(args, 'log', None))

# check for the key_dir directory
if hasattr(args, 'key_dir'):
    if not os.path.exists(args.key_dir):
        try:
            os.makedirs(args.key_dir, 0o700)
        except:
            logging.error('Cannot create the "key" directory %s', args.key_dir)
        else:
            logging.error('The "key" directory don\'t exists. Create directory %s', args.key_dir)
    if not os.path.isdir(args.key_dir):
        logging.error('The "key" directory don\'t exists: %s', args.key_dir)
        mainparser.error('error: path "{}" is not a directory'.format(args.key_dir))
    if not os.access(args.key_dir, os.R_OK|os.W_OK):
        logging.error('The "key" directory is not accessible for READ and WRITE: %s', args.key_dir)
        mainparser.error('error: need read and write access to "{}"'.format(args.key_dir))

# check for the config file
if hasattr(args, 'config'):
    if args.config==None and os.path.exists(CONFIG_FILE):
        args.config=CONFIG_FILE # the default file exists, use it
        logging.debug('Use config file %s', args.config)
    if args.config!=None and (not os.path.exists(args.config) or not os.access(args.config, os.R_OK)):
        logging.error('The config file don\'t exists or cannot be read: %s', args.config)
        mainparser.error('The config file don\'t exists or cannot be read: %s'.format(args.config))

sys.exit(args.func(args))