#!/usr/bin/env python
# Copyright (c) 2014-2016 Sine Nomine Associates
#
# 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.
#
# THE SOFTWARE IS PROVIDED 'AS IS' AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

import os
import sys
import logging
import argparse
import gettext
import time

try:
    import argcomplete
except ImportError:
    argcomplete = None

import afsutil.system
import afsutil.build
import afsutil.service
import afsutil.transarc
import afsutil.rpm
from afsutil.keytab import Keytab
from afsutil.cell import Cell

logger = logging.getLogger('afsutil')

def _cd(chdir):
    """Change directory if given and returning the old directory."""
    savedir = None
    if chdir is not None:
        logger.info("Changing to directory %s", chdir)
        savedir = os.getcwd()
        os.chdir(chdir)
    return savedir

def _get_log_level(args):
    if getattr(args, 'verbose', False):
        level = logging.DEBUG
    elif getattr(args, 'quiet', False):
        level = logging.ERROR
    else:
        level = logging.INFO
    return level

def _log_options(args):
    options = {}
    options['level'] = _get_log_level(args)
    log = getattr(args, 'log', None)
    if log:
        options['filename'] = log
    fmt = getattr(args, 'log_format', None)
    if fmt is None:
        # Timestamp log records when writing to file, by default.
        if log:
           fmt = '%(asctime)s %(levelname)s %(message)s'
        else:
           fmt = '%(message)s'
    options['format'] = fmt
    return options

def _optlists2dict(options):
    """Covert the option list of lists to a dict."""
    names = {}
    if not options:
        return names
    for optlist in options:
        for o in optlist:
            name,value = o.split('=', 1)
            names[name.strip()] = value.strip()
    return names

def version(**kwargs):
    sys.stdout.write("Version %s\n" % (afsutil.__version__))

def getdeps(**kwargs):
    afsutil.build.getdeps(**kwargs)

def check(**kwargs):
    result = 0
    if not afsutil.system.check_hosts_file():
        if kwargs['fix_hosts']:
            sys.stdout.write("Attempting to fix /etc/hosts file.\n")
            afsutil.system.fix_hosts_file()
        else:
            result = 1
    if not afsutil.system.check_host_address():
        result = 2
    return result

def build(**kwargs):
    afsutil.build.build(**kwargs)

def package(**kwargs):
    afsutil.build.package(**kwargs)

def modreload(**kwargs):
    afsutil.build.modreload(**kwargs)

def install(**kwargs):
    dist = kwargs['dist']
    if dist == 'transarc':
        installer = afsutil.transarc.TransarcInstaller(**kwargs)
    elif dist == 'rpm':
        installer = afsutil.rpm.RpmInstaller(**kwargs)
    else:
        raise AssertionError("Unsupported 'dist' option")
    installer.install()

def remove(**kwargs):
    dist = kwargs['dist']
    if dist == 'transarc':
        installer = afsutil.transarc.TransarcInstaller(**kwargs)
    elif dist == 'rpm':
        installer = afsutil.rpm.RpmInstaller(**kwargs)
    else:
        raise AssertionError("Unsupported 'dist' option")
    installer.remove()

def start(**kwargs):
    afsutil.service.start(**kwargs)

def stop(**kwargs):
    afsutil.service.stop(**kwargs)

def fakekey(**kwargs):
    cell = kwargs['cell']
    realm = kwargs['realm']
    keytab = kwargs['keytab']
    enctype = kwargs['enctype']
    secret = kwargs['secret']
    if realm is None:
        realm = cell.upper()
    principal = "afs/%s@%s" % (cell, realm)
    k = Keytab()
    logger.info("Adding fake key for principal '%s', enctype='%s'.", principal, enctype)
    k.add_fake_key(principal, enctype=enctype, secret=secret)
    logger.info("Writing keytab %s.", keytab)
    k.write(keytab)

def setkey(**kwargs):
    keytab = kwargs['keytab']
    k = Keytab.load(keytab)
    k.set_service_key(**kwargs)

