import os
import sys
import subprocess
from multiprocessing import Process, Manager, Pool, TimeoutError, freeze_support, active_children
from sys import platform
from time import sleep
import uuid 
import yaml
from util import find_data_file
from util import fprint
from util import kill
from util import run_cmd
import taskbartool
import util
import netstat
import ssh
import auth
import panel
import block

badapps = [756, 278670]
badips = ["208.59.79.12",]
displaydata = None
settings = None
netdata_res = None
procdata_res = None
killme = None
ppanel = None
datafile = ""
#print(datafile)
config = None
interval = 10
win32 = platform == "win32"
linux = platform == "linux" or platform == "linux2"
macos = platform == "darwin"
# Get unique system values
if win32:
    sysid = hex(uuid.getnode())
    datafile += sysid
    datafile += "gendata.csv"
    # Python is running as Administrator (so netstat can get filename, to block, etc), 
    # so we use this to see who is actually logged in
    # it's very hacky
    startupinfo = subprocess.STARTUPINFO()
    #if not getattr(sys, "frozen", False):
    startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW # hide powershell window
    res = subprocess.check_output(["WMIC", "ComputerSystem", "GET", "UserName"], universal_newlines=True, startupinfo=startupinfo)
    _, username = res.strip().rsplit("\n", 1)
    userid, sysdom = username.rsplit("\\", 1)

if linux:
    sysid = hex(uuid.getnode())
    #fprint(sysid)
    datafile += sysid
    datafile += "gendata.csv"
    res = subprocess.check_output(["who",], universal_newlines=True)
    userid = res.strip().split(" ")[0]
    sysdom = subprocess.check_output(["hostname",], universal_newlines=True).strip()
    #fprint(hostname)

def netstat_done(res):
    fprint("netstat done, processing")
    procdata_res = pool.apply_async(netstat.process, (res,), callback=process_done)
    #netstat.process(res)

def process_done(res):
    if settings["running"] == True:
        fprint("uploading to sftp...")
        #ssh.sftp_send_data(res, config, datafile)
        procdata_res = pool.apply_async(ssh.sftp_send_data, (config, datafile, 'send'), callback=upload_done)
    
def upload_done(res):
    settings["block"] = True

def login_done(res):
    if not res:
        fprint("Login failure")
        settings["message"] = "Login failure"
    else:
        fprint("Login result in main: " + str(res))
        settings["loggedin"] = res
        settings["continueui"] = True

def blockdata_done(res):
    global settings
    fprint("FINISHED downloading block data")
    tmpkill = settings["kill"]
    settings["kill"] = False
    #block_res = pool.apply_async(block.block_conn, (config, datafile, res, settings))
    block_pids, block_ips, block_data = block.block_conn(config, datafile, res)
    tmplist = settings["badapps"]
    
    for x in block_pids:
        fprint(x)
        if not x in tmplist:
            tmplist.append(x)
    settings["badapps"] = tmplist
    fprint(settings["badapps"])
    
    tmplist = settings["badips"]
    for x in block_ips:
        fprint(x)
        if not x in tmplist:
            tmplist.append(x)
    settings["badips"] = tmplist
    fprint(settings["badips"])

    settings["kill"] = tmpkill

    tmplist = settings["badlines"]
    for x in block_data:
        fprint(x)
        if not x in tmplist:
            tmplist.append(x)
    settings["badlines"] = tmplist
    fprint(settings["badlines"])

    settings["newdata"] = True
        
def killall():
    kids = active_children()
    for kid in kids:
        kid.kill()
    fprint("Every child has been killed")
    os.kill(os.getpid(), 9) # dirty kill of self



