#! /usr/bin/python3
# -*- coding: utf-8 -*-
#
# SSH Authorized keys updater
#
# This script act as a client for SSH Authorized keys API:
#
#   https://gitlab.easter-eggs.com/brenard/ssh-auth-keys-api
#
# Author: Benjamin Renard <brenard@easter-eggs.com>
# Date: Mon, 10 Dec 2012 18:04:24 +0100
# Source: http://gitlab.easter-eggs.com/brenard/ssh-auth-keys-updater
# License:
#
#   You can use it under GNU GENERAL PUBLIC LICENSE Version 3 policy
#
#   For more information on GNU GENERAL PUBLIC LICENSE Version 3 policy,
#   please refer to COPYING.
#

""" SSH Authorized keys updater """

import sys
import os
import pwd
import logging
import json
import re
import urllib.request
import difflib
import configparser
import argparse
import subprocess

default_config_file = "/etc/ssh-auth-keys-updater.conf"

######
# DO #
######

parser = argparse.ArgumentParser(description=__doc__)

parser.add_argument(
    '-j', '--just-try',
    action="store_true",
    dest="justtry",
    help="Enable just-try mode")

parser.add_argument(
    '-v', '--verbose',
    action="store_true",
    dest="verbose",
    help="Enable verbose mode")

parser.add_argument(
    '-d', '--debug',
    action="store_true",
    dest="debug",
    help="Enable debug mode")

parser.add_argument(
    '-l', '--log-file',
    action="store",
    type=str,
    dest="logfile",
    help="Log file path")

parser.add_argument(
    '-C', '--console',
    action="store_true",
    dest="console",
    help="Log on console even if log file is defined in " +
    "configuration file")

parser.add_argument(
    '-U', '--api-url',
    action="store",
    type=str,
    dest="apiurl",
    help="API URL")

parser.add_argument(
    '-c', '--config',
    action="store",
    type=str,
    dest="config",
    help="Configuration file path (Default: %s)"
         % default_config_file,
    default=default_config_file)

options = parser.parse_args()

config = configparser.ConfigParser()
if options.config:
    config.read(options.config)

logformat = ('%(asctime)s - ssh-auth-keys-updater'
             '- %(levelname)s - %(message)s')
if options.debug:
    loglevel = logging.DEBUG
elif options.verbose:
    loglevel = logging.INFO
else:
    loglevels = {
        'CRITICAL': logging.CRITICAL,
        'ERROR': logging.ERROR,
        'WARNING': logging.WARNING,
        'INFO': logging.INFO,
        'DEBUG': logging.DEBUG,
    }
    if config.get('log', 'level', fallback='WARNING') in loglevels:
        loglevel = loglevels.get(config.get('log', 'level',
                                            fallback='WARNING'))
    else:
        print("Invalid config log level. Set WARNING level.")
        loglevel = logging.WARNING

if options.logfile:
    logging.basicConfig(filename=options.logfile,
                        level=loglevel, format=logformat)
elif config.get('log', 'file', fallback=None) and not options.console:
    logging.basicConfig(filename=config.get('log', 'file'),
                        level=loglevel, format=logformat)
else:
    logging.basicConfig(level=loglevel, format=logformat)

# Check configuration

# Load allowed users
if config.get('ssh', 'allowed_users', fallback=None):
    allowed_users = {}
    # Check that allowed users really exists on POSIX passwd database
    for user in config.get('ssh', 'allowed_users').split():
        try:
            allowed_users[user] = pwd.getpwnam(user)
        except KeyError:
            logging.warning(
                "User %s from configuration file (allowed_users) "
                "doesn't exists in passwd POSIX database.", user)
            if user in allowed_users:
                del allowed_users[user]
    if not allowed_users:
        logging.info('No allowed user loaded from configuration file. Stop.')
        sys.exit(0)
    logging.debug('Allowed user(s) loaded from configuration file: %s',
                  ', '.join(allowed_users))
else:
    logging.critical(
        'You must provide SSH allowed users using configuration file.')
    sys.exit(1)

# Load/detect authorized_keys file path
auth_keys_path = config.get('ssh', 'authorized_keys_file_path', fallback=None)
if auth_keys_path:
    if not re.search('{user}', auth_keys_path) and \
         not re.search('{home}', auth_keys_path):
        logging.critical(
            'Invalid authorized_keys_file_path configuration parameter (%s):'
            ' must contain keyword "{user}" or "{home}".', auth_keys_path)
        sys.exit(1)
