#!/opt/dosis2/bin/python3-venv

import os
import socket
import subprocess
import sys
import pwd
import time
import xml.etree.ElementTree as ET
import psycopg2
import paramiko

DOSIS_CONFIG_PATH = '/opt/dosis2/etc/custom/dosisConfig.xml'
DPG_FUSE_KEY_DIR = '/home/dPG/.ssh/fuse/'
VALID_DB_SERVER_IP = '127.0.0.1'
AUTO_SSHFS_CONF = '/etc/auto.sshfs'
AUTO_MASTER_CONF = '/etc/auto.master'
AUTO_MASTER_OPTION = '/- /etc/auto.sshfs'
ROOT_KNOWN_HOSTS = '/root/.ssh/known_hosts'

# requires dafe_name {d_name} and dafe ip {d_ip}
RUN_IMAGES_MOUNT_OPTIONS_TEMPLATE = '/mnt/towers/{d_name}/runs -fstype=fuse,rw,allow_other,IdentityFile=/home/dPG/.ssh/fuse/id_rsa :sshfs\#dPG@{d_ip}\:/opt/dosis2/replay.run'
CANISTER_IMAGES_MOUNT_OPTIONS_TEMPLATE = '/mnt/towers/{d_name}/canisters -fstype=fuse,rw,allow_other,IdentityFile=/home/dPG/.ssh/fuse/id_rsa :sshfs\#dPG@{d_ip}\:/opt/dosis2/tmp.out'

# requires dafe_name {d_name}
RUN_IMAGES_MOUNT_TEMPLATE = '/tmp/towers/{d_name}/runs/'
CANISTER_IMAGES_MOUNT_TEMPLATE = '/tmp/towers/{d_name}/canisters/'


if not os.geteuid()==0:
    sys.exit('update handler must be run as root!')


def generate_sshfs_key(key_path):
    cmd = list(f"ssh-keygen -b 2048 -t rsa -f {key_path} -q -N".split(" "))
    cmd.append("")
    try:
        subprocess.run(cmd, check=True)
        print('sshfs key generated successfully')
        return True
    except subprocess.CalledProcessError as e:
        print(f"Error generating sshfs key: {e}")
        sys.exit()


def copy_sshfs_key(host, user, password, sshfs_key):
    try:
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(hostname=host, username=user, password=password)
        ssh.exec_command('mkdir -p /home/dPG/.ssh && touch /home/dPG/.ssh/authorized_keys')
        with open(sshfs_key, 'r') as f:
            sshfs_content = f.read().strip()
        with ssh.open_sftp().open('/home/dPG/.ssh/authorized_keys', 'r') as t:
            existing_keys = t.read().decode().splitlines()
        if sshfs_content not in existing_keys:
            with ssh.open_sftp().open('/home/dPG/.ssh/authorized_keys', 'a') as t:
                t.write(sshfs_content + '\n')
        print('credentials copied successfully')
    except paramiko.AuthenticationException:
        print("Authentication failed, please verify credentials")
    except paramiko.SSHException as sshException:
        print("Unable to establish SSH connection %s" % sshException)
    except Exception as e:
        print("Credential copy error: %s" % e)
    finally:
        ssh.close()


def get_active_towers(db_name, db_user, db_host):
    try:
        tuple_keys = ("tag", "site_ip", "active")
        with psycopg2.connect(dbname=db_name, user=db_user, host=db_host) as conn:
            with conn.cursor() as cur:
                cur.execute("""SELECT tag, site_ip, active FROM equipment.dafe""")
                rows = cur.fetchall()
                active_towers = [dict(zip(tuple_keys, row)) for row in rows]
        return [tower for tower in active_towers if tower.get('tag') != 'SYSTEM' and tower.get('active')]
    except (psycopg2.Error, Exception) as e:
        print(f"An error occurred: {e}")
        return None


def get_general_property_from_dosis_config(general_property):
    try:
        xml_path = './section[@name="general"]/property[@name="%s"]/value' % general_property
        dosis_config_tree = ET.parse(DOSIS_CONFIG_PATH)
        return dosis_config_tree.findtext(xml_path)
    except (FileNotFoundError, ET.ParseError):
        return None


