ff/ippigeon.py

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
from notification import send_notification
import taskbartool
import util
import netstat
import ssh
import auth
import panel
import block


history = list()
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:
        fprint("uploading to sftp...")
        #ssh.sftp_send_data(config, datafile, 'send')
        
        procdata_res = pool.apply_async(ssh.sftp_send_data, (config, datafile, 'send'), callback=upload_done)
    
def upload_done(res):
    settings["block"] = True
    tmpstat = settings["stats"]
    tmpstat[2] += 1
    settings["stats"] = tmpstat

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, goodct = block.block_conn(config, datafile, res)
    tmpstat = settings["stats"]
    tmpstat[1] += goodct
    if tmpstat[0] > 0 and goodct > 0:
        tmpstat[4] = 1.0 / (goodct * 100.0 / tmpstat[0])
    else: 
        tmpstat[4] = 0.0
    tmpstat[3] += 1
    settings["stats"] = tmpstat

    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 readstat_done(res):
    settings["stats"] = res
    fprint("Read stats!" + str(settings["stats"]))


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")
        util.clear_fwll() # clear the firewall rules before shutdown
        killall()
    #print(res.get(timeout=1))
    if counter == 0: # runs every INTERVAL
        #fprint("start loop")
        if settings["stats"][1] > 0:
            util.write_stats(settings["stats"])
        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, settings["appendbad"]), 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:
            send_notification("Killing PID " + str(x))
            kill(x)

    if settings["fwll"] == True:
        global history
        tmplist = settings["badlines"]
        tmpstat = settings["stats"]
        tmpstat[0] += len(tmplist) 
        settings["stats"] = tmpstat
        settings["badlines"] = list()
        for line in tmplist:
            badproto = line[1]
            badip = line[4]
            badport = line[5]
            
            if (badip, badport) not in history:
                fprint("Firewalling " + badip + ":" + str(badport))
                send_notification("Firewalling " + badip + ":" + str(badport))
                if win32:
                    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 linux:
                    cmd = "nft add rule ip ippigeon output ip daddr " + badip + " " + badproto.lower() + " dport " + str(badport) + " drop"
                    run_cmd(cmd)
            else:
                history.append((badip, badport))
        settings["badapps"] = list()
        
    
    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 
            # setup shared data variables
            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
            settings["appendbad"] = list()
            tmp = list()
            tmp.append(["N/A", "TCP", "N/A", "N/A", "20.112.52.29", "5000", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A"])
            tmp.append(["N/A", "TCP", "N/A", "N/A", "20.81.111.85", "80", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A"])
            tmp.append(["N/A", "TCP", "N/A", "N/A", "100.115.71.78", "5000", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A"])
            tmp.append(["N/A", "TCP", "N/A", "N/A", "100.115.71.78", "5000", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A"])
            tmp.append(["N/A", "TCP", "N/A", "N/A", "174.143.130.167", "443", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A"])
            tmp.append(["N/A", "TCP", "N/A", "N/A", "216.47.134.203", "443", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A"])
            tmp.append(["N/A", "TCP", "N/A", "N/A", "34.111.83.189", "443", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A"])
            settings["appendbad"] = tmp
            # connections blocked, total connections allowed, count of data sent, data recieved, ratio blocked-unblocked
            settings["stats"] = [0, 0, 0, 0, 0.0]

            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)
            if linux:
                # clear existing nftables entries, create new table
                run_cmd("nft delete table ip ippigeon")
                run_cmd("nft create table ip ippigeon")
                run_cmd("nft add chain ip ippigeon output { type filter hook output priority 0 \; policy accept\; }")
                #run_cmd("nft add chain ippigeon filter")
                
            pool.apply_async(util.read_stats, callback=readstat_done)
            # launch loop - non-blocking!
            counter = 0
            while(keeprunning):
                mainloop(pool)