else:
    # Try to detect authorized_keys file path from sshd_config file
    sshd_config = config.get('ssh', 'sshd_config',
                             fallback='/etc/ssh/sshd_config')
    if not os.path.isfile(sshd_config):
        logging.critical(
            'SSH configuration file not found (%s). Fail to auto-detect '
            'authorized_keys file path.', sshd_config)
        sys.exit(1)

    # Dump sshd config
    try:
        result = subprocess.run(
            ['sshd', '-f', sshd_config, '-T'],
            capture_output=True, check=True)
    except subprocess.CalledProcessError:
        logging.critical(
            "Fail to dump SSH configuration: can't detect AuthorizedKeysFile path",
            exc_info=True)
        sys.exit(1)

    # Detect authorizedkeysfile value
    authorizedkeysfile_line = re.compile(r'^authorizedkeysfile (.*)\s*$')
    sshd_auth_keys_path = False
    for line in result.stdout.splitlines():
        m = authorizedkeysfile_line.match(line.decode('utf8', 'ignore'))
        if m:
            sshd_auth_keys_path = m.group(1)
            break

    if not sshd_auth_keys_path:
        logging.critical(
            'Fail to auto-detect authorized_keys file path from SSH '
            'configuration file (%s): no AuthorizedKeysFile directive found.',
            sshd_config)
        sys.exit(1)
    if '%u' not in sshd_auth_keys_path and '%h' not in sshd_auth_keys_path:
        logging.critical(
            'Invalid AuthorizedKeysFile value retreived from SSH configuration'
            'file (%s): must contain keyword "%%u" or "%%h".', sshd_auth_keys_path)
        sys.exit(1)
    logging.info('SSH AuthorizedKeys file path detected: %s',
                 sshd_auth_keys_path)
    auth_keys_path = sshd_auth_keys_path.replace('%u', '{user}') \
                                        .replace('%h', '{home}')

# Load API URL
if options.apiurl:
    api_url = options.apiurl
elif config.get('api', 'url', fallback=None):
    api_url = config.get('api', 'url')
else:
    logging.critical(
        'You must provide API URL using configuration file '
        '(or -c parameter).')
    sys.exit(1)

# Load API Auth token
if config.get('api', 'auth_token', fallback=None):
    api_auth_token = config.get('api', 'auth_token')
else:
    logging.critical(
        'You must provide API Authorization token using '
        'configuration file.')
    sys.exit(1)

# Retreive API data
try:
    req = urllib.request.Request(api_url, {}, {'X-AUTHORIZATION': api_auth_token})
    resp = urllib.request.urlopen(req)
    txt_data = resp.read()
except Exception:
    logging.fatal(
        'Fail to retreive data from API (URL: %s)',
        api_url, exc_info=True)
    sys.exit(1)

# Parse API data
try:
    data = json.loads(txt_data)
    logging.debug('API decoded return: %s', data)
except Exception:
    logging.debug('API return: %s', txt_data)
    logging.fatal('Fail to parse API return', exc_info=True)
    sys.exit(1)

# Check API return data content
if not isinstance(data, dict) or 'access' not in data or not isinstance(data['access'], dict):
    logging.critical('Invalid API return')
    sys.exit(1)

# Define default auth keys file/dir owner, group and mode
if re.search('{home}', auth_keys_path):
    logging.debug('Auth key file is inside user home directory')
    default_auth_keys_file_owner = '{user_uid}'
    default_auth_keys_file_group = '{user_gid}'
    default_auth_keys_file_mode = '0640'
    default_auth_keys_dir_mode = '0750'
else:
    logging.debug('Auth key file is outside user home directory')
    default_auth_keys_file_owner = '0'
    default_auth_keys_file_group = '0'
    default_auth_keys_file_mode = '0644'
    default_auth_keys_dir_mode = '0755'


def get_auth_keys_file_path(user):
    """ Compute auth key file path """
    global auth_keys_path, allowed_users
    return auth_keys_path.strip() \
                         .format(user=user,
                                 home=allowed_users[user].pw_dir)


