#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Main program of the Client written for CIRCLE Cloud.
The Client job is to help the ease of use of the cloud system.
"""

import os, sys, ctypes
import argparse
import subprocess
import logging
import tempfile
import time
import win32crypt
import windowsclasses
from shutil import copy, which
from collections import OrderedDict


logging.basicConfig(
    level = logging.INFO,
    format='%(asctime)s %(name)s: [%(levelname)s] %(message)s',
    datefmt='%Y-%m-%d %H:%M',
    handlers = [
        logging.FileHandler(tempfile.gettempdir() + "\\circle_client.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("circle_client")


def is_administrator():
    return ctypes.windll.shell32.IsUserAnAdmin() != 0

def custom_protocol_register(custom_protocol):
    """
    Custom protocol register based on RegistryHandler module
    Can raise AttributeError on wrongly provided dictionary chain

    Keyword arguments:
    @param dict_chain -- OrderedDictionary chain of the registry tree
    """
    logger.debug(custom_protocol)
    if not isinstance(custom_protocol, OrderedDict):
        raise AttributeError
    try:
        custom_arguments = windowsclasses.Struct()
        custom_arguments.registry = "HKCR"
        handler = windowsclasses.RegistryHandler(custom_arguments)
        handler.create_registry_from_dict_chain(custom_protocol, True)
    except:
        raise


def install_client(dir, exe):
    logger.info("Copy \"%s\" to \"%s\"" % (exe, os.path.abspath(dir)))
    try:
        os.makedirs(os.path.abspath(dir), exist_ok=True)
        copy(exe, dir)
        logger.info("Creating custom URL protocol handler for circle protocol")
        custom_protocol = OrderedDict([
            ('circle', ["default",
						"URL:circle Protocol",
						"URL Protocol",
						""]),
			('circle\\URL Protocol', ""),
            ('circle\\DefaultIcon', "cloud.ico"),
			('circle\\shell', {'open': {
				'command': (
					"\"%(dir)s\\%(exe)s\" " % { 'dir': os.path.abspath(dir), 'exe': os.path.basename(exe) } + r'"%1"')
			}})
		])
        custom_protocol_register(custom_protocol)
    except Exception as e:
        logger.error(e)
        logger.error("URL Protocol handler installation aborted!")
    if which("putty.exe") == None: 
        logger.warning("*** No putty.exe found in PATH ****");
    else:
        logger.info("---- Done ----")


def connect(vm):
    """
    Handles to connection to the Virtual Machines from the local
    machine

    Keyword arguments:
    vm.protocol -- SSH, NX and RDP possible
    vm.host     -- Address of the Virtual Computer
    vm.port     -- The port where we can access the Virtual Computer
    vm.user     -- Username used for the connection
    vm.password -- Password used for the connection
    """
    command = None
    logger.debug('Connect methods started')
    if vm.protocol == "SSH":
        logger.info('SSH protocol received')
        arguments = str("-ssh -P %(port)s -pw %(password)s"
                        " %(user)s@%(host)s" % {
                            'port': vm.port,
                            'password': vm.password,
                            'user': vm.user,
                            'host': vm.host})
        logger.info('Popen: putty.exe')
        logger.debug('Determined arguments: %s', arguments)
        subprocess.Popen("putty.exe " + arguments).wait()
    elif vm.protocol == "RDP":
        logger.debug('RDP protocol received')
        full_address = "full address:s:%s:%s" % (vm.host, vm.port)
        user = "username:s:%s" % vm.user
        config_file = "%s%s" % (
            tempfile.gettempdir(),
            "\\circle_" + vm.user + str(int(time.time() * 1000)) + ".rdp")
        logger.info('Creating config file %s' % config_file)
        password = win32crypt.CryptProtectData(vm.password.encode('utf-16-le'), 
            'psw', None, None, None, 0).hex()
        config = RPD_template.format(vm.user, vm.host, vm.port, password)
        logger.debug(config)
        f = open(config_file, 'w')
        f.write(config)
        f.close()
        logger.debug('Popen the config file: %s', config_file)
        subprocess.Popen("mstsc.exe " + config_file, shell=False, 
            stdout=subprocess.PIPE).communicate()  # mstsc.exe forks but does not close the stdout
        os.remove(config_file)

RPD_template = """username:s:{}
full address:s:{}:{}
authentication level:i:0
password 51:b:{}"""

class Struct:
    """
    A struct used for parameter passing

    Keyword arguments:
    state    -- State of the Virtual Computer (running, etc..)
    protocol -- SSH, NX and RDP possible
    host     -- Address of the Virtual Computer
    port     -- The port where we can access the Virtual Computer
    user     -- Username used for the connection
    password -- Password used for the connection
    """
    pass


def parse_arguments():
    """
    Argument parser, based on the argparse module

    Keyword arguments:
    @return args -- arguments given by console
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "uri", type=str, 
        help="Specific uri for conneting to the VM", nargs='?',
        default=None)
    parser.add_argument(
        "--install", metavar='DIR',
        help="Install to DIR (default: %(const)s)", type=str, nargs='?',
        const='c:\\Program Files\\circle',
        default=None)
    parser.add_argument(
        "--loglevel", type=str,
        help="The level of logging severity",
        choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
        default=None)
    return parser

def main():
    """
    Main program
    """
    if frozen := getattr(sys, 'frozen', False):
        """
        If the application is run as a bundle, the PyInstaller bootloader
        extends the sys module by a flag frozen=True and sets the app 
        path into variable _MEIPASS'.
        """
        application_path = sys._MEIPASS
    else:
        application_path = os.path.dirname(os.path.abspath(__file__))
    """
    print("path=%s frozen=%s" % (application_path, frozen))
    print( 'sys.argv[0] is', sys.argv[0] )
    print( 'sys.executable is', sys.executable )
    print( 'os.getcwd is', os.getcwd() )
    """
    parser = parse_arguments()
    args = parser.parse_args()
    if args.loglevel == None:
        level = os.environ.get('LOGLEVEL', 'INFO')
    else:
        level = args.loglevel
    logger.setLevel(level)
    logger.info("Started")
    try:
        if args.install != None:
            if not frozen:
                logger.error("Use the bundled version (circle_client.exe) instead\n")
            else:
                if not is_administrator():
                    logger.error("No Administrator rights to install\n")
                else:
                    install_client(args.install, sys.executable)
        elif args.uri is not None:
            vm = Struct()
            try:
                x, vm.protocol, vm.user, vm.password, vm.host, vm.port = args.uri.split(':', 5)
                logger.debug(("Received the following URI: %(handler)s:"
                    "%(protocol)s:%(user)s:password:%(host)s:%(port)s" % {
                    'handler': x,
                    'protocol': vm.protocol,
                    'user': vm.user,
                    'host': vm.host,
                    'port': vm.port
                     }))
                vm.protocol = vm.protocol.upper()
                connect(vm)
            except ValueError:
                logger.error("No valid URL passed as argumnet: %s\n", args.uri);
                parser.print_help()
            except FileNotFoundError as e:
                logger.error("Cannot connect: %s\n", e);
        else:
            parser.print_help()
    except Exception as e:
        logger.exception("An exception %s was raised", e)

if __name__ == "__main__":
    main()