def do_first_login(host, user, password_file, known_hosts_file, timeout=60):
    try:
        client = paramiko.SSHClient()
        client.load_system_host_keys(known_hosts_file)
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(
            hostname=host,
            username=user,
            key_filename=password_file,
            timeout=timeout
        )

        host_key = client.get_transport().get_remote_server_key()
        host_key_line = f'{host} {host_key.get_name()} {host_key.get_base64()}'
        with open(known_hosts_file, 'a+') as f:
            f.seek(0)
            existing_host_keys = f.read().splitlines()
            if host_key_line not in existing_host_keys:
                f.write(host_key_line + '\n')

        stdin, stdout, stderr = client.exec_command(command='whoami', timeout=timeout)
        result = stdout.channel.recv_exit_status()
        time.sleep(3)
        client.close()
        return result
    except paramiko.BadHostKeyException as e:
        print(e)
    except paramiko.AuthenticationException as e:
        print(e)
    except paramiko.SSHException as err:
        print(err)


def ensure_auto_master_option():
    try:
        with open(AUTO_MASTER_CONF, 'r+') as auto_master:
            contents = auto_master.read()
            if AUTO_MASTER_OPTION in contents.splitlines():
                return
            if contents and not contents.endswith('\n'):
                auto_master.write('\n')
            auto_master.write(AUTO_MASTER_OPTION + '\n')
    except FileNotFoundError:
        with open(AUTO_MASTER_CONF, 'w') as auto_master:
            auto_master.write(AUTO_MASTER_OPTION + '\n')



def create_file_shares(active_towers_list):
    # add sshfs rules to auto.master
    ensure_auto_master_option()

    # handle the rest of the fileshare setup
    try:
        os.makedirs('/root/.ssh/', exist_ok=True)
        os.makedirs(os.path.dirname(ROOT_KNOWN_HOSTS), exist_ok=True)
        open(ROOT_KNOWN_HOSTS, 'a').close()

        with open(AUTO_SSHFS_CONF, 'w') as asc:
            for at in active_towers_list:
                d_name = at.get('tag').split('_')[1]
                # Create mount directories
                run_mnt_path = RUN_IMAGES_MOUNT_TEMPLATE.format(d_name=d_name)
                os.makedirs(run_mnt_path, exist_ok=True)
                chown(run_mnt_path, 'dPG')
                canister_mnt_path = CANISTER_IMAGES_MOUNT_TEMPLATE.format(d_name=d_name)
                os.makedirs(canister_mnt_path, exist_ok=True)
                chown(canister_mnt_path, 'dPG')

                # write options to auto.sshfs
                asc.write(RUN_IMAGES_MOUNT_OPTIONS_TEMPLATE.format(d_name=d_name, d_ip=at.get('site_ip'))+'\n')
                asc.write(CANISTER_IMAGES_MOUNT_OPTIONS_TEMPLATE.format(d_name=d_name, d_ip=at.get('site_ip'))+'\n')

                # copy ssh key to target tower
                public_key = os.path.join(DPG_FUSE_KEY_DIR, 'id_rsa.pub')
                print(public_key)
                copy_sshfs_key(at.get('site_ip'), 'dPG', '1$Dosis', public_key)

                # login to the towers at least once to validate key copy.
                private_key = os.path.join(DPG_FUSE_KEY_DIR, 'id_rsa')
                if do_first_login(at.get('site_ip'), 'dPG', private_key, ROOT_KNOWN_HOSTS) == 0:
                    print(f"successfully logged into {at.get('tag')}")
                else:
                    # if login fails, you may need to clear out key from known_hosts:
                    # ssh-keygen -f "/root/.ssh/known_hosts" -R "<site_ip>"
                    print(f"failed to log into {at.get('tag')}")
    except (IOError, Exception) as e:
        print(f"An error occurred: {e}")


def chown(target, username):
    pwnam = pwd.getpwnam(username)
    os.chown(target, pwnam.pw_uid, pwnam.pw_gid)


if __name__ == '__main__':
    # check if it's a site-server and get active towers.
    local_db_ip = get_general_property_from_dosis_config('DBServer')
    if VALID_DB_SERVER_IP == local_db_ip and 'ss' in socket.gethostname().lower():
        active_towers = get_active_towers('dosis2', 'dpg', local_db_ip)
        if active_towers is None:
            raise ValueError('Unable to load active towers from equipment.dafe')
    else:
        raise ValueError('This should only be run on a site-server')

    # check for pre-requisite software
    if not os.path.exists('/usr/bin/sshfs'):
        raise ValueError('Package \'sshfs\' not found')

    # generate a new ssh key for the sshfs mount
    key_path = os.path.join(DPG_FUSE_KEY_DIR, 'id_rsa')
    if not os.path.exists(key_path):
        os.makedirs(DPG_FUSE_KEY_DIR, exist_ok=True)
        generate_sshfs_key(key_path)
        chown(DPG_FUSE_KEY_DIR, 'dPG')
        chown(key_path, 'dPG')
        chown(key_path + '.pub', 'dPG')
    else:
        print('pre-existing sshfs key found')

    # create shares
    create_file_shares(active_towers)
