# Author: Tatiana Kholkina <holkina@selectel.ru>

"""
Custom SSH Authorized Keys
--------------------------
**Summary:** ssh authorized keys management

This module allows to add and delete SSH authorized keys.

Authorized keys for the default user/first user defined in ``users`` can be
specified using ``custom_ssh_authorized_keys``. Keys should be specified as a
list of public keys. The list of keys will be updated after an instance reboot.
To delete the key just remove it from the list and reboot an instance.
If ``custom_ssh_authorized_keys`` is empty all keys added via this config key
will be removed. If ``custom_ssh_authorized_keys`` is not presented in config
nothing will be changed.

Root login can be enabled/disabled using the ``disable_root`` config key. Root
login options can be manually specified with ``disable_root_opts``. If
``disable_root_opts`` is specified and contains the string ``$USER``,
it will be replaced with the username of the default user. By default,
root login is disabled, and root login opts are set to::

    no-port-forwarding,no-agent-forwarding,no-X11-forwarding

.. note::
    keys will be marked with a prefix to correctly determine them

.. note::
    if one of provided keys is invalid no changes will be applied

**Internal name:** ``cc_custom_ssh_authorized_keys``

**Module frequency:** always

**Supported distros:** all

**Config keys**::

    disable_root: <true/false>
    disable_root_opts: <disable root options string>
    custom_ssh_authorized_keys:
        - ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAGEA3FSyQwBI6Z+nCSjUU ...
        - ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAQEA3I7VUf2l5gSn5uavROsc5HRDpZ ...
"""

import os

from cloudinit.distros import ug_util
from cloudinit.log import logging
from cloudinit.settings import PER_ALWAYS
from cloudinit import ssh_util
from cloudinit import util

try:
    from cloudinit.config.cc_ssh import DISABLE_ROOT_OPTS as DISABLE_USER_OPTS
except ImportError:
    from cloudinit.ssh_util import DISABLE_USER_OPTS

LOG = logging.getLogger(__name__)

frequency = PER_ALWAYS

KEY_PREFIX = 'cloud-init'


def handle(_name, cfg, cloud, log, _args):
    if "custom_ssh_authorized_keys" in cfg:
        # Get user to add authorized_keys to (the same as in cc_ssh)
        (users, _groups) = ug_util.normalize_users_groups(cfg, cloud.distro)
        (user, _user_config) = ug_util.extract_default(users)
        disable_root = util.get_cfg_option_bool(cfg, "disable_root", True)
        disable_root_opts = util.get_cfg_option_str(cfg, "disable_root_opts",
                                                    DISABLE_USER_OPTS)

        content = cfg["custom_ssh_authorized_keys"] or []

        if not user:
            user = 'root'
        try:
            apply_credentials(content, user, disable_root, disable_root_opts)
        except Exception:
            util.logexc(log, "Applying ssh credentials failed!")


def apply_credentials(new_keys, user, disable_root, disable_root_opts):
    if user:
        setup_user_keys(new_keys, user)
    else:
        if disable_root:
            if not user:
                user = "NONE"
            key_prefix = disable_root_opts.replace('$USER', user)
        else:
            key_prefix = ''

        setup_user_keys(new_keys, 'root', options=key_prefix)


def setup_user_keys(new_keys, username, options=None):
    # Make sure the users .ssh dir is setup accordingly
    (ssh_dir, pwent) = ssh_util.users_ssh_info(username)
    if not os.path.isdir(ssh_dir):
        util.ensure_dir(ssh_dir, mode=0o700)
        util.chownbyid(ssh_dir, pwent.pw_uid, pwent.pw_gid)

    # Turn the 'update' keys given into actual entries.
    parser = ssh_util.AuthKeyLineParser()
    key_entries = []
    for k in new_keys:
        key_entries.append(parser.parse(str(k), options=options))

    if not validate_and_add_prefix(key_entries):
        LOG.warning("Keys added via custom_ssh_authorized_keys "
                    "are invalid. Changes are not applied.")
        return

    # Extract existing keys from .ssh/authorized_keys
    (auth_key_fn, auth_key_entries) = ssh_util.extract_authorized_keys(username)
    with util.SeLinuxGuard(ssh_dir, recursive=True):
        content = update_authorized_keys(auth_key_entries, key_entries)
        util.ensure_dir(os.path.dirname(auth_key_fn), mode=0o700)
        util.write_file(auth_key_fn, content, mode=0o600)
        util.chownbyid(auth_key_fn, pwent.pw_uid, pwent.pw_gid)


def update_authorized_keys(existing_keys, new_keys):
    # Remove keys with prefix
    keys_to_add = list(filter(lambda k: not k.comment.startswith(KEY_PREFIX),
                              existing_keys))
    for key in new_keys:
        keys_to_add.append(key)

    # Now format them back to strings...
    lines = [str(b) for b in keys_to_add]

    # Ensure it ends with a newline
    lines.append('')
    return '\n'.join(lines)


def validate_and_add_prefix(keys):
    for key in keys:
        if not key.valid():
            return False
        key.comment = '%(prefix)s-%(comment)s' % {
            'prefix': KEY_PREFIX,
            'comment': key.comment
        }
    return True