def newcell(**kwargs):
    akimpersonate = kwargs.get('akimpersonate', False)
    aklog = kwargs.get('aklog', None)
    if aklog:
        afsutil.cli.AKLOG = aklog
    kwargs['admins'] = [kwargs['admin']] # for Cell()
    kwargs['user'] = kwargs['admin']     # for akimpersonate

    if not afsutil.system.is_running('bosserver'):
        logger.warning("bosserver is not running! trying to start it.")
        start(components=['server'])
        time.sleep(2) # Give the server a chance to start.
    cell = Cell(**kwargs)
    cell.newcell()
    # A cache manager is required to setup the root and top level volumes!
    if kwargs['noclient']:
        logger.warning("Skipping root volume setup; --no-client was given.")
        return
    if not os.path.isfile(kwargs['keytab']):
        logger.error("Skipping root volume setup; keytab %s not found" % (kwargs['keytab']))
        return
    if afsutil.system.afs_mountpoint() is None:
        logger.warning("afs is not running! trying to start it.")
        start(components=['client'])
    k = Keytab.load(kwargs['keytab'])
    if akimpersonate:
        k.akimpersonate(**kwargs)
    else:
        _kinit_aklog(**kwargs)
    options = _optlists2dict(kwargs['options'])
    afsd_options = options.get('afsd', '')
    dynroot = '-dynroot' in afsd_options
    cell.mount_root_volumes(dynroot)
    cell.create_top_volumes(kwargs['top'])

def addfs(**kwargs):
    cell = Cell.current(**kwargs)
    cell.add_fileserver(kwargs['hostname'])

def _akimpersonate(**kwargs):
    keytab = kwargs['keytab']
    k = Keytab.load(keytab)
    k.akimpersonate(**kwargs)

def _kinit_aklog(**kwargs):
    user = kwargs['user']
    cell = kwargs['cell']
    realm = kwargs['realm']
    keytab = kwargs['keytab']
    if not keytab:
        keytab = "/tmp/afs.keytab"
    if not os.path.exists(keytab):
        raise ValueError("Keytab file not found. Please specify --keytab")
    if not realm:
        realm = cell.upper()
    output = afsutil.cli.kinit('-V', '-k', '-t', keytab, user.replace('.','/'))
    for line in output.splitlines():
        logger.info(line)
    output = afsutil.cli.aklog('-d', '-c', cell, '-k', realm)
    for line in output.splitlines():
        logger.info(line)

def login(**kwargs):
    if os.geteuid() == 0:
        logger.warning("Running afsutil login as root! Your regular user will not have a token.")
    akimpersonate = kwargs['akimpersonate']
    aklog = kwargs['aklog']
    kinit = kwargs['kinit']
    if aklog:
        afsutil.cli.AKLOG = aklog
    if kinit:
        afsutil.cli.KINIT = kinit
    if akimpersonate:
        _akimpersonate(**kwargs)
    else:
        _kinit_aklog(**kwargs)

# Customize the argument parser to display the full help instead
# of the abbreviated usage message on errors. Perhaps overkill,
# but feels friendlier since this is a custom tool.
class ArgParser(argparse.ArgumentParser):
    def __init__(self, *args, **kwargs):
        argparse.ArgumentParser.__init__(self, *args, **kwargs)
    def error(self, message):
        self.print_help(sys.stderr)
        self.exit(2, gettext.gettext('\n%s: error: %s\n') % (self.prog, message))

