You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
245 lines
8.6 KiB
Python
245 lines
8.6 KiB
Python
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
|
|
import yaml
|
|
from fabric import Connection
|
|
|
|
displaydata = None
|
|
settings = None
|
|
netdata_res = None
|
|
procdata_res = None
|
|
killme = 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)
|
|
|
|
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()
|
|
#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"
|
|
completed = subprocess.run(["powershell", "-Command", cmd], capture_output=True, startupinfo=startupinfo)
|
|
#completed = subprocess.run(["powershell", "-WindowStyle", "hidden", "-Command", cmd], capture_output=True, startupinfo=startupinfo)
|
|
|
|
return completed
|
|
|
|
def setup_child():
|
|
sys.stdout = Logger(filename=find_data_file("output.log"))
|
|
sys.stderr = Logger(filename=find_data_file("output.log"))
|
|
|
|
def netstat():
|
|
setup_child()
|
|
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,), callback=process_done)
|
|
#process_netstat(res)
|
|
#print(procdata_res.get())
|
|
#netdata_res = pool.apply_async(netstat)
|
|
|
|
def process_done(res):
|
|
print("uploading to sftp...")
|
|
sftp_connect(4)
|
|
#procdata_res = pool.apply_async(sftp_connect, (res,))
|
|
|
|
def sftp_connect(res):
|
|
print("Sending data over SFTP")
|
|
c = Connection(host=config['sftp']['host'], user=config['sftp']['user'], port=config['sftp']['port'], connect_kwargs={"key_filename": find_data_file(config['sftp']['keyfile']),})
|
|
c.put(find_data_file(datafile), remote=config['sftp']['filepath']['send'])
|
|
command = 'ls ' + config['sftp']['filepath']['send']
|
|
c.run(command)
|
|
|
|
|
|
|
|
#cnopts = pysftp.CnOpts(knownhosts=find_data_file('known_hosts'))
|
|
|
|
#with pysftp.Connection(config['sftp']['host'], username=config['sftp']['user'], private_key=find_data_file(config['sftp']['keyfile']), cnopts=cnopts) as sftp:
|
|
#with sftp.cd(config['sftp']['filepath']['send']):
|
|
#sftp.put(find_data_file(datafile))
|
|
|
|
def process_netstat(data):
|
|
setup_child()
|
|
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(datafile), "w", newline="") as f:
|
|
writer = csv.writer(f)
|
|
writer.writerows(output2)
|
|
print("done creating csv")
|
|
|
|
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)
|
|
|
|
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")
|
|
self.terminal.write(message)
|
|
|
|
def flush(self):
|
|
print("")
|
|
|
|
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.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)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|