import os
import yaml

HOST_ID_CONF = "host_id.conf"
USER_CA_CONF = "user_id.conf"
USER_CA_KEYS = "user-ca-keys"
AUTH_KEYS_CONF = "authorized_keys.conf"

ETC = "/etc/ssh"
CONFIG = "sshd_config"
CONFIG_D = "sshd_config.d"
LOGFILE = "/tmp/vmadm.script"
USER_DATA = "/var/lib/cloud/instance/user-data.txt"


def etc_join(*paths):
    return os.path.join(etc, *paths)


def log(msg):
    logfile.write(msg)
    logfile.write("\n")
    logfile.flush()


logfile = open(LOGFILE, "w")
log("vmadm cloud-init script starting")

if os.environ.get("VMADM_TESTING"):
    filename = "smoke/user-data"
    etc = "x"
else:
    filename = USER_DATA
    etc = ETC

key_types = ("rsa", "dsa", "ecdsa", "ed25519")

log(f"loading user-data from {filename}")
obj = yaml.safe_load(open(filename))

ssh_keys = obj.get("ssh_keys", {})
user_ca_pubkey = obj.get("user_ca_pubkey", {})
allow_authorized_keys = obj.get("allow_authorized_keys", True)

keys = []
certs = []

for key_type in key_types:
    filename = etc_join(f"ssh_host_{key_type}_key.pub")
    if os.path.exists(filename):
        log(f"removing {filename}")
        os.remove(filename)
    else:
        log(f"file {filename} does not exist")

for key_type in key_types:
    key = ssh_keys.get(f"{key_type}_private")
    cert = ssh_keys.get(f"{key_type}_certificate")
    log(f"key {key_type} {key}")
    log(f"cert {key_type} {cert }")

    if key:
        filename = etc_join(f"ssh_host_{key_type}_key")
        log(f"writing key {filename}")
        keys.append(filename)
        with open(filename, "w") as f:
            f.write(key)

    if cert:
        filename = etc_join(f"ssh_host_{key_type}_key-cert.pub")
        log(f"writing cert {filename}")
        certs.append(filename)
        with open(filename, "w") as f:
            f.write(cert)

user_ca_filename = etc_join(USER_CA_KEYS)
if user_ca_pubkey:
    with open(user_ca_filename, "w") as f:
        f.write(user_ca_pubkey)

config = etc_join(CONFIG)
data = ""
if os.path.exists(config):
    data = open(config).read()

with open(config, "w") as f:
    f.write(data)

log(f"configuring sshd {config}")
log(f"keys {keys}")
log(f"certs {certs}")

config_d = etc_join(CONFIG_D)
log(f"config.d {CONFIG_D}")
if not os.path.exists(config_d):
    log(f"mkdir {config_d}")
    os.mkdir(config_d)

host_id_conf = etc_join(CONFIG_D, HOST_ID_CONF)
log(f"write {host_id_conf}")
with open(host_id_conf, "w") as f:
    for filename in keys:
        log(f"hostkey {filename}")
        f.write(f"hostkey {filename}\n")
    for filename in certs:
        log(f"hostcert {filename}")
        f.write(f"hostcertificate {filename}\n")

if user_ca_pubkey:
    user_ca_conf = etc_join(CONFIG_D, USER_CA_CONF)
    log(f"write {user_ca_conf}")
    with open(user_ca_conf, "w") as f:
        log(f"trustedusercakeys {user_ca_filename}")
        f.write(f"trustedusercakeys {user_ca_filename}\n")

if not allow_authorized_keys:
    authz_keys_conf = etc_join(CONFIG_D, AUTHZ_KEYS_CONF)
    log(f"write {authz_keys_conf}")
    with open(auth_keys_conf, "w") as f:
        f.write("authorizedkeysfile none\n")

log("vmadm cloud-init script ending")
logfile.close()