# Copyright (C) 2014 by Solly Ross
# Copyright (C) 2010 by the Massachusetts Institute of Technology.
# All rights reserved.

# Export of this software from the United States of America may
#   require a specific license from the United States Government.
#   It is the responsibility of any person or organization contemplating
#   export to obtain such a license before exporting.
#
# WITHIN THAT CONSTRAINT, permission to use, copy, modify, and
# distribute this software and its documentation for any purpose and
# without fee is hereby granted, provided that the above copyright
# notice appear in all copies and that both that copyright notice and
# this permission notice appear in supporting documentation, and that
# the name of M.I.T. not be used in advertising or publicity pertaining
# to distribution of the software without specific, written prior
# permission.  Furthermore if you modify this software you must label
# your software as modified software and not distribute it in such a
# fashion that it might be confused with the original M.I.T. software.
# M.I.T. makes no representations about the suitability of
# this software for any purpose.  It is provided "as is" without express
# or implied warranty.

# Changes from original: modified to work with Python's unittest
import copy
import os
import shutil
import signal
import socket
import string
import sys
import subprocess
import tempfile
import logging

import six

from k5test import _utils


_LOG = logging.getLogger(__name__)


def _cfg_merge(cfg1, cfg2):
    if not cfg2:
        return cfg1
    if not cfg1:
        return cfg2
    result = copy.deepcopy(cfg1)
    for key, value2 in cfg2.items():
        if value2 is None or key not in result:
            result[key] = copy.deepcopy(value2)
        else:
            value1 = result[key]
            if isinstance(value1, dict):
                if not isinstance(value2, dict):
                    raise TypeError("value at key '{key}' not dict: "
                                    "{type}".format(key=key,
                                                    type=type(value2)))
                result[key] = _cfg_merge(value1, value2)
            else:
                result[key] = copy.deepcopy(value2)
    return result


_default_krb5_conf = {
    'libdefaults': {
        'default_realm': '$realm',
        'dns_lookup_kdc': 'false'},
    'realms': {
        '$realm': {
            'kdc': '$hostname:$port0',
            'admin_server': '$hostname:$port1',
            'kpasswd_server': '$hostname:$port2'}}}


_default_kdc_conf = {
    'realms': {
        '$realm': {
            'database_module': 'db',
            'iprop_port': '$port4',
            'key_stash_file': '$tmpdir/stash',
            'acl_file': '$tmpdir/acl',
            'dictfile': '$tmpdir/dictfile',
            'kadmind_port': '$port1',
            'kpasswd_port': '$port2',
            'kdc_ports': '$port0',
            'kdc_tcp_ports': '$port0',
            'database_name': '$tmpdir/db'}},
    'dbmodules': {
        'db_module_dir': os.path.join(_utils.find_plugin_dir(),
                                      'kdb'),
        'db': {'db_library': 'db2', 'database_name': '$tmpdir/db'}},
    'logging': {
        'admin_server': 'FILE:$tmpdir/kadmind5.log',
        'kdc': 'FILE:$tmpdir/kdc.log',
        'default': 'FILE:$tmpdir/others.log'}}