# Handle user auth keys file update
for user in data.get('access'):
    if user not in allowed_users:
        logging.warning('User "%s" from API data not allowed', user)
        continue

    if not isinstance(data.get('access').get(user), dict):
        logging.warning('Invalid API access data for user %s', user)
        continue

    auth_keys_file = get_auth_keys_file_path(user)

    auth_keys_content = [
        "# FILE MANAGE BY SSH-AUTH-KEYS-UPDATER. DO NOT EDIT IT MANUALLY, " +
        " ALL CHANGES WILL BE OVERRIDEN !\n",
        "\n"
    ]
    for ruser in data.get('access').get(user):
        if not isinstance(data.get('access').get(user).get(ruser), list):
            logging.warning(
                'Invalid API access data for user %s and remote user %s',
                user, ruser)
            continue
        auth_keys_content.append("# User %s\n" % ruser)
        for key in data.get('access').get(user).get(ruser):
            auth_keys_content.append(key+"\n")
        auth_keys_content.append("\n")

    logging.debug(
        'Auth keys file %s content:\n%s', auth_keys_file,
        "".join(auth_keys_content))

    # Check auth keys file directory exists and create it if not (paying
    # attention about directory rights)
    auth_keys_dir = os.path.dirname(auth_keys_file)
    if not os.path.isdir(auth_keys_dir):
        mode = config.get('ssh', 'authorized_keys_dir_mode',
                          fallback=default_auth_keys_dir_mode)
        if not re.match('^[0-2][0-7]{3}$', mode):
            mode = '0700'
            logging.warning(
                'Invalid authorized_keys_dir_mode parameter retreived '
                'from configuration file. Use %s as default.', mode)
        try:
            logging.info(
                "Create auth keys directory %s (mode: %s)",
                auth_keys_dir, mode)
            if not options.justtry:
                os.mkdir(auth_keys_dir, int(mode, 8))
            else:
                logging.debug("Just-try mode, not really created it")
        except Exception:
            logging.error(
                "Error creating auth keys directory %s",
                auth_keys_dir, exc_info=True)
            continue
        logging.info(
            "Auth keys file %s doesn't exists. Create it with content:\n%s",
            auth_keys_file, "".join(auth_keys_content))
    elif os.path.isfile(auth_keys_file):
        # Check if update is need
        diff = list(difflib.unified_diff(
                    open(auth_keys_file, 'r').readlines(), auth_keys_content,
                    'current', 'new'))
        if diff:
            logging.info(
                "Some changes have to be put on auth keys file %s:\n%s",
                auth_keys_file, "".join(diff))
        else:
            logging.debug("Auth keys file %s is uptodate", auth_keys_file)
            continue
    else:
        logging.info(
            "Auth keys file %s doesn't exists. Create it with content:\n%s",
            auth_keys_file, "".join(auth_keys_content))

    # Create/update auth keys file
    try:
        if not options.justtry:
            fd = open(auth_keys_file, 'w')
            fd.writelines(auth_keys_content)
            fd.close()
        else:
            logging.debug("Just-try mode, not really created/updated it.")
        logging.info('Auth keys file %s updated', auth_keys_file)

        # Set auth key file owner/group
        logging.debug('Set owner/group and mode of auth keys file')
        uid = config.get('ssh', 'authorized_keys_file_owner',
                         fallback=default_auth_keys_file_owner) \
                    .format(user_uid=allowed_users[user].pw_uid,
                            user_gid=allowed_users[user].pw_gid) \
                    .strip()
        gid = config.get('ssh', 'authorized_keys_file_group',
                         fallback=default_auth_keys_file_group) \
                    .format(user_uid=allowed_users[user].pw_uid,
                            user_gid=allowed_users[user].pw_gid) \
                    .strip()
        logging.debug('Convert owner uid: %s', uid)
        logging.debug('Convert group gid: %s', gid)

        if not re.match('^[0-9]+$', uid):
            logging.warning(
                "Invalid authorized_keys_file_owner parameter retreived "
                "from configuration file. Can't set owner and group of "
                "auth key file.")
        elif not re.match('^[0-9]+$', gid):
            logging.warning(
                "Invalid authorized_keys_file_group parameter retreived "
                "from configuration file. Can't set owner and group of "
                "auth key file.")
        else:
            logging.debug(
                "Set auth key file %s owner and group to %s:%s",
                auth_keys_file, uid, gid)
            if not options.justtry:
                os.chown(auth_keys_file, int(uid), int(gid))
            else:
                logging.debug(
                    "Just-try mode, not really setting owner and group.")

        # Set auth key file mode
        mode = config.get('ssh', 'authorized_keys_file_mode',
                          fallback=default_auth_keys_file_mode)
        if not re.match('^[0-2][0-7]{3}$', mode):
            mode = default_auth_keys_file_mode
            logging.warning(
                "Invalid authorized_keys_file_mode parameter retreived "
                "from configuration file. Use %s as default.", mode)
        logging.debug(
            "Set auth key file %s mode to %s", auth_keys_file, mode)
        if not options.justtry:
            os.chmod(auth_keys_file, int(mode, 8))
        else:
            logging.debug("Just-try mode, not really setting mode.")
    except Exception:
        logging.error(
            "Error occured writing auth key file %s content",
            auth_keys_file, exc_info=True)

for user in allowed_users:
    if user not in data.get('access'):
        logging.warning("No SSH key retreived for user %s", user)

        auth_keys_file = get_auth_keys_file_path(user)

        if os.path.isfile(auth_keys_file):
            logging.debug("Delete auth key file %s", auth_keys_file)
            if not options.justtry:
                try:
                    os.remove(auth_keys_file)
                except Exception:
                    logging.error(
                        "Error occured removing auth key file %s",
                        auth_keys_file, exc_info=True)
            else:
                logging.debug(
                    "Just-try mode, not really deleting auth key file.")
        else:
            logging.debug(
                "Auth key file not found for this user (%s).",
                auth_keys_file)
