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 or macos: 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: #settings["newdata"] = 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 util.clear_fwll() 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)