Code: Select all
mgmt.send_command('kill {0}'.format(session['Username']))
Code: Select all
#!/usr/bin/python3
import sys
import os
import psutil
import requests
import requests.packages.urllib3.util.connection as urllib3_cn
import time
import json
import socket
from config import API_URL
sys.dont_write_bytecode = True
sys.path.insert(0, "/usr/local/vpn")
import OpenVPN
import PPP
NICS = ['eth0']
TOKEN = None
def allowed_gai_family():
family = socket.AF_INET # force IPv4
return family
urllib3_cn.allowed_gai_family = allowed_gai_family
def request(backend, data=None) -> dict:
dst = '{0}{1}'.format(API_URL, backend)
headers = {
'content-type': 'application/json',
'x-t4v-apiversion': '1',
'x-t4v-locale': 'en_en'
}
if TOKEN:
headers['Authorization'] = 'Bearer {0}'.format(TOKEN)
req = None
timeout = (5, 10)
try:
if data is not None:
req = requests.post(dst, data=json.dumps(data), headers=headers, timeout=timeout)
else:
req = requests.get(dst, headers=headers, timeout=timeout)
req.raise_for_status()
d = json.loads(req.text)
except Exception as ex:
d = {'Code': -1, 'Error': 'Error parsing request: {0}\nResponse:\n{1}'.format(str(ex), req.text if req else None)}
if d['Code'] != 1000:
raise Exception(d['Error'])
return d
# Gather data
info = {}
timestamp = int(time.time())
# CPU info
cpu_load = psutil.cpu_percent(60)
# BW Info
c_bytes_recv, c_bytes_sent, c_packets_recv, c_packets_sent = [0, 0, 0, 0]
bw_in = bw_out = pps_in = pps_out = 0
if os.path.exists('/tmp/.bw'):
with open('/tmp/.bw', 'r') as bw_file:
l_time, l_bytes_recv, l_bytes_sent, l_packets_recv, l_packets_sent = (int(i) for i in bw_file.readline().split())
net_io_counters = psutil.net_io_counters(pernic=True)
for nic in net_io_counters:
if nic in NICS:
c_bytes_recv = net_io_counters[nic].bytes_recv
c_bytes_sent = net_io_counters[nic].bytes_sent
c_packets_recv = net_io_counters[nic].packets_recv
c_packets_sent = net_io_counters[nic].packets_sent
break
if l_time < timestamp:
t_diff = timestamp - l_time
bw_in = (c_bytes_recv - l_bytes_recv) / t_diff / 1024 / 1024 * 8 # Mbps
bw_out = (c_bytes_sent - l_bytes_sent) / t_diff / 1024 / 1024 * 8 # Mbps
pps_in = (c_packets_recv - l_packets_recv) / t_diff
pps_out = (c_packets_sent - l_packets_sent) / t_diff
with open('/tmp/.bw', 'w') as bw_file:
bw_file.write(' '.join(map(str, [timestamp, c_bytes_recv, c_bytes_sent, c_packets_recv, c_packets_sent])))
bw_load = (bw_in + bw_out) * 100 / 1000
if bw_load > 100:
bw_load = 100
# OpenVPN status
status = 0
for proc in psutil.process_iter():
try:
pinfo = proc.as_dict(attrs=['pid', 'name'])
except psutil.NoSuchProcess:
pass
else:
if pinfo['name'] == 'openvpn':
status = 1
break
# OpenVPN sessions
mgmts = []
sessions = []
for port in [5555, 5556]:
mgmt = OpenVPN.Management('127.0.0.1', port)
sessions += mgmt.parse_status(mgmt.send_command('status 3'))
mgmts.append(mgmt)
max_attempts = 5
response = {}
for attempt in range(0, max_attempts):
if attempt > 0:
time.sleep(5)
try:
attempt += 1
response = request('auth', {'GrantType': 'client_credentials', 'ClientID': 'OpenVPNServer'})
TOKEN = response['AccessToken']
response = request('vpn/logicals', {
'Load': int(max([cpu_load, bw_load])),
'Status': status,
'Sessions': sessions
})
break
except Exception as ex:
if attempt == max_attempts:
raise
if 'Doomed' in response:
for session in response['Doomed']:
for mgmt in mgmts:
# Killl openvpn session
mgmt.send_command('kill {0}'.format(session['Username']))
# Kill pptp sessions also
PPP.disconnect(session['Username'])