def main():
    nproc = afsutil.system.nproc() # default --jobs
    parser = ArgParser()
    subparsers = parser.add_subparsers(title="sub-commands", metavar="<command>")

    def _common_opts(p):
        p.add_argument('-v', '--verbose', action='store_true', help='print debug messages')
        p.add_argument('-q', '--quiet', action='store_true', help='do not print info messages')
        p.add_argument('-l', '--log', help='log file location (default: stderr)', default=None)
        p.add_argument('--log-format', help='logging format: see pydoc logging.Formatter', default=None)

    p = subparsers.add_parser('version', help='print version information')
    p.set_defaults(fn=version, must_be_root=False)

    p = subparsers.add_parser('getdeps', help='install build dependencies')
    _common_opts(p)
    p.add_argument('--creds', help='location of repo creds (maybe url)', default='/root/creds')
    p.set_defaults(fn=getdeps, must_be_root=True)

    p = subparsers.add_parser('check', help='check system')
    _common_opts(p)
    p.add_argument('--fix-hosts', help='fix /etc/hosts file by replacing loopback entry.', action='store_true')
    p.set_defaults(fn=check, must_be_root=False)

    p = subparsers.add_parser('build', help='build binaries')
    _common_opts(p)
    p.add_argument('--chdir', help='change to directory', default=None)
    p.add_argument('--cf', help='configure options, e.g., --cf="--enable-debug"', default=None)
    p.add_argument('--target', help='make target (default: all)', default='all')
    p.add_argument('--no-clean', dest='clean', help='do not run git clean', action='store_false')
    p.add_argument('--no-transarc-paths', dest='transarc_paths', help='do not use transarc paths', action='store_false')
    p.add_argument('--no-modern-kmod-name', dest='modern_kmod_name', help='do not name the kmod openafs.ko on linux',
                       action='store_false')
    p.add_argument('-j', '--jobs', help='parallel build jobs (default: %d)' % nproc, default=nproc)
    p.add_argument('--srcdir', help='source code directory (default: .)', default='.')
    p.add_argument('--tarball', help='path and file name of dest tarball (default: openafs-<sysname>.tar.gz)', default=None)
    p.set_defaults(fn=build, must_be_root=False)

    p = subparsers.add_parser('reload', help='reload the kernel module after a build')
    _common_opts(p)
    p.add_argument('--chdir', help='change to directory', default=None)
    p.set_defaults(fn=modreload, must_be_root=True)

    p = subparsers.add_parser('package', help='build packages')
    _common_opts(p)
    p.add_argument('--chdir', help='change to directory', default=None)
    p.add_argument('--source', help='build source package only', action='store_true')
    p.add_argument('--no-clean', dest='clean', help='do not run git clean', action='store_false')
    p.set_defaults(fn=package, must_be_root=False)

    p = subparsers.add_parser('install', help='install binaries')
    _common_opts(p)
    p.add_argument('--chdir', help='change to directory', default=None)
    p.add_argument('--dist', choices=['transarc', 'rpm'], help='distribution type', default='transarc')
    p.add_argument('--dir', help='distribution directory', default=None)
    p.add_argument('--components', metavar='<name>', nargs='+', choices=['client', 'server'],
                   help='components to install', default=['client', 'server']),
    p.add_argument('--cell', help='cell name (default: localcell)', default='localcell')
    p.add_argument('--hosts', nargs='+', help='cell service db hosts (default: this host)', default=None)
    p.add_argument('--realm', help='realm name (default: LOCALCELL)', default=None)
    p.add_argument('--csdb', help='path to CellServDB.dist file for client (default: none)', default=None)
    p.add_argument('--force', action='store_true', help='overwrite existing files')
    p.add_argument('-o', '--options', nargs='+', help="command line args: <name>=<value>", action='append')
    p.set_defaults(fn=install, must_be_root=True)

    p = subparsers.add_parser('remove', help="remove binaries")
    _common_opts(p)
    p.add_argument('--dist', choices=['transarc', 'rpm'], help='distribution type', default='transarc')
    p.add_argument('--components', metavar='<name>', nargs='+', choices=['client', 'server'],
                   help='components to remove', default=['client', 'server']),
    p.add_argument('--purge', action='store_true', help='remove config and data too')
    p.set_defaults(fn=remove, must_be_root=True)

    p = subparsers.add_parser('start', help="start afs services")
    _common_opts(p)
    p.add_argument('components', metavar='<name>', nargs='*',
                   help='values: server, client (default: all)',
                   default=afsutil.service.COMPONENTS)
    p.set_defaults(fn=start, must_be_root=True)

    p = subparsers.add_parser('stop', help="stop afs services")
    _common_opts(p)
    p.add_argument('components', metavar='<name>', nargs='*',
                   help='values: server, client (default: all)',
                   default=afsutil.service.COMPONENTS)
    p.set_defaults(fn=stop, must_be_root=True)

    p = subparsers.add_parser('fakekey', help='generate a keytab file for testing')
    _common_opts(p)
    p.add_argument('--keytab', help='keytab file to be created (default: /tmp/afs.keytab)', default='/tmp/afs.keytab')
    p.add_argument('--cell', help='cell name (default: localcell)', default='localcell')
    p.add_argument('--realm', help='realm name (default: LOCALCELL)', default=None)
    p.add_argument('--enctype', help='encryption type (default: aes256-cts-hmac-sha1-96)', default='aes256-cts-hmac-sha1-96')
    p.add_argument('--secret', help='passphrase (default: random)', default=None)
    p.set_defaults(fn=fakekey, must_be_root=False)

    p = subparsers.add_parser('setkey', help='add a service key from a keytab file')
    _common_opts(p)
    p.add_argument('--keytab', help="keytab file (default: /tmp/afs.keytab)", default="/tmp/afs.keytab")
    p.add_argument('--cell', help="cell name (default: detect)", default=None)
    p.add_argument('--realm', help="realm name (default: detect)", default=None)
    p.add_argument('--confdir', help="server config directory (default: /usr/afs/etc)", default="/usr/afs/etc")
    p.add_argument('--format', dest='kformat', help="key file format (default: detect)",
                   choices=['transarc', 'rxkad-k5', 'extended'], default=None)
    p.set_defaults(fn=setkey, must_be_root=True)

    p = subparsers.add_parser('newcell', help='setup a new cell')
    _common_opts(p)
    p.add_argument('--akimpersonate', help="print a ticket for admin user", action='store_true')
    p.add_argument('--keytab', help="keytab file (default: /tmp/afs.keytab)", default="/tmp/afs.keytab")
    p.add_argument('--cell', help="cell name (default: localcell)", default='localcell')
    p.add_argument('--realm', help="realm name (default: detect)", default=None)
    p.add_argument('--admin', help="admin username (default: 'admin')", default='admin')
    p.add_argument('--db', nargs='+', help="cell database hosts (default: local hostname)", default=None)
    p.add_argument('--fs', nargs='+', help="cell fileserver hosts (default: local hostname)", default=None)
    p.add_argument('--top', nargs='+', help="top level volumes (default: none)", default=[])
    p.add_argument('--no-client', dest='noclient', help='do not mount root and top level volumes', action='store_true')
    p.add_argument('--aklog', help="path to aklog program (default: detect)", default=None)
    p.add_argument('-o', '--options', nargs='+', help="command line args: <name>=<value>", action='append')
    p.set_defaults(fn=newcell, must_be_root=True)

    p = subparsers.add_parser('addfs', help='add a new fileserver to the cell')
    _common_opts(p)
    p.add_argument('hostname', help="fileserver hostname")
    p.add_argument('--keytab', help="keytab file (default: /tmp/afs.keytab)", default="/tmp/afs.keytab")
    p.add_argument('-o', '--options', nargs='+', help="command line args; <name>=<value>", action='append')
    p.set_defaults(fn=addfs, must_be_root=True)

    p = subparsers.add_parser('login', help='obtain token with a keytab')
    _common_opts(p)
    p.add_argument('--akimpersonate', help="print a ticket for the service key in the keytab", action='store_true')
    p.add_argument('--keytab', help="keytab file (default: /tmp/afs.keytab)", default="/tmp/afs.keytab")
    p.add_argument('--user', help="user name (default: admin)", default="admin")
    p.add_argument('--cell', help="cell name (default: detect with --akimpersonate)", default="localcell")
    p.add_argument('--realm', help="realm name (default: detect with --akimpersonate)", default="LOCALCELL")
    p.add_argument('--aklog', help="path to the aklog program (default: detect)", default=None)
    p.add_argument('--kinit', help="path to the kinit program (default: detect)", default=None)
    p.set_defaults(fn=login, must_be_root=False)

    if argcomplete:
        argcomplete.autocomplete(parser)
    args = parser.parse_args()
    logging.basicConfig(**_log_options(args))
    if args.must_be_root:
        if os.geteuid() != 0:
            sys.stderr.write("afsutil: Must run as root!\n")
            sys.exit(1)
    savedir = _cd(getattr(args, 'chdir', None))
    try:
        kwargs = vars(args)
        code = args.fn(**kwargs)
    except afsutil.system.CommandFailed as e:
        code = 1
        logger.exception(e)
        sys.stderr.write("Command failed: %s, code %d\n" % (e.cmd, e.code))
        sys.stderr.write("output:\n")
        sys.stderr.write("%s\n" % (e.out))
        sys.stderr.write("error:\n")
        sys.stderr.write("%s\n" % (e.err))
    except Exception as e:
        code = 1
        logger.exception(e)
        sys.stderr.write("%s\n" % (e))
    finally:
        _cd(savedir)
    sys.exit(code)

if __name__ == "__main__":
    main()

