import glob
import logging
import os
import shutil
import sys

import boto3


def create_proxy_user(proxy_user):
    os.system(f'adduser --disabled-password --gecos "" {proxy_user}')
    os.makedirs(f"/home/{proxy_user}/.ssh", exist_ok=True)


def copy_authorized_keys_for_valid_user(proxy_user):
    valid_proxy_users = ["fivetran", "looker", "sme-rds"]
    if proxy_user in valid_proxy_users:
        shutil.copy(
            f"./{proxy_user}/configs/authorized_keys",
            f"/home/{proxy_user}/.ssh/authorized_keys",
        )

        os.chmod(f"/home/{proxy_user}/.ssh/authorized_keys", 0o600)
        for dirpath, dirnames, filenames in os.walk(f"/home/{proxy_user}"):
            shutil.chown(dirpath, user=proxy_user, group=proxy_user)
            for filename in filenames:
                shutil.chown(
                    os.path.join(dirpath, filename), user=proxy_user, group=proxy_user
                )


def remove_existing_ssh_host_keys():
    host_key_files = glob.glob("/etc/ssh/ssh_host_*key*")
    for host_key_file in host_key_files:
        os.remove(host_key_file)


def retrieve_and_set_sshd_host_keys(environment, service_name):
    key_types = ["rsa", "ecdsa", "ed25519"]
    client = boto3.client("secretsmanager")

    for key_type in key_types:
        key = client.get_secret_value(
            SecretId=f"{environment}/{service_name}/SSH_HOST_{key_type.upper()}_KEY"
        )["SecretString"]
        pub_key = client.get_secret_value(
            SecretId=f"{environment}/{service_name}/SSH_HOST_{key_type.upper()}_KEY_PUB"
        )["SecretString"]
        with open(f"/etc/ssh/ssh_host_{key_type}_key", "w") as f:
            f.write(key)
        with open(f"/etc/ssh/ssh_host_{key_type}_key.pub", "w") as f:
            f.write(pub_key)
        os.chmod(f"/etc/ssh/ssh_host_{key_type}_key", 0o600)
        os.chmod(f"/etc/ssh/ssh_host_{key_type}_key.pub", 0o644)


def update_sshd_config_with_db_host(db_host):
    if not db_host:
        logging.error(
            "PERMIT_DB_HOST environment variable is not set. SSHD PermitOpen option requires a database host."
        )
        sys.exit(1)
    with open("/etc/ssh/sshd_config", "r") as file:
        data = file.read()
    data = data.replace("db_host", db_host)
    with open("/etc/ssh/sshd_config", "w") as file:
        file.write(data)


# Entry point. Check if proxy_user is set
if __name__ == "__main__":
    proxy_user = os.getenv("PROXY_USER")
    if not proxy_user:
        logging.error(
            "PROXY_USER environment variable is not set. SSHD requires a proxy user"
        )
        sys.exit(1)

    create_proxy_user(proxy_user)
    copy_authorized_keys_for_valid_user(proxy_user)
    remove_existing_ssh_host_keys()
    environment = os.getenv("Environment")
    # log service_name
    service_name = os.getenv("SERVICE_NAME")
    logging.info(f"Service Name: {os.getenv('SERVICE_NAME')}")
    retrieve_and_set_sshd_host_keys(environment, service_name)
    db_host = os.getenv("PERMIT_DB_HOST")
    update_sshd_config_with_db_host(db_host)

    os.system("/usr/sbin/sshd -D -e")