class K5Realm(object):
    """An object representing a functional krb5 test realm."""
    def __init__(self, realm='KRBTEST.COM', portbase=61000,
                 krb5_conf=None, kdc_conf=None, create_kdb=True,
                 krbtgt_keysalt=None, create_user=True, get_creds=True,
                 create_host=True, start_kdc=True, start_kadmind=False,
                 existing=None, **paths):

        if existing is not None:
            self.tmpdir = existing
            self.is_existing = True
        else:
            self.tmpdir = tempfile.mkdtemp(suffix='-krbtest')
            self.is_existing = False

        self.realm = realm
        self.portbase = portbase
        self.user_princ = 'user@' + self.realm
        self.admin_princ = 'user/admin@' + self.realm
        self.host_princ = 'host/%s@%s' % (self.hostname, self.realm)
        self.nfs_princ = 'nfs/%s@%s' % (self.hostname, self.realm)
        self.krbtgt_princ = 'krbtgt/%s@%s' % (self.realm, self.realm)
        self.keytab = os.path.join(self.tmpdir, 'keytab')
        self.client_keytab = os.path.join(self.tmpdir, 'client_keytab')
        self.ccache = os.path.join(self.tmpdir, 'ccache')
        self.kadmin_ccache = os.path.join(self.tmpdir, 'kadmin_ccache')
        self._krb5_conf = _cfg_merge(_default_krb5_conf, krb5_conf)
        self._kdc_conf = _cfg_merge(_default_kdc_conf, kdc_conf)
        self._kdc_proc = None
        self._kadmind_proc = None
        krb5_conf_path = os.path.join(self.tmpdir, 'krb5.conf')
        kdc_conf_path = os.path.join(self.tmpdir, 'kdc.conf')
        self.env = self._make_env(krb5_conf_path, kdc_conf_path)

        self._daemons = []

        self._init_paths(**paths)

        self._devnull = open(os.devnull, 'r')

        if existing is None:
            self._create_conf(self._krb5_conf, krb5_conf_path)
            self._create_conf(self._kdc_conf, kdc_conf_path)
            self._create_acl()
            self._create_dictfile()

            if create_kdb:
                self.create_kdb()
            if krbtgt_keysalt and create_kdb:
                self.run_kadminl('cpw -randkey -e %s %s' %
                                 (krbtgt_keysalt, self.krbtgt_princ))
            if create_user and create_kdb:
                self.addprinc(self.user_princ, self.password('user'))
                self.addprinc(self.admin_princ, self.password('admin'))
            if create_host and create_kdb:
                self.addprinc(self.host_princ)
                self.extract_keytab(self.host_princ, self.keytab)
            if start_kdc and create_kdb:
                self.start_kdc()
            if start_kadmind and create_kdb:
                self.start_kadmind()

        if (get_creds and ((create_kdb and create_user and start_kdc)
                           or self.is_existing)):
            self.kinit(self.user_princ, self.password('user'))
            self.klist()

    def _discover_path(self, name, default, paths):
        stderr_out = getattr(subprocess, 'DEVNULL', subprocess.PIPE)
        try:
            path = subprocess.check_output(['which', name],
                                           stderr=stderr_out).strip()
            path = path.decode(sys.getfilesystemencoding() or
                               sys.getdefaultencoding())
            _LOG.debug("Using discovered path for {name} ({path}".format(
                name=name, path=path))
            return path
        except subprocess.CalledProcessError as e:
            path = paths.get(name, default)
            _LOG.debug("Using default path for {name} ({path}): {err}".format(
                name=name, path=path, err=e))
            return path

    def _init_paths(self, **paths):
        self.kdb5_util = self._discover_path('kdb5_util',
                                             '/usr/sbin/kdb5_util', paths)
        self.krb5kdc = self._discover_path('krb5kdc',
                                           '/usr/sbin/krb5kdc', paths)
        self.kadmin_local = self._discover_path('kadmin_local',
                                                '/usr/sbin/kadmin.local',
                                                paths)
        self.kprop = self._discover_path('kprop', '/usr/sbin/kprop', paths)
        self.kadmind = self._discover_path('kadmind',
                                           '/usr/sbin/kadmind', paths)

        self._kinit = self._discover_path('kinit', '/usr/bin/kinit', paths)
        self._klist = self._discover_path('klist', '/usr/bin/klist', paths)

    def _create_conf(self, profile, filename):
        with open(filename, 'w') as conf_file:
            for section, contents in profile.items():
                conf_file.write('[%s]\n' % section)
                self._write_cfg_section(conf_file, contents, 1)

    def _write_cfg_section(self, conf_file, contents, indent_level):
        indent = '\t' * indent_level
        for name, value in contents.items():
            name = self._subst_cfg_value(name)
            if isinstance(value, dict):
                # A dictionary value yields a list subsection.
                conf_file.write('%s%s = {\n' % (indent, name))
                self._write_cfg_section(conf_file, value, indent_level + 1)
                conf_file.write('%s}\n' % indent)
            elif isinstance(value, list):
                # A list value yields multiple values for the same name.
                for item in value:
                    item = self._subst_cfg_value(item)
                    conf_file.write('%s%s = %s\n' % (indent, name, item))
            elif isinstance(value, six.string_types):
                # A string value yields a straightforward variable setting.
                value = self._subst_cfg_value(value)
                conf_file.write('%s%s = %s\n' % (indent, name, value))
            elif value is not None:
                raise TypeError("Unknown config type at key '{key}': "
                                "{type}".format(key=name, type=type(value)))

    @property
    def hostname(self):
        return socket.getfqdn()

    def _subst_cfg_value(self, value):
        template = string.Template(value)
        return template.substitute(realm=self.realm,
                                   tmpdir=self.tmpdir,
                                   hostname=self.hostname,
                                   port0=self.portbase,
                                   port1=self.portbase + 1,
                                   port2=self.portbase + 2,
                                   port3=self.portbase + 3,
                                   port4=self.portbase + 4,
                                   port5=self.portbase + 5,
                                   port6=self.portbase + 6,
                                   port7=self.portbase + 7,
                                   port8=self.portbase + 8,
                                   port9=self.portbase + 9)

    def _create_acl(self):
        filename = os.path.join(self.tmpdir, 'acl')
        with open(filename, 'w') as acl_file:
            acl_file.write('%s *\n' % self.admin_princ)
            acl_file.write('kiprop/%s@%s p\n' % (self.hostname, self.realm))

    def _create_dictfile(self):
        filename = os.path.join(self.tmpdir, 'dictfile')
        with open(filename, 'w') as dict_file:
            dict_file.write('weak_password\n')

    def _make_env(self, krb5_conf_path, kdc_conf_path):
        env = {}
        env['KRB5_CONFIG'] = krb5_conf_path
        env['KRB5_KDC_PROFILE'] = kdc_conf_path or os.devnull
        env['KRB5CCNAME'] = self.ccache
        env['KRB5_KTNAME'] = self.keytab
        env['KRB5_CLIENT_KTNAME'] = self.client_keytab
        env['KRB5RCACHEDIR'] = self.tmpdir
        env['KPROPD_PORT'] = six.text_type(self.kprop_port())
        env['KPROP_PORT'] = six.text_type(self.kprop_port())
        return env

    def run(self, args, env=None, input=None, expected_code=0):
        if env is None:
            env = self.env

        if input:
            infile = subprocess.PIPE
        else:
            infile = self._devnull

        proc = subprocess.Popen(args, stdin=infile, stdout=subprocess.PIPE,
                                stderr=subprocess.STDOUT, env=env)
        if input:
            inbytes = input.encode()
        else:
            inbytes = None
        (outdata, blank_errdata) = proc.communicate(inbytes)
        code = proc.returncode
        cmd = ' '.join(args)
        outstr = outdata.decode()
        _LOG.debug('[OUTPUT FROM `{args}`]\n{output}\n'.format(args=cmd,
                                                               output=outstr))
        if code != expected_code:
            raise Exception("Unexpected return code "
                            "for command `{args}`: {code}".format(args=cmd,
                                                                  code=code))

        return outdata

    def __del__(self):
        self._devnull.close()

    def kprop_port(self):
        return self.portbase + 3

    def server_port(self):
        return self.portbase + 5

    def create_kdb(self):
        self.run([self.kdb5_util, 'create', '-W', '-s', '-P', 'master'])

    def _start_daemon(self, args, env, sentinel):
        proc = subprocess.Popen(args, stdin=self._devnull,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.STDOUT, env=env)
        cmd = ' '.join(args)
        while True:
            line = proc.stdout.readline().decode()
            if line == "":
                code = proc.wait()
                raise Exception('`{args}` failed to start '
                                'with code {code}'.format(args=cmd,
                                                          code=code))
            else:
                _LOG.debug('[OUTPUT FROM `{args}`]\n'
                           '{output}\n'.format(args=cmd, output=line))

            if sentinel in line:
                break

        self._daemons.append(proc)

        return proc

    def start_kdc(self, args=[], env=None):
        if env is None:
            env = self.env
        assert(self._kdc_proc is None)
        self._kdc_proc = self._start_daemon([self.krb5kdc, '-n'] + args, env,
                                            'starting...')

    def _stop_daemon(self, proc):
        proc.terminate()
        proc.communicate()
        self._daemons.remove(proc)

    def stop_kdc(self):
        assert(self._kdc_proc is not None)
        self._stop_daemon(self._kdc_proc)
        self._kdc_proc = None

    def start_kadmind(self, env=None):
        if env is None:
            env = self.env
        assert(self._kadmind_proc is None)
        dump_path = os.path.join(self.tmpdir, 'dump')
        self._kadmind_proc = self._start_daemon([self.kadmind, '-nofork', '-W',
                                                 '-p', self.kdb5_util,
                                                 '-K', self.kprop,
                                                 '-F', dump_path], env,
                                                'starting...')

    def stop_kadmind(self):
        assert(self._kadmind_proc is not None)
        self.stop_daemon(self._kadmind_proc)
        self._kadmind_proc = None

    def stop(self):
        if self._kdc_proc:
            self.stop_kdc()
        if self._kadmind_proc:
            self.stop_kadmind()

        if self.tmpdir and not self.is_existing:
            shutil.rmtree(self.tmpdir)

    def addprinc(self, princname, password=None):
        if password:
            self.run_kadminl('addprinc -pw %s %s' % (password, princname))
        else:
            self.run_kadminl('addprinc -randkey %s' % princname)

    def extract_keytab(self, princname, keytab):
        self.run_kadminl('ktadd -k %s -norandkey %s' % (keytab, princname))

    def kinit(self, princname, password=None, flags=[], verbose=True,
              **keywords):
        if password:
            input = password + "\n"
        else:
            input = None

        cmd = [self._kinit]
        if verbose:
            cmd.append('-V')
        cmd.extend(flags)
        cmd.append(princname)
        return self.run(cmd, input=input, **keywords)

    def klist(self, ccache=None, **keywords):
        if ccache is None:
            ccache = self.ccache
        ccachestr = ccache
        if len(ccachestr) < 2 or ':' not in ccachestr[2:]:
            ccachestr = 'FILE:' + ccachestr
        return self.run([self._klist, ccache], **keywords)

    def klist_keytab(self, keytab=None, **keywords):
        if keytab is None:
            keytab = self.keytab
        output = self.run([self._klist, '-k', keytab], **keywords)
        return output

    def run_kadminl(self, query, env=None):
        return self.run([self.kadmin_local, '-q', query], env=env)

    def password(self, name):
        """Get a weakly random password from name, consistent across calls."""
        return name + six.text_type(os.path.basename(self.tmpdir))

    def prep_kadmin(self, princname=None, pw=None, flags=[]):
        if princname is None:
            princname = self.admin_princ
            pw = self.password('admin')
        return self.kinit(princname, pw,
                          flags=['-S', 'kadmin/admin',
                                 '-c', self.kadmin_ccache] + flags)

    def run_kadmin(self, query, **keywords):
        return self.run([self.kadmin, '-c', self.kadmin_ccache, '-q', query],
                        **keywords)

    def special_env(self, name, has_kdc_conf, krb5_conf=None, kdc_conf=None):
        krb5_conf_path = os.path.join(self.tmpdir, 'krb5.conf.%s' % name)
        krb5_conf = _cfg_merge(self._krb5_conf, krb5_conf)
        self._create_conf(krb5_conf, krb5_conf_path)
        if has_kdc_conf:
            kdc_conf_path = os.path.join(self.tmpdir, 'kdc.conf.%s' % name)
            kdc_conf = _cfg_merge(self._kdc_conf, kdc_conf)
            self._create_conf(kdc_conf, kdc_conf_path)
        else:
            kdc_conf_path = None
        return self._make_env(krb5_conf_path, kdc_conf_path)

    def kill_daemons(self):
        # clean up daemons
        for proc in self._daemons:
            os.kill(proc.pid, signal.SIGTERM)
