#!/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-2020 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 urllib.request
import getpass
import logging
import logging.config
import re

from pwd import getpwnam
from grp import getgrnam

__VERSION__ = '0.0.4.3'
__AUTHOR__ = 'Mike Gabriel (X2Go Project) <mike.gabriel@das-netzwerkteam.de>'

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

### The following is a code duplication of x2gobroker.loggers and x2gobroker.defaults.
### Normally, we would avoid that. However, this is to make this script independent from
### the python-x2gobroker package (and its manifold python module dependencies).

if 'X2GOBROKER_DAEMON_USER' in os.environ:
    X2GOBROKER_DAEMON_USER=os.environ['X2GOBROKER_DAEMON_USER']
else:
    X2GOBROKER_DAEMON_USER="x2gobroker"
if 'X2GOBROKER_DAEMON_GROUP' in os.environ:
    X2GOBROKER_DAEMON_GROUP=os.environ['X2GOBROKER_DAEMON_GROUP']
else:
    X2GOBROKER_DAEMON_GROUP="x2gobroker"
if 'X2GOBROKER_DEBUG' in os.environ:
    X2GOBROKER_DEBUG = ( os.environ['X2GOBROKER_DEBUG'].lower() in ('1', 'on', 'true', 'yes', ) )
else:
    X2GOBROKER_DEBUG = False
# the home directory of the user that the daemon/cgi runs as
X2GOBROKER_HOME = os.path.normpath(os.path.expanduser('~{broker_uid}'.format(broker_uid=X2GOBROKER_DAEMON_USER)))

logger_root = logging.getLogger()
stderr_handler = logging.StreamHandler(sys.stderr)
stderr_handler.setFormatter(logging.Formatter(fmt='%(asctime)s - %(name)s - %(levelname)s - %(message)s', datefmt=''))

# all loggers stream to stderr...
logger_root.addHandler(stderr_handler)

logger_broker = logging.getLogger('broker')
logger_broker.addHandler(stderr_handler)
logger_broker.propagate = 0

logger_error = logging.getLogger('error')
logger_error.addHandler(stderr_handler)
logger_error.propagate = 0

# raise log level to DEBUG if requested...
if 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 »PubKey Authorizer«\'s environment...')
logger_broker.info('  X2GOBROKER_DEBUG: {value}'.format(value=X2GOBROKER_DEBUG))
logger_broker.info('  X2GOBROKER_DAEMON_USER: {value}'.format(value=X2GOBROKER_DAEMON_USER))
logger_broker.info('  X2GOBROKER_DAEMON_GROUP: {value}'.format(value=X2GOBROKER_DAEMON_GROUP))

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

