ff/util.py

import inspect
import sys
import subprocess
import os
from sys import platform
import time as t
from time import sleep
import uuid
import csv

win32 = platform == "win32"
linux = platform == "linux" or platform == "linux2"
macos = platform == "darwin"
datafile = ""
logMsg = ""
logCont = ""

settings = None

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(sysdom)
    #fprint("d")

def time():
    return int(t.time())

def kill(pid):
    setup_child()
    if pid > 0:
        fprint("Killing PID " + str(pid), settings)
        os.kill(int(pid), 9)
        fprint("Signal 9 sent to PID " + str(pid), settings)

def fprint(msg, settings = None):
    #if not getattr(sys, "frozen", False):
    setup_child()
    try:
        frm = inspect.stack()[1]
        
        mod = inspect.getmodule(frm[0])
        logMsg = '[' + mod.__name__ + ":" + frm.function + ']:' + str(msg)
        
        print(logMsg)
        if (settings is not None):
            tmpList = settings["logMsg"]
            tmpList.append(logMsg)
            settings["logMsg"] = tmpList
    except Exception as e:
        try:
            print('[????:' + frm.function + ']:', str(msg))
            print('[util:fprint]: ' + str(e))
        except:
            print('[????]:', str(msg))
        
        
   # else:
        #print(msg)

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_cmd(cmd):
    if win32:
        startupinfo = subprocess.STARTUPINFO()
        #print("DICKS")
        #if not getattr(sys, "frozen", False):
        #    print("test")
        #    
        #completed = subprocess.run(["powershell", "-Command", cmd], capture_output=True, startupinfo=startupinfo)
        #else:
        #    print("alt")
        startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW    # , "-WindowStyle", "hidden"
        fprint("running PS command: " + cmd, settings)
        completed = subprocess.run(["powershell", "-Command", cmd], capture_output=True, startupinfo=startupinfo)
        fprint("ran PS command successfully", settings)
        #completed = subprocess.run(["powershell", "-WindowStyle", "hidden", "-Command", cmd], capture_output=True, startupinfo=startupinfo)
        return completed
    if linux or macos:
        fprint("running sh command: " + cmd, settings)
        completed = subprocess.run(["sh", "-c", cmd], capture_output=True)
        fprint("ran sh command successfully", settings)
        return completed

def setup_child(sets=None):
    if not getattr(sys, "frozen", False):
        sys.stdout = Logger(filename=find_data_file("output.log"))
        sys.stderr = Logger(filename=find_data_file("output.log"))
    if sets is not None:
        settings = sets

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="")

def clear_fwll():
    if win32:
        run_cmd('Remove-NetFirewallRule -Group "IPPigeon"')
    if linux:
        run_cmd("nft delete table ip ippigeon")


def write_stats(stats):
    fprint("Writing stats", settings)
    tmp = list()
    tmp.append(["connections blocked", "connections allowed", "data uploaded", "data recieved", "block ratio"])
    tmp.append(stats)
    
    with open(find_data_file("stats.csv"), "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerows(tmp)
    fprint("Done writing stats", settings)

def read_stats():
    with open(find_data_file("stats.csv"), newline='') as csvfile:
        csvreader = csv.reader(csvfile, delimiter=',', quotechar='|')
        header = True
        fprint(csvreader, settings)
        data = list()
        for line in csvreader:
            fprint(line, settings)
            if header:
                header = False
                continue
            data = line
        for idx in range(len(data) - 1):
            data[idx] = int(data[idx])
        data[len(data) - 1] = float(data[len(data) - 1])
        return data