# 
"""
SSH
---
**Summary:** configure ssh and ssh keys
    update ssh keys
    ssh_authorized_keys:
        - ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAGEA3FSyQwBI6Z+nCSjUU ...
        - ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAQEA3I7VUf2l5gSn5uavROsc5HRDpZ ...
"""
import glob
import os
import sys
from cloudinit.distros import ug_util
from cloudinit.ssh_util import *
from cloudinit import util
from cloudinit.settings import PER_ALWAYS

frequency = PER_ALWAYS

def update_authorized_keys(old_entries, keys, prev_keys, log):
    to_add = list([k for k in keys if k.valid()])
    to_sub = []
    log.debug('old entries %s, pre_keys %s , new keys %s' %(len(old_entries), len(prev_keys), len(keys)))
    for i in range(0, len(prev_keys)):
        ent = prev_keys[i]
        if not ent.valid():
            continue
        found = 0
        for k in keys:
            if k.base64 == ent.base64:
                found = 1
                break
        if not found:
            to_sub.append(ent)

    for i in range(0, len(old_entries)):
        ent = old_entries[i]
        if not ent.valid():
            continue
        # Replace those with the same base64
        for k in keys:
            if k.base64 == ent.base64:
                # Replace it with our better one
                ent = k
                # Don't add it later
                if k in to_add:
                    to_add.remove(k)
        old_entries[i] = ent

    # Now append any entries we did not match above
    for key in to_add:
        old_entries.append(key)

    for key in to_sub:
        for i in old_entries:
            if i.valid() and key.base64 == i.base64:
                log.debug("key removd %s" %str(i))
                old_entries.remove(i)

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

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

def update_ssh_keys(username, keys, cloud, log):
    (ssh_dir, pwent) = 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)
    parser = AuthKeyLineParser()
    key_entries = []
    for k in keys:
        key_entries.append(parser.parse(str(k)))

    # Extract the old and make the new
    user_pre = "%s-pre_ssh_keys" %username
    pre_ssh_fn = os.path.join(cloud.get_cpath('data'), user_pre)
    log.debug("pre ssh file %s" %pre_ssh_fn)
    if os.path.exists(pre_ssh_fn):
        prev_ssh_keys = parse_authorized_keys(pre_ssh_fn)
    else:
        prev_ssh_keys = []
    (auth_key_fn, auth_key_entries) = extract_authorized_keys(username)
    with util.SeLinuxGuard(ssh_dir, recursive=True):
        content = update_authorized_keys(auth_key_entries, key_entries, prev_ssh_keys, log)
        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)

    #update pre ssh keys
    lines = [str(b) for b in keys]
    lines.append('')
    keys_content = '\n'.join(lines)
    util.write_file(pre_ssh_fn, keys_content)

def do_remove_pre_key(cloud, log):
    user = 'root'
    user_pre = "%s-pre_ssh_keys" %user
    pre_ssh_fn = os.path.join(cloud.get_cpath('data'), user_pre)
    util.del_file(pre_ssh_fn)
    log.debug("remove user %s pre ssh keys" %user)
    sys_info = util.system_info()
    if sys_info['variant'] == "ubuntu":
        user = 'ubuntu'
        ssh_fn = "%s-pre_ssh_keys" %user
        prev_fn = os.path.join(cloud.get_cpath('data'), ssh_fn)
        util.del_file(prev_fn)
        log.debug("remove user %s pre ssh keys" %user)

def handle(_name, cfg, cloud, log, _args):
    try:
        user = util.get_user(cfg, cloud)
        ver = cloud.get_cloudinit_version()
        if not ver:
            log.debug("no cloudinit version")
            return
        pre_ver_fn = os.path.join(cloud.get_cpath('data'), 'cloudinit_version')
        if not os.path.exists(pre_ver_fn):
            log.debug("no pre cloudinit version found")
            do_remove_pre_key(cloud, log)
        else:
            pre_ver=util.load_file(pre_ver_fn)
            if pre_ver != ver:
               do_remove_pre_key(cloud, log)

        util.write_file(pre_ver_fn, ver)
        if not user:
            (users, _groups) = ug_util.normalize_users_groups(cfg, cloud.distro)
            (user, _user_config) = ug_util.extract_default(users)
            log.debug("use default user %s" %user)

        keys = cloud.get_public_ssh_keys() or []
        if "ssh_authorized_keys" in cfg:
            cfgkeys = cfg["ssh_authorized_keys"]
            keys.extend(cfgkeys)
        update_ssh_keys(user, set(keys), cloud, log)
        if user != 'root':
            update_ssh_keys('root', set(keys), cloud, log)

        if len(keys) >= 1:
            handle_ssh_pwauth(
                cfg.get('ssh_pwauth'), service_cmd=cloud.distro.init_cmd,
                service_name=cloud.distro.get_option('ssh_svcname', 'ssh'))
        else:
            handle_ssh_pwauth(
                not cfg.get('ssh_pwauth'), service_cmd=cloud.distro.init_cmd,
                service_name=cloud.distro.get_option('ssh_svcname', 'ssh'))

    except Exception:
        util.logexc(log, "Updateing ssh credentials failed!")
        raise
# vi: ts=4 expandtab