def mainloop(pool):
    # worker pool: netstat, netstat cleanup, upload, download, ui tasks
    global config
    global counter
    global netdata_res
    global procdata_res
    global rawdata
    global killme
    global ppanel
    #print(killme)
    if killme.value > 0:
        #print("killing")
        killall()
    #print(res.get(timeout=1))
    if counter == 0: # runs every INTERVAL
        #fprint("start loop")
        if netdata_res is None or netdata_res.ready():
            #rawdata = netdata_res.get()
            #procdata_res = pool.apply_async(process_netstat, (rawdata))
            fprint("netstat starting")
            netdata_res = pool.apply_async(netstat.start, callback=netstat_done)
            #fprint(netdata_res.successful())
    
    # runs every 50ms
    
    if settings["continueui"] == True:
        settings["continueui"] = False
        if ppanel is not None:
            # login panel is already open
            ppanel.terminate()
            ppanel = Process(target=panel.openwindow, args=(displaydata,settings,killme))
            ppanel.start()  

    if settings["showui"] == True:
        settings["showui"] = False
        ppanel = Process(target=panel.openwindow, args=(displaydata,settings,killme))
        ppanel.start()   

    if settings["login"] == True:
        login_res = pool.apply_async(auth.login, (config, settings["username"], settings["password"], sysid), callback=login_done)
        #fprint(auth.login(config, settings["username"], settings["password"], sysid))
        settings["login"] = False
        
    
    if settings["block"] == True and settings["running"] == True:
        blockdata_res = pool.apply_async(block.get_blocklist, (config,), callback=blockdata_done)
        #block.get_blocklist(config)
        settings["block"] = False

    if config["core"]["level"] == 0:
        settings["kill"] = False
        settings["fwll"] = False

    if config["core"]["level"] == 1:
        settings["kill"] = True
        settings["fwll"] = False

    if config["core"]["level"] == 2:
        settings["kill"] = False
        settings["fwll"] = True

    if config["core"]["level"] == 3:
        settings["kill"] = True
        settings["fwll"] = True

    if settings["kill"] == True:
        tmplist = settings["badapps"]
        settings["badapps"] = list()
        for x in tmplist:
            kill(x)

    if settings["fwll"] == True:
        tmplist = settings["badlines"]
        settings["badlines"] = list()
        for line in tmplist:
            badproto = line[1]
            badip = line[4]
            badport = line[5]
            fprint("Firewalling " + badip + ":" + str(badport))
            cmd = 'New-NetFirewallRule -DisplayName "IPPigeon Security Rule ' + badip + ':' + str(badport) + '" -Group "IPPigeon" -Direction Outbound -LocalPort Any -Protocol ' + badproto + ' -Action Block -RemoteAddress ' + badip + ' -RemotePort ' + str(badport)
            run_cmd(cmd)
        
    
    if settings["applyconfig"] == True:
        settings["applyconfig"] = False
        config = settings["config"]
        #fprint(settings["config"])
        with open(find_data_file('config.yml'), 'w') as filewrite:
            #global config
            yaml.dump(config, filewrite)
        fprint("Config saved!")

    sleep(interval / (interval * config["core"]["clockspeed"]))
    counter += 1
    if counter == interval * config["core"]["clockspeed"]:
        counter = 0


class Logger(object):
    def __init__(self, filename="output.log"):
        self.log = open(filename, "a")
        self.terminal = sys.stdout

    def write(self, message):
        self.log.write(message)
        #close(filename)
        #self.log = open(filename, "a")
        try:
            self.terminal.write(message)
        except:
            sleep(0)
        
    def flush(self):
        print("",end="")

if __name__ == '__main__':
    freeze_support() # required if packaged into single EXE
    # create manager to share data to me, background, foreground
    # create worker pool
    
    sys.stdout = Logger(filename=find_data_file("output.log"))
    sys.stderr = Logger(filename=find_data_file("output.log"))
    with Manager() as manager:
        with Pool(processes=5) as pool:
        
            with open(find_data_file('config.yml'), 'r') as fileread:
                #global config
                config = yaml.safe_load(fileread)
                #print(config['sftp']['host'])
            interval = config['core']['interval']
            displaydata = manager.list(range(2)) # data to be printed
            settings = manager.dict() # configuration 
            settings["login"] = False
            settings["loggedin"] = False
            settings["showui"] = False
            settings["continueui"] = False
            settings["killbox"] = list()
            settings["badapps"] = list()
            settings["badips"] = list()
            settings["badlines"] = list()
            settings["block"] = False
            settings["kill"] = False
            settings["config"] = config
            settings["applyconfig"] = False
            settings["fwll"] = 0
            settings["running"] = config["core"]["autostart"]
            settings["newdata"] = False
            killme = manager.Value('d', 0)
            #killme = False
            # launch background UI app as process
            p = Process(target=taskbartool.background, args=(displaydata,settings,killme))
            p.start()   
            #p.join() # not a foreground job, so let's not join it
            keeprunning = True

            # initial setup
            #netdata_res = pool.apply_async(netstat, callback=netstat_done)


            # launch loop - non-blocking!
            counter = 0
            while(keeprunning):
                mainloop(pool)