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 import taskbartool import util import netstat import ssh import auth import panel import block badapps = [756, 278670] 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): fprint("uploading to sftp...") #ssh.sftp_send_data(res, config, datafile) procdata_res = pool.apply_async(ssh.sftp_send_data, (config, datafile, 'send')) 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): fprint("FINISHED downloading block data") #block_res = pool.apply_async(block.block_conn, (config, datafile, res)) block.block_conn(config, datafile, res) 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 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: blockdata_res = pool.apply_async(block.get_blocklist, (config,), callback=blockdata_done) #block.get_blocklist(config) settings["block"] = False #fprint(settings["killbox"]) if len(settings["killbox"]) > 0: fprint("Kill opportunity!") for proc in settings["killbox"]: pool.apply_async(kill, (proc,)) settings["killbox"].remove(proc) 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 Pool(processes=5) as pool: with Manager() as manager: with open(find_data_file('config.yml'), 'r') as file: #global config config = yaml.safe_load(file) #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"] = badapps settings["block"] = True 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)