ff/ippigeon.py

import taskbartool
import os
import subprocess
from multiprocessing import Process, Manager, Pool, TimeoutError, freeze_support, active_children
from sys import platform
from time import sleep
import time
import csv
import uuid 
import sys

displaydata = None
settings = None
netdata_res = None
procdata_res = None
killme = None

interval = 5
win32 = platform == "win32"
linux = platform == "linux" or platform == "linux2"
macos = platform == "darwin"
# Get unique system values
if win32:
    sysid = hex(uuid.getnode())
    # Python is running as Administrator (so netstat can get filename, to block, etc), 
    # so we use this to see who is actually logged in
    startupinfo = subprocess.STARTUPINFO()
    startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
    res = subprocess.check_output(["WMIC", "ComputerSystem", "GET", "UserName"], universal_newlines=True, startupinfo=startupinfo)
    _, username = res.strip().rsplit("\n", 1)
    userid, sysdom = username.rsplit("\\", 1)

def find_data_file(filename):
    if getattr(sys, "frozen", False):
        # The application is frozen
        datadir = os.path.dirname(sys.executable)
    else:
        # The application is not frozen
        # Change this bit to match where you store your data files:
        datadir = os.path.dirname(__file__)
    return os.path.join(datadir, filename)

def run_ps(cmd):
    if win32:
        startupinfo = subprocess.STARTUPINFO()
        startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
        completed = subprocess.run(["powershell", "-WindowStyle", "hidden", "-Command", cmd], capture_output=True, startupinfo=startupinfo)
        return completed

def netstat():
    print("netstat started")
    if win32:
        data = run_ps("netstat -n -o -b")
        return data

def netstat_done(res):
    print("netstat done")
    procdata_res = pool.apply_async(process_netstat, (res,))
    #process_netstat(res)
    #print(procdata_res.get())
    #netdata_res = pool.apply_async(netstat)

def process_netstat(data):
    print("netstat processing")
    if win32:
        #output = data.stdout
        #print(output)
        output = data.stdout.decode().split('\r\n') # split stdout into lines
        #print(output)
        if output[0].find("The requested operation requires elevation.") >= 0:
            #print("test3")
            raise PermissionError("Unable to acquire netstat data without admin!")
        #print("test2")
        output2 = list()
        output2.append([sysid, userid, sysdom, int( time.time() )]) # add metadata
        #print(output2)
        procname = ""
        """for x in range(4, len(output)):
                tmp = output[x].split(" ")
                print(tmp)
                tmp = [i for i in output[x] if i]
                print(tmp)
                print(len(tmp))
                if len(len(tmp) == 1):
                    procname = tmp[0]
                    print(x)
                else:
                    print(x)
                    output2[x] = list()
                    output2[x].append(procname)
                    output2[x].append(output[x].split(" "))
                    output2[x] = [i for i in output2[x] if i]
        output2 = [i for i in output2 if i]
        print(output2)"""
        x = len(output) - 1 # start at the end because filename comes after connection
        procname = "Unknown" # if the very last connection happens to have no file (yes, it's possible), we can say unknown
        while x > 3:
            string = output[x]
            #print("LINE: ", string)
            string_split = string.split(" ")
            string_split = [i for i in string_split if i]
            if string.find("Can not obtain ownership information") >= 0: # Higher privilige than us, must be system
                procname = "Windows System"
            elif string.find("]") >= 0 and string.find("[") == 1: # generic [file.exe]
                procname = string[2:-1]
            elif len(string_split) == 5: # actual netstat line
                tmp = [procname,] # add executable name first
                tmp.extend(string.split(" "))
                tmp = [i for i in tmp if i]
                #print(tmp)
                output2.append(tmp)
            #else: # In case of an extra line above file, or an empty line, ignore it
                #print("Garbage data", string)
            x = x - 1
        #output2 = output2[2:]
        print(output2)
        with open(find_data_file("gendata.csv"), "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerows(output2)
        print("done")

def killall():
    kids = active_children()
    for kid in kids:
        kid.kill()
    print("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
    print("start loop")
    
    global netdata_res
    global procdata_res
    global rawdata
    global killme
    print(killme)
    if killme.value > 0:
        #print("killing")
        killall()
    #print(res.get(timeout=1))
    if netdata_res is None or netdata_res.ready():
        #rawdata = netdata_res.get()
        #procdata_res = pool.apply_async(process_netstat, (rawdata))
        print("netstat starting")
        netdata_res = pool.apply_async(netstat, callback=netstat_done)
    sleep(interval)

if __name__ == '__main__':
    freeze_support() # required if packaged into single EXE
    # create manager to share data to me, background, foreground
    # create worker pool

    with Pool(processes=5) as pool:
        with Manager() as manager:
            displaydata = manager.list(range(2)) # data to be printed
            settings = manager.list(range(20)) # configuration
            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!
            while(keeprunning):
                mainloop(pool)