Added psutil workaround in case it's not installed.

This commit is contained in:
Anton Hvornum 2019-04-10 10:56:04 +02:00
parent fe9dfc4ffd
commit 2037b6c0c1
1 changed files with 60 additions and 2 deletions

View File

@ -1,6 +1,6 @@
#!/usr/bin/python3 #!/usr/bin/python3
import traceback import traceback
import psutil, os, re, struct, sys, json import os, re, struct, sys, json
import urllib.request, urllib.parse, ssl import urllib.request, urllib.parse, ssl
from glob import glob from glob import glob
#from select import epoll, EPOLLIN, EPOLLHUP #from select import epoll, EPOLLIN, EPOLLHUP
@ -9,6 +9,64 @@ from collections import OrderedDict as oDict
from subprocess import Popen, STDOUT, PIPE from subprocess import Popen, STDOUT, PIPE
from time import sleep from time import sleep
try:
import psutil
except:
## Time to monkey patch in all the stats and psutil fuctions if it isn't installed.
class mem():
def __init__(self, free, percent=-1):
self.free = free
self.percent = percent
class disk():
def __init__(self, size, free, percent):
self.size = size
self.free = free
self.percent = percent
class iostat():
def __init__(self, interface, bytes_sent=0, bytes_recv=0):
self.interface = interface
self.bytes_recv = int(bytes_recv)
self.bytes_sent = int(bytes_sent)
def __repr__(self, *args, **kwargs):
return f'iostat@{self.interface}[bytes_sent: {self.bytes_sent}, bytes_recv: {self.bytes_recv}]'
class psutil():
def cpu_percent(interval=0):
## This just counts the ammount of time the CPU has spent. Find a better way!
with cmd("grep 'cpu ' /proc/stat | awk '{usage=($2+$4)*100/($2+$4+$5)} END {print usage}'") as output:
for line in output:
return float(line.strip().decode('UTF-8'))
def virtual_memory():
with cmd("grep 'MemFree: ' /proc/meminfo | awk '{free=($2)} END {print free}'") as output:
for line in output:
return mem(float(line.strip().decode('UTF-8')))
def disk_usage(partition):
disk_stats = os.statvfs(partition)
free_size = disk_stats.f_bfree * disk_stats.f_bsize
disk_size = disk_stats.f_blocks * disk_stats.f_bsize
percent = (100/disk_size)*free_size
return disk(disk_size, free_size, percent)
def net_if_addrs():
interfaces = {}
for root, folders, files in os.walk('/sys/class/net/'):
for name in folders:
interfaces[name] = {}
return interfaces
def net_io_counters(pernic=True):
data = {}
for interface in psutil.net_if_addrs().keys():
with cmd("grep '{interface}:' /proc/net/dev | awk '{{recv=$2}}{{send=$10}} END {{print send,recv}}'".format(interface=interface)) as output:
for line in output:
data[interface] = iostat(interface, *line.strip().decode('UTF-8').split(' ',1))
return data
## FIXME: dependency checks (fdisk, lsblk etc) ## FIXME: dependency checks (fdisk, lsblk etc)
rootdir_pattern = re.compile('^.*?/devices') rootdir_pattern = re.compile('^.*?/devices')
@ -287,7 +345,7 @@ if __name__ == '__main__':
print('[E] Failed to setup a yubikey password, is it plugged in?') print('[E] Failed to setup a yubikey password, is it plugged in?')
exit(1) exit(1)
print(args) print(json.dumps(args, indent=4))
if not os.path.isfile(args['pwfile']): if not os.path.isfile(args['pwfile']):
#PIN = '0000' #PIN = '0000'