# Copyright (C) 2021-2030 Baidu Ltd.
#
# Author: dudeli<dudeli@baidu.com>
#
# This file is part of cloud-init. See LICENSE file for license information.

"""
Update authorized_keys file
-------------
**Summary:** Update authorized_keys

Update authorized_keys for user if necessary.
This a function for baidu cloud custom.

**Internal name:** ``cc_update_authorized_keys``

**Module frequency:** per always

**Supported distros:** all

**Config keys**::

    add_authrized_keys:
            - key1
            - key2
    del_authrized_keys:
            - key3
            - key4
"""

import os

from cloudinit import log as logging
from cloudinit import util
from cloudinit import ssh_util
from cloudinit.settings import PER_ALWAYS

frequency = PER_ALWAYS

LOG = logging.getLogger(__name__)


def handle(_name, cfg, cloud, log, args):
    """
    handle reset password module

    Args:
    _name: module name
    cfg: cloud config
    cloud: cloudinit object wrapper
    log: logger
    args: extra args

    Returns:
    """
    if 'add_authorized_keys' not in cfg and 'del_authorized_keys' not in cfg:
        log.debug(("Skipping module named %s,"
                   " no authorized keys cfg in configuration"), _name)
        return

    # support only root for now
    user_name = 'root'

    # delete keys
    if 'del_authorized_keys' in cfg:
        delcfg = cfg['del_authorized_keys']
        (ssh_dir, pwent) = ssh_util.users_ssh_info(user_name)
        (key_fn, key_entries) = ssh_util.extract_authorized_keys(user_name)
        parser = ssh_util.AuthKeyLineParser()
        if os.path.isdir(ssh_dir):
            dkeys = []
            dkeys.extend(delcfg)
            for dk in dkeys:
                dk_entry = parser.parse(str(dk), options=None)
                for k in key_entries:
                    if k.base64 == dk_entry.base64:
                        key_entries.remove(k)

            with util.SeLinuxGuard(ssh_dir, recursive=True):
                content = ssh_util.update_authorized_keys(key_entries,
                                                          key_entries)
                util.ensure_dir(os.path.dirname(key_fn), mode=0o700)
                util.write_file(key_fn, content, mode=0o600)
                util.chownbyid(key_fn, pwent.pw_uid, pwent.pw_gid)

    # add keys
    if 'add_authorized_keys' in cfg:
        addcfg = cfg['add_authorized_keys']
        akeys = []
        akeys.extend(addcfg)
        ssh_util.setup_user_keys(akeys, user_name)

# vi: ts=4 expandtab