#!/usr/bin/env python3

import configparser
import signal
import sys
import time
import subprocess
from types import NoneType
import paramiko
import base64
import struct
import os
import json
import threading
import socket
import shutil

# List of running threads
threads = []
running = False

def sighandler(sig, frame):
    global threads
    global running

    print("\rStopping server")
    running = False
    for t in threads:
        t.stop();

class FileTransferSF(threading.Thread):
    def __init__(self, channel, fname, identifier, fsize):
        threading.Thread.__init__(self)
        self.channel = channel
        self.fname = fname
        self.identifier = identifier
        self.fsize = fsize
        self.running = True
    def stop(self):
        self.running = False
    def run(self):
        with open(f"{self.identifier}/{self.fname}", 'wb') as f:
            fsize = self.fsize
            while fsize>0 and self.running:
                fdata = self.channel.recv(1024)
                f.write(fdata)
                fsize -= 1024

class FileTransferRF(threading.Thread):
    def __init__(self, channel, fname, identifier):
        threading.Thread.__init__(self)
        self.channel = channel
        self.fname = fname
        self.identifier = identifier
        self.running = True
    def stop(self):
        self.running = False
    def run(self):
        with open(f"{self.identifier}/{self.fname}", 'rb') as f:
            stat = os.fstat(f.fileno())
            print('  -> fsize', stat.st_size)
            fsize = struct.pack('>q', stat.st_size)
            i = stat.st_size
            self.channel.sendall(b'OK'+fsize)
            while i>0 and self.running:
                fdata = f.read(1024)
                self.channel.sendall(fdata)
                i -= 1024

