#!/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'

# 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')
        with open(sshfs_key, 'r') as f:
            sshfs_content = f.read()
        with ssh.open_sftp().open('/home/dPG/.ssh/authorized_keys', 'a') as t:
            t.write(sshfs_content)
        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()
        with open(known_hosts_file, 'a') as f:
            f.write(f'{host} {host_key.get_name()} {host_key.get_base64()}\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(e)



def create_file_shares(active_towers_list):
    # add sshfs rules to auto.master
    try:
        with open(AUTO_MASTER_CONF, 'r') as read_amc:
            if AUTO_MASTER_OPTION not in read_amc.read():
                with open(AUTO_MASTER_CONF, 'a') as append_amc:
                    append_amc.write(AUTO_MASTER_OPTION+'\n')
    except FileNotFoundError:
        with open(AUTO_MASTER_CONF, 'w') as write_amc:
            write_amc.write(AUTO_MASTER_OPTION+'\n')

    # handle the rest of the fileshare setup
    try:
        os.makedirs('/root/.ssh/', exist_ok=True)
        with open('/root/.ssh/known_hosts', 'w') as f:
            pass

        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/.ssh/known_hosts') == 0:
                    print(f"successfully logged into {at.get('tag')}")
                else:
                    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)
    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')
    if not os.path.exists('/usr/bin/sshpass'):
        raise ValueError('Package \'sshpass\' not found')
    if not os.path.exists('/usr/bin/ssh-copy-id'):
        raise ValueError('Package \'ssh-copy-id\' 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)