if __name__ == '__main__':

    common_options = [
        {'args':['-t','--broker-url'], 'default': None, 'help': 'The URL of the X2Go Session Broker that we want to retrieve public keys from. The common pattern for this URL is http(s)://<broker_hostname>:<port>/pubkeys/.', },
    ]
    p = argparse.ArgumentParser(description='X2Go Session Broker (PubKey Installer)',\
                                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)

    print ()
    cmdline_args = p.parse_args()

    if cmdline_args.broker_url is None:
        logger_error.error('Cannot proceed without having an URL specified. Use --broker-url as cmdline parameter.')
        logger_error.error('Exiting...')
        sys.exit(-2)

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

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

    logger_broker.info('Authorizing access to this X2Go server for X2Go Session Broker')
    logger_broker.info('at URL {url}'.format(url=cmdline_args.broker_url))

    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))

    tmpfile_name, httpmsg = urllib.request.urlretrieve(cmdline_args.broker_url)
    tmpfile = open(tmpfile_name, 'rb')
    new_pubkeys_raw = [ k for k in tmpfile.read().decode().split('\n') if k ]

    i = 0
    new_pubkeys = []
    for new_pubkey in new_pubkeys_raw:

        if not new_pubkey:
            # fully ignore empty lines
            continue

        if re.match(r'^#.*', new_pubkey):
            # fully ignore commented out lines
            continue

        # check key integrity!
        is_key = False
        if re.match(r'.*ssh-dss AAAAB3NzaC1kc3MA.*', new_pubkey):
            is_key = True
        elif re.match(r'.*ssh-rsa AAAAB3NzaC1yc2EA.*', new_pubkey):
            is_key = True

        if not is_key:
           logger_broker.error('The broker returned something that does not look like SSH RSA/DSA keys.')
           logger_broker.error('Check the URL {url}'.format(url=cmdline_args.broker_url))
           logger_broker.error('manually from a webbrowser.')
           sys.exit(-1)

        i += 1
        new_pubkeys.append(new_pubkey)

    if i == 1:
        logger_broker.info('  Found {n} public key at URL {url}'.format(n=len(new_pubkeys), url=cmdline_args.broker_url))
    elif i > 1:
        logger_broker.info('  Found {n} public keys at URL {url}'.format(n=len(new_pubkeys), url=cmdline_args.broker_url))
    else:
        logger_broker.info('  No public keys found at URL {url}'.format(url=cmdline_args.broker_url))
        sys.exit(0)

    tmpfile.close()

    append_newline = ""
    try:
        read_authorized_keys = open('{home}/.ssh/authorized_keys'.format(home=broker_home), 'rb')
        _content = read_authorized_keys.read()
        if _content and _content[-1] != 10:
            append_newline = '\n'
        already_authorized_keys = _content.decode().split('\n')
        read_authorized_keys.close()
    except IOError:
        already_authorized_keys = []

    already_authorized_keys = [ k for k in already_authorized_keys if k ]

    append_authorized_keys = open('{home}/.ssh/authorized_keys'.format(home=broker_home), 'ab')

    if append_newline:
        logger_broker.warning('  The file {authorized_keys} does not end with a newline character. Adding it.'.format(authorized_keys='{home}/.ssh/authorized_keys'.format(home=broker_home)))
        append_authorized_keys.write(append_newline)

    to_be_removed = []
    for new_pubkey in new_pubkeys:

        # legacy support for authorized_keys files containing SSH keys without options...
        # if the remote server provides an already present pubkey with options, replace the
        # non-option key in the authorized_keys file...
        keytype, pubkey, owner = new_pubkey.rsplit(" ", 2)
        keyopts = ""
        if " " in keytype:
            keyopts, keytype = keytype.rsplit(" ", 1)
        for authorized_key in already_authorized_keys:
            if authorized_key.endswith(" ".join([keytype, pubkey, owner])) and not authorized_key.startswith(keyopts):
                to_be_removed.append(authorized_key)

        if new_pubkey not in already_authorized_keys:
            append_authorized_keys.write('{k}\n'.format(k=new_pubkey).encode())
            logger_broker.info('  Adding new public key (counter={i}) to {authorized_keys}.'.format(i=i, authorized_keys='{home}/.ssh/authorized_keys'.format(home=broker_home)))
        else:
            logger_broker.warning('  Skipping new public key (counter={i}), already in {authorized_keys}.'.format(i=i, authorized_keys='{home}/.ssh/authorized_keys'.format(home=broker_home)))

    append_authorized_keys.close()

    if to_be_removed:
        cleanup_authorized_keys = open('{home}/.ssh/authorized_keys'.format(home=broker_home), 'r+')
        lines = cleanup_authorized_keys.readlines()
        cleanup_authorized_keys.seek(0)
        i = 0
        for line in lines:
            i += 1
            line = line.rstrip("\n")
            if line not in to_be_removed:
                cleanup_authorized_keys.write(line+"\n")
            else:
                logger_broker.info('  Dropping public key (counter={i}) with deprecated or no options from {authorized_keys}.'.format(i=i, authorized_keys='{home}/.ssh/authorized_keys'.format(home=broker_home)))
        cleanup_authorized_keys.truncate()
        cleanup_authorized_keys.close()

    if i == 0:
        logger_broker.error('No public SSH key was processed.')
        logger_broker.error('Check the URL {url}'.format(url=cmdline_args.broker_url))
        logger_broker.error('manually from a webbrowser.')
    else:
        # set proper file permissions
        os.chown('{home}/.ssh/authorized_keys'.format(home=broker_home), broker_uidnumber, broker_gidnumber)
        os.chmod('{home}/.ssh/authorized_keys'.format(home=broker_home), 0o0644)

        logger_broker.info('Completed successfully: X2Go Session Broker\'s PubKey Authorizer.'.format(url=cmdline_args.broker_url))
