class TShapper(Thread):
N_TOKENS = 8000
SLEEP_INTERVAL = 2 #Seconds
FILTER_DELETE_INTERVAL = 30 #Seconds
OLD_DEVICES_TIMEOUT = 300 #Seconds
USE_SNMP = False
def __init__(self):
Thread.__init__(self)
logging.basicConfig(filename='trottle.log',level=logging.DEBUG)
logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s',\
datefmt='%m/%d/%Y %I:%M:%S %p')
self.stopped = False
self.active_filters = 0
self.filter_delete_counter = 0
self.tokens = list()
self.devices = dict()
self.max_devices = {
'Number': 0,
'Time' : datetime.now()
}
self.speed = {
'Down' : 0,
'Up' : 0,
'Last_Update' : datetime.now()
}
self.filter = Filter()
self.generate_tokens()
if self.USE_SNMP:
self.SNMPDetector = SNMPDetector()
self.SNMPDetector.start()
def generate_tokens(self):
for i in range(1, self.N_TOKENS+1):
self.tokens.append(i)
def get_token(self):
return self.tokens.pop(0)
def release_token(self, token):
self.tokens.append(token)
def stop(self):
self.stopped = True
#Deletes the TC Rules
self.filter.destroy_tc_rules()
if self.USE_SNMP:
self.SNMPDetector.stop()
def restart(self):
#Deletes the TC Rules
self.filter.destroy_tc_rules()
#Restarts Objects
self.__init__()
#Starts the new TC Rules
self.filter.init_tc_rules()
def run(self):
#Starts the new TC Rules
self.filter.init_tc_rules()
#Updates the TC Rules
while not self.stopped:
try:
self.update()
except KernelError as e:
cprint("KERNEL ERROR: Restarting ", 'red')
logging.warning("KernelError:"+str(e))
sleep(2)
self.restart()
sleep(self.SLEEP_INTERVAL)
def update(self):
#cprint("Update", 'white')
if self.USE_SNMP:
updated_devices = self.get_devices_adresses_from_snmp()
else:
updated_devices = self.get_devices_adresses()
for device_mac, ips in updated_devices.items():
self.update_device(device_mac, ips)
self.clean_old_devices()
if len(self.devices) > self.max_devices['Number']:
self.max_devices['Number'] = len(self.devices)
self.max_devices['Time'] = datetime.now()
cprint("Current Clients/Max:"+str(len(self.devices))+"/"+\
str(self.max_devices['Number'])+" at "+\
str(self.max_devices['Time'])+" TokensUsed:"+\
str(self.N_TOKENS - len(self.tokens))+" ActiveFilters:"+\
str(self.active_filters)+self.get_speed(), 'green')
def clean_old_devices(self):
for device_mac, obj in self.devices.items():
if (datetime.now() - obj["last_seen"]).total_seconds() > self.OLD_DEVICES_TIMEOUT:
#obj = self.devices[device_mac]
token = obj["token"]
#Delete Filters
for ip in obj["ips"]:
self.filter.tc_del_filter(token, ip, obj)
self.active_filters -= 2
#Delete Class
self.filter.tc_del_class(token, obj)
del self.devices[device_mac]
self.release_token(token)
def update_device(self, device_mac, ips):
#Add new Device
if device_mac not in self.devices.keys():
obj = { "mac":device_mac,
"token":self.get_token(),
"ips":ips,
"last_seen":datetime.now(),
"prefs": { "lan": dict(), "wan": dict() }
}
self.devices[device_mac] = obj
self.filter.tc_add_device(obj)
for ip in ips:
self.filter.tc_add_filter(obj['token'], ip, obj)
self.active_filters += 2
#Modify Existing Device Rules
else:
old_ips = set(self.devices[device_mac]["ips"])
discovered_ips = set(ips)
ips_to_delete = old_ips - discovered_ips
ips_to_add = discovered_ips - old_ips
self.devices[device_mac]["ips"] = list(discovered_ips)
self.devices[device_mac]["last_seen"] = datetime.now()
if len(ips_to_add) > 0:
for ip in ips_to_add:
self.filter.tc_add_filter(\
self.devices[device_mac]['token'],
ip, self.devices[device_mac])
self.active_filters += 2
#if len(ips_to_delete) > 0 and self.filter_delete_counter > self.FILTER_DELETE_INTERVAL:
# for ip in ips_to_delete:
# self.filter.tc_del_filter(self.devices[device_mac]['token'], ip, self.devices[device_mac])
# self.active_filters -= 2
# self.filter_delete_counter = 0
#else:
# self.filter_delete_counter += 1
def print_devices(self):
for device, obj in self.devices.items():
print "\n"+device
for name, value in obj.items():
if name == "prefs":
print "\t\t"+name+":"
for x, v in value["lan"].items():
print "\t\t\tlan\t", x, v
for x, v in value["wan"].items():
print "\t\t\twan\t", x, v
else:
print "\t\t"+name+":", value
#Grabs device adresses
def get_devices_adresses(self):
proc = subprocess.Popen('ip neigh show dev '+self.filter.LAN_INTERFACE,\
shell=True, stdout=subprocess.PIPE)
output = proc.communicate()
lines = output[0].split('\n')
clients = dict()
for l in lines:
var = l.split()
if len(var) == 4:
ip = var[0]
mac = var[2]
state = var[3]
if ip[:6] == "fe80::":
continue
if not mac in clients.keys():
clients[mac] = [ip]
else:
clients[mac].append(ip)
return clients
def get_devices_adresses_from_snmp(self):
return self.SNMPDetector.get_adresses("2.4Ghz")
def get_speed(self):
proc = subprocess.Popen('ifconfig '+self.filter.WAN_INTERFACE+\
" | grep 'RX bytes'", shell=True,\
stdout=subprocess.PIPE)
output = proc.communicate()
line = output[0].split('\n')[0]
new_Down = int(line.split(':')[1].split()[0])
new_Up = int(line.split(':')[2].split()[0])
new_Time = datetime.now()
interval = (new_Time-self.speed["Last_Update"]).total_seconds()
down_Speed = round(((new_Down-self.speed["Down"])*8/interval)/1000000, 2)
up_Speed = round(((new_Up-self.speed["Up"])*8/interval)/1000000, 2)
self.speed = { 'Down' : new_Down,
'Up' : new_Up,
'Last_Update' : new_Time
}
return " Down/Up: "+str(down_Speed) + "/" + str(up_Speed)+" Mbps"