class SSHServer(paramiko.ServerInterface):
    def __init__(self, authorized):
        self.event = threading.Event()
        self.authorized = authorized
        self.identifier = ''
        self.subprocesses = []
        self.threads = []

    def stop(self):
        self.event.set()
        for s in self.subprocesses:
            s.kill()
        for t in self.threads:
            if type(t) is not NoneType:
                t.stop()
                t.join()

    def check_channel_request(self, kind, chanid):
        if kind == 'session':
            return paramiko.OPEN_SUCCEEDED

    def check_auth_publickey(self, username, key):
        keyascii = key.get_base64()
        for auth in self.authorized:
            authascii = auth.split(' ')[1]
            if authascii==keyascii:
                return paramiko.AUTH_SUCCESSFUL
        return paramiko.AUTH_FAILED

    def get_allowed_auths(self, username):
        return 'publickey'

    def rstr(self, b):
        l = struct.unpack('>I', b[:4])[0]
        return (bytes.decode(b[4:4+l], 'utf-8'), b[4+l:])

    def sstr(self, s):
        return struct.pack('>I', len(s)) + s.encode('utf-8')

    def check_channel_exec_request(self, channel, command):
        try:
            command = base64.decodebytes(command)
            cmd = command[:2]
            data = command[2:]

            # Identifier
            if cmd==b'id':
                identifier = struct.unpack('>q', data[:8])[0]
                self.identifier = str(identifier)
                print('>', identifier)
                # Create directory
                if os.path.exists(str(identifier)):
                    shutil.rmtree(str(identifier))
                os.mkdir(str(identifier))
            
            # Exit
            elif cmd==b'ex':
                print('<', self.identifier)
                self.event.set()

            # Config
            elif cmd==b'cf':
                cnf, data = self.rstr(data)
                self.config = configparser.ConfigParser()
                self.config.read_dict(json.loads(cnf))

            # List files
            elif cmd==b'ls':
                dr, data = self.rstr(data)
                print('ls', dr)
                if not os.path.exists(f"{self.identifier}/{dr}"):
                    channel.sendall(b'ERFile not found')
                es = []
                for f in os.listdir(f'{self.identifier}/{dr}'):
                    if os.path.isfile(f'{self.identifier}/{dr}/{f}'):
                        df = 'f'
                    else:
                        df = 'd'
                    es.append(f'{df}{f}')
                channel.sendall(b'OK' + self.sstr('\n'.join(es)))

            # Send file
            elif cmd==b'sf':
                fname, data = self.rstr(data)
                fsize = struct.unpack('>q', data)[0]
                print('>>', fname, fsize)
                os.makedirs(os.path.dirname(f"{self.identifier}/{fname}"), exist_ok=True)
                channel.sendall(b'OK\n')
                t = FileTransferSF(channel, fname, self.identifier, fsize)
                self.threads.append(t)
                t.start()

            # Receive file
            elif cmd==b'rf':
                fname, data = self.rstr(data)
                print('<<', fname, self.identifier)
                if not os.path.exists(f"{self.identifier}/{fname}"):
                    channel.sendall(b'ERFile not found')
                else:
                    t = FileTransferRF(channel, fname, self.identifier)
                    self.threads.append(t)
                    t.start()

            # Execute rbuild
            elif cmd==b'do':
                args, data = self.rstr(data)
                print('[]', args)
                with open(f"{self.identifier}/project.cfg", "w") as f:
                    self.config.write(f)
                p = subprocess.Popen(f"rbuild -c project.cfg {args}", shell=True, cwd=f'{self.identifier}', stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
                self.subprocesses.append(p)
                res = p.wait()
                channel.sendall(struct.pack('>I', res))

            return True

        except Exception:
            global running
            if running:
                print("ERROR: Unknown error")
            return False

class Connection(threading.Thread):
    def __init__(self, sock, addr, host_key, authorized):
        threading.Thread.__init__(self)
        self.sock = sock
        self.addr = addr
        self.host_key = host_key
        self.authorized = authorized
        self.running = False
        print("Connection from", addr)

    def stop(self):
        self.server.event.set()
        self.server.stop()

    def clean(self):
        pass

    def run(self):
        transport = paramiko.Transport(self.sock)
        transport.set_gss_host(socket.getfqdn(""))
        transport.load_server_moduli()
        transport.add_server_key(self.host_key)
        server = SSHServer(self.authorized)
        transport.start_server(server=server)
        self.server = server
        while not server.event.is_set():
            if not transport.is_alive():
                print("Connection", self.addr, "is broken from other end")
                server.stop()
                break
            time.sleep(0.2)
        else:
            print("Connection", self.addr, "closed")
        if server.identifier!='':
            pass
            shutil.rmtree(server.identifier, True)
        transport.close()

def print_help():
    print("Unified FPGA synthesizer frontend - remote execution server\r\n(c) Joppe Blondel - 2022\r\n")
    print(f"Usage: {sys.argv[0]} [ OPTIONS ] host port privkey pubkey authorized_hosts")
    print("")
    print("Options:")
    print("  -h                 Show this help message")

if __name__=="__main__":
    # Parse arguments
    i = 1
    host = ''
    port = ''
    pubkey = ''
    privkey = ''
    authorized_f = ''
    while i<len(sys.argv):
        if sys.argv[i]=='-h':
            print_help()
            exit(0)
        else:
            if host=='':
                host = sys.argv[i]
            elif port=='':
                port = sys.argv[i]
            elif privkey=='':
                privkey = sys.argv[i]
            elif pubkey=='':
                pubkey = sys.argv[i]
            elif authorized_f=='':
                authorized_f = sys.argv[i]
        i += 1

    signal.signal(signal.SIGINT, sighandler)

    # Load SSH settings
    host_key = paramiko.RSAKey(filename=privkey)
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind((host, int(port)))
    sock.settimeout(0.2)
    sock.listen(100)
    # Get authorized hosts
    with open(authorized_f, 'r') as f:
        authorized = f.read().split('\n')

    running = True
    while running:
        try:
            client, addr = sock.accept()
            conn = Connection(client, addr, host_key, authorized)
            conn.start()
            threads.append(conn)
        except TimeoutError as e:
            pass

    for t in threads:
        t.join()
        t.clean()