#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# vim:fenc=utf-8

# This file is part of the  X2Go Project - https://www.x2go.org
# Copyright (C) 2012-2019 by Mike Gabriel <mike.gabriel@das-netzwerkteam.de>
#
# X2Go Session Broker is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# X2Go Session Broker is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program; if not, write to the
# Free Software Foundation, Inc.,
# 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA.

import os
import sys
import setproctitle
import argparse
import logging
import paramiko

import x2gobroker._paramiko
x2gobroker._paramiko.monkey_patch_paramiko()

from pwd import getpwnam
from grp import getgrnam

try:
    import x2gobroker.defaults
except ImportError:
    sys.path.insert(0, os.path.join(os.getcwd(), '..'))
    import x2gobroker.defaults

import x2gobroker._paramiko
x2gobroker._paramiko.monkey_patch_paramiko()

import x2gobroker.utils

supported_key_types = ('RSA', 'DSA')

PROG_NAME = os.path.basename(sys.argv[0])
PROG_OPTIONS = sys.argv[1:]
setproctitle.setproctitle("%s %s" % (PROG_NAME, " ".join(PROG_OPTIONS)))

from x2gobroker import __VERSION__
from x2gobroker import __AUTHOR__
from x2gobroker.loggers import logger_broker, logger_error

if os.geteuid() == 0:
    # propagate msgs for  the broker logger to the root logger (i.e. to stderr)
    logger_broker.propagate = 1
    logger_error.propagate = 1

# raise log level to DEBUG if requested...
if x2gobroker.defaults.X2GOBROKER_DEBUG:
    logger_broker.setLevel(logging.DEBUG)
else:
    logger_broker.setLevel(logging.INFO)

logger_broker.info('X2Go Session Broker ({version}), written by {author}'.format(version=__VERSION__, author=__AUTHOR__))
logger_broker.info('Setting up the key generator\'s environment...')
logger_broker.info('  X2GOBROKER_DEBUG: {value}'.format(value=x2gobroker.defaults.X2GOBROKER_DEBUG))
logger_broker.info('  X2GOBROKER_DAEMON_USER: {value}'.format(value=x2gobroker.defaults.X2GOBROKER_DAEMON_USER))
logger_broker.info('  X2GOBROKER_DAEMON_GROUP: {value}'.format(value=x2gobroker.defaults.X2GOBROKER_DAEMON_GROUP))

# check effective UID the broker runs as and complain appropriately...
if os.geteuid() != 0:
    logger_error.error('X2Go Session Broker\'s key generator has to run with root privileges. Exiting...')
    sys.exit(-1)

if __name__ == '__main__':

    common_options = [
        {'args':['-t','---key-type'], 'default': 'RSA', 'help': 'Choose a key type for the X2Go Session Broker pub/priv SSH key pair (available: RSA, DSA).', },
        {'args':['-f','--force'], 'default': False, 'action': 'store_true', 'help': 'Enforce the creation of a public/private key pair. WARNING: This will overwrite earlier created keys.', },
    ]
    p = argparse.ArgumentParser(description='X2Go Session Broker (Key Generator)',\
                                formatter_class=argparse.RawDescriptionHelpFormatter, \
                                add_help=True, argument_default=None)
    p_common = p.add_argument_group('common parameters')

    for (p_group, opts) in ( (p_common, common_options), ):
        for opt in opts:
            args = opt['args']
            del opt['args']
            p_group.add_argument(*args, **opt)

    cmdline_args = p.parse_args()

    if cmdline_args.key_type.upper() not in supported_key_types:
        logger_error.error('Unknown key type »{key_type}«. Possible key types are RSA and DSA. Exiting...'.format(key_type=cmdline_args.key_type.upper()))
        sys.exit(-2)

    broker_uid = x2gobroker.defaults.X2GOBROKER_DAEMON_USER
    broker_uidnumber = getpwnam(broker_uid).pw_uid
    broker_gid = x2gobroker.defaults.X2GOBROKER_DAEMON_GROUP
    broker_gidnumber = getgrnam(broker_gid).gr_gid
    broker_home = x2gobroker.defaults.X2GOBROKER_HOME

    if not os.path.exists(broker_home):
        logger_error.error('The home directory {home} of user {user} does not exists. Cannot continue. Exiting...'.format(home=broker_home, user=broker_uid))
        sys.exit(-2)

    logger_broker.info('Creating pub/priv key pair for X2Go Session Broker...')
    if not os.path.exists('{home}/.ssh'.format(home=broker_home)):
        os.mkdir('{home}/.ssh'.format(home=broker_home))
        os.chown('{home}/.ssh'.format(home=broker_home), broker_uidnumber, broker_gidnumber)
        os.chmod('{home}/.ssh'.format(home=broker_home), 0o0750)
        logger_broker.info('  Created {home}/.ssh'.format(home=broker_home))

    # generate key pair
    if cmdline_args.key_type.upper() == 'RSA':
        key = paramiko.RSAKey.generate(2048)
        id_file = 'id_rsa'
    elif cmdline_args.key_type.upper() == 'DSA':
        key = paramiko.DSSKey.generate(1024)
        id_file = 'id_dsa'

    logger_broker.info('  The {key_type} key has been generated, fingerprint: {fingerprint}'.format(key_type=cmdline_args.key_type.upper(), fingerprint=x2gobroker.utils.get_key_fingerprint_with_colons(key)))

    if os.path.exists('{home}/.ssh/{id_file}'.format(home=broker_home, id_file=id_file)) and not cmdline_args.force:
        logger_broker.error('  Private key {home}/.ssh/{id_file} exists. Use --force to overwrite'.format(home=broker_home, id_file=id_file))
        logger_broker.error('  the file with a newly generated key.')
        logger_broker.error('  Make sure that broker agents get the new public key file deployed with the')
        logger_broker.error('  x2gobroker-pubkeyauthorizer tool, if you enforce key re-generation.')
        logger_broker.error('Exiting...')
        sys.exit(-3)
    elif os.path.exists('{home}/.ssh/{id_file}'.format(home=broker_home, id_file=id_file)):
        logger_broker.warn('  WARNING: you requested to overwrite existing key files!!!')

    key.write_private_key_file('{home}/.ssh/{id_file}'.format(home=broker_home, id_file=id_file))
    os.chown('{home}/.ssh/{id_file}'.format(home=broker_home, id_file=id_file), broker_uidnumber, broker_gidnumber)
    os.chmod('{home}/.ssh/{id_file}'.format(home=broker_home, id_file=id_file), 0o0600)
    logger_broker.info('  Private key written to file {key_file}'.format(key_file='{home}/.ssh/{id_file}'.format(home=broker_home, id_file=id_file)))

    pubkey_file = open('{home}/.ssh/{id_file}.pub'.format(home=broker_home, id_file=id_file),'w')
    if id_file == 'id_rsa':
        pubkey_file.write("ssh-rsa " +key.get_base64())
    elif id_file == 'id_dsa':
        pubkey_file.write("ssh-dss " +key.get_base64())
    pubkey_file.close()
    os.chown('{home}/.ssh/{id_file}.pub'.format(home=broker_home, id_file=id_file), broker_uidnumber, broker_gidnumber)
    os.chmod('{home}/.ssh/{id_file}.pub'.format(home=broker_home, id_file=id_file), 0o0644)
    logger_broker.info('  Public key written to file {key_file}'.format(key_file='{home}/.ssh/{id_file}.pub'.format(home=broker_home, id_file=id_file)))

    logger_broker.info('Key file generation has been successful!')

