From 5ba5e6eda17d34fc3fe327ac1ee11678f557d062 Mon Sep 17 00:00:00 2001 From: "Nicholas St. Germain" Date: Fri, 4 Jan 2019 15:30:27 -0600 Subject: [PATCH] Feature added: Unifi. #79 --- Varken.py | 8 ++++- data/varken.example.ini | 11 ++++++ varken/__init__.py | 2 +- varken/cisco.py | 5 ++- varken/helpers.py | 6 +++- varken/iniparser.py | 55 +++++++++++++----------------- varken/structures.py | 11 ++++++ varken/tautulli.py | 2 +- varken/unifi.py | 75 +++++++++++++++++++++++++++++++++++++++++ 9 files changed, 137 insertions(+), 38 deletions(-) create mode 100644 varken/unifi.py diff --git a/Varken.py b/Varken.py index 1951088..add9f13 100644 --- a/Varken.py +++ b/Varken.py @@ -11,6 +11,7 @@ from argparse import ArgumentParser, RawTextHelpFormatter from logging import getLogger, StreamHandler, Formatter, DEBUG from varken.ombi import OmbiAPI +from varken.unifi import UniFiAPI from varken.cisco import CiscoAPI from varken import VERSION, BRANCH from varken.sonarr import SonarrAPI @@ -135,8 +136,13 @@ if __name__ == "__main__": ASA = CiscoAPI(firewall, DBMANAGER) schedule.every(firewall.get_bandwidth_run_seconds).seconds.do(threaded, ASA.get_bandwidth) + if CONFIG.unifi_enabled: + for server in CONFIG.unifi_servers: + UNIFI = UniFiAPI(server, DBMANAGER) + schedule.every(server.get_usg_stats_run_seconds).seconds.do(threaded, UNIFI.get_usg_stats) + # Run all on startup - SERVICES_ENABLED = [CONFIG.ombi_enabled, CONFIG.radarr_enabled, CONFIG.tautulli_enabled, + SERVICES_ENABLED = [CONFIG.ombi_enabled, CONFIG.radarr_enabled, CONFIG.tautulli_enabled, CONFIG.unifi_enabled, CONFIG.sonarr_enabled, CONFIG.ciscoasa_enabled, CONFIG.sickchill_enabled] if not [enabled for enabled in SERVICES_ENABLED if enabled]: vl.logger.error("All services disabled. Exiting") diff --git a/data/varken.example.ini b/data/varken.example.ini index 6a8f67a..88d6d3c 100644 --- a/data/varken.example.ini +++ b/data/varken.example.ini @@ -5,6 +5,7 @@ tautulli_server_ids = 1 ombi_server_ids = 1 ciscoasa_server_ids = false sickchill_server_ids = false +unifi_server_ids = false [influxdb] url = influxdb.domain.tld @@ -95,3 +96,13 @@ outside_interface = WAN ssl = false verify_ssl = false get_bandwidth_run_seconds = 300 + +[unifi-1] +url = unifi.domain.tld:8443 +username = ubnt +password = ubnt +site = default +usg_name = MyRouter +ssl = false +verify_ssl = false +get_usg_stats_run_seconds = 300 \ No newline at end of file diff --git a/varken/__init__.py b/varken/__init__.py index 90e1972..452f1ad 100644 --- a/varken/__init__.py +++ b/varken/__init__.py @@ -1,2 +1,2 @@ -VERSION = 1.7 +VERSION = "1.6.1" BRANCH = 'pre-nightly' diff --git a/varken/cisco.py b/varken/cisco.py index ac04166..12621c2 100644 --- a/varken/cisco.py +++ b/varken/cisco.py @@ -7,7 +7,6 @@ from varken.helpers import connection_handler class CiscoAPI(object): def __init__(self, firewall, dbmanager): - self.now = datetime.now(timezone.utc).astimezone().isoformat() self.dbmanager = dbmanager self.firewall = firewall # Create session to reduce server web thread load, and globally define pageSize for all requests @@ -32,7 +31,7 @@ class CiscoAPI(object): self.session.headers = {'X-Auth-Token': post} def get_bandwidth(self): - self.now = datetime.now(timezone.utc).astimezone().isoformat() + now = datetime.now(timezone.utc).astimezone().isoformat() endpoint = '/api/monitoring/device/interfaces/' + self.firewall.outside_interface if not self.session.headers: @@ -50,7 +49,7 @@ class CiscoAPI(object): "tags": { "interface": self.firewall.outside_interface }, - "time": self.now, + "time": now, "fields": { "upload_bitrate": get['outputBitRate'], "download_bitrate": get['inputBitRate'] diff --git a/varken/helpers.py b/varken/helpers.py index 5d693c1..1db2691 100644 --- a/varken/helpers.py +++ b/varken/helpers.py @@ -90,7 +90,8 @@ def rfc1918_ip_check(ip): return rfc1918_ip -def connection_handler(session, request, verify): +def connection_handler(session, request, verify, as_is_reply=False): + air = as_is_reply s = session r = request v = verify @@ -114,6 +115,8 @@ def connection_handler(session, request, verify): if get.headers['X-Auth-Token']: return get.headers['X-Auth-Token'] + if air: + return get except InvalidSchema: logger.error("You added http(s):// in the config file. Don't do that.") @@ -123,6 +126,7 @@ def connection_handler(session, request, verify): except ConnectionError as e: logger.error('Cannot resolve the url/ip/port. Check connectivity. Error: %s', e) + return return_json diff --git a/varken/iniparser.py b/varken/iniparser.py index 4ad040b..c7a81ce 100644 --- a/varken/iniparser.py +++ b/varken/iniparser.py @@ -4,9 +4,9 @@ from os.path import join, exists from re import match, compile, IGNORECASE from configparser import ConfigParser, NoOptionError, NoSectionError -from varken.helpers import clean_sid_check, rfc1918_ip_check -from varken.structures import SickChillServer from varken.varkenlogger import BlacklistFilter +from varken.structures import SickChillServer, UniFiServer +from varken.helpers import clean_sid_check, rfc1918_ip_check from varken.structures import SonarrServer, RadarrServer, OmbiServer, TautulliServer, InfluxServer, CiscoASAFirewall @@ -15,7 +15,7 @@ class INIParser(object): self.config = None self.data_folder = data_folder self.filtered_strings = None - self.services = ['sonarr', 'radarr', 'ombi', 'tautulli', 'sickchill', 'ciscoasa'] + self.services = ['sonarr', 'radarr', 'ombi', 'tautulli', 'sickchill', 'ciscoasa', 'unifi'] self.logger = getLogger() self.influx_server = InfluxServer() @@ -167,29 +167,26 @@ class INIParser(object): url = self.url_check(self.config.get(section, 'url'), section=section) apikey = None - if service != 'ciscoasa': + if service not in ['ciscoasa', 'unifi']: apikey = self.config.get(section, 'apikey') scheme = 'https://' if self.config.getboolean(section, 'ssl') else 'http://' - verify_ssl = self.config.getboolean(section, 'verify_ssl') if scheme != 'https://': verify_ssl = False - if service == 'sonarr': + if service in ['sonarr', 'radarr']: queue = self.config.getboolean(section, 'queue') + queue_run_seconds = self.config.getint(section, 'queue_run_seconds') + if service == 'sonarr': missing_days = self.config.getint(section, 'missing_days') - future_days = self.config.getint(section, 'future_days') missing_days_run_seconds = self.config.getint(section, 'missing_days_run_seconds') - future_days_run_seconds = self.config.getint(section, 'future_days_run_seconds') - queue_run_seconds = self.config.getint(section, 'queue_run_seconds') - server = SonarrServer(id=server_id, url=scheme + url, api_key=apikey, verify_ssl=verify_ssl, missing_days=missing_days, future_days=future_days, missing_days_run_seconds=missing_days_run_seconds, @@ -197,12 +194,7 @@ class INIParser(object): queue=queue, queue_run_seconds=queue_run_seconds) if service == 'radarr': - queue = self.config.getboolean(section, 'queue') - - queue_run_seconds = self.config.getint(section, 'queue_run_seconds') - get_missing = self.config.getboolean(section, 'get_missing') - get_missing_run_seconds = self.config.getint(section, 'get_missing_run_seconds') server = RadarrServer(id=server_id, url=scheme + url, api_key=apikey, verify_ssl=verify_ssl, @@ -212,18 +204,17 @@ class INIParser(object): if service == 'tautulli': fallback_ip = self.config.get(section, 'fallback_ip') + get_stats = self.config.getboolean(section, 'get_stats') get_activity = self.config.getboolean(section, 'get_activity') get_activity_run_seconds = self.config.getint(section, 'get_activity_run_seconds') - - get_stats = self.config.getboolean(section, 'get_stats') - get_stats_run_seconds = self.config.getint(section, 'get_stats_run_seconds') invalid_wan_ip = rfc1918_ip_check(fallback_ip) if invalid_wan_ip: - self.logger.error('Invalid failback_ip [%s] set for %s-%s!', fallback_ip, service, server_id) + self.logger.error('Invalid fallback_ip [%s] set for %s-%s!', fallback_ip, service, + server_id) exit(1) server = TautulliServer(id=server_id, url=scheme + url, api_key=apikey, @@ -233,17 +224,13 @@ class INIParser(object): get_stats_run_seconds=get_stats_run_seconds) if service == 'ombi': + issue_status_counts = self.config.getboolean(section, 'get_issue_status_counts') request_type_counts = self.config.getboolean(section, 'get_request_type_counts') - - request_type_run_seconds = self.config.getint(section, 'request_type_run_seconds') - request_total_counts = self.config.getboolean(section, 'get_request_total_counts') - request_total_run_seconds = self.config.getint(section, 'request_total_run_seconds') - - issue_status_counts = self.config.getboolean(section, 'get_issue_status_counts') - issue_status_run_seconds = self.config.getint(section, 'issue_status_run_seconds') + request_type_run_seconds = self.config.getint(section, 'request_type_run_seconds') + request_total_run_seconds = self.config.getint(section, 'request_total_run_seconds') server = OmbiServer(id=server_id, url=scheme + url, api_key=apikey, verify_ssl=verify_ssl, request_type_counts=request_type_counts, @@ -255,20 +242,18 @@ class INIParser(object): if service == 'sickchill': get_missing = self.config.getboolean(section, 'get_missing') - get_missing_run_seconds = self.config.getint(section, 'get_missing_run_seconds') server = SickChillServer(id=server_id, url=scheme + url, api_key=apikey, verify_ssl=verify_ssl, get_missing=get_missing, get_missing_run_seconds=get_missing_run_seconds) - if service == 'ciscoasa': + if service in ['ciscoasa', 'unifi']: username = self.config.get(section, 'username') - password = self.config.get(section, 'password') + if service == 'ciscoasa': outside_interface = self.config.get(section, 'outside_interface') - get_bandwidth_run_seconds = self.config.getint(section, 'get_bandwidth_run_seconds') server = CiscoASAFirewall(id=server_id, url=scheme + url, verify_ssl=verify_ssl, @@ -276,8 +261,16 @@ class INIParser(object): outside_interface=outside_interface, get_bandwidth_run_seconds=get_bandwidth_run_seconds) - getattr(self, f'{service}_servers').append(server) + if service == 'unifi': + site = self.config.get(section, 'site').lower() + usg_name = self.config.get(section, 'usg_name') + get_usg_stats_run_seconds = self.config.getint(section, 'get_usg_stats_run_seconds') + server = UniFiServer(id=server_id, url=scheme + url, verify_ssl=verify_ssl, site=site, + username=username, password=password, usg_name=usg_name, + get_usg_stats_run_seconds=get_usg_stats_run_seconds) + + getattr(self, f'{service}_servers').append(server) except NoOptionError as e: self.logger.error('Missing key in %s. Error: %s', section, e) self.rectify_ini() diff --git a/varken/structures.py b/varken/structures.py index 0672f50..296217a 100644 --- a/varken/structures.py +++ b/varken/structures.py @@ -86,6 +86,17 @@ class CiscoASAFirewall(NamedTuple): verify_ssl: bool = False +class UniFiServer(NamedTuple): + get_usg_stats_run_seconds: int = 30 + id: int = None + password: str = 'ubnt' + site: str = None + url: str = 'unifi.domain.tld:8443' + username: str = 'ubnt' + usg_name: str = None + verify_ssl: bool = False + + # Shared class Queue(NamedTuple): downloadId: str = None diff --git a/varken/tautulli.py b/varken/tautulli.py index 7d5ec78..3678703 100644 --- a/varken/tautulli.py +++ b/varken/tautulli.py @@ -52,7 +52,7 @@ class TautulliAPI(object): geodata = self.geoiphandler.lookup(session.ip_address_public) except (ValueError, AddressNotFoundError): if self.server.fallback_ip: - # Try the failback ip in the config file + # Try the fallback ip in the config file try: geodata = self.geoiphandler.lookup(self.server.fallback_ip) except AddressNotFoundError as e: diff --git a/varken/unifi.py b/varken/unifi.py new file mode 100644 index 0000000..0701b59 --- /dev/null +++ b/varken/unifi.py @@ -0,0 +1,75 @@ +from time import time +from logging import getLogger +from requests import Session, Request +from datetime import datetime, timezone + +from varken.helpers import connection_handler + + +class UniFiAPI(object): + def __init__(self, server, dbmanager): + self.dbmanager = dbmanager + self.server = server + # Create session to reduce server web thread load, and globally define pageSize for all requests + self.session = Session() + self.logger = getLogger() + + self.get_cookie() + + def __repr__(self): + return f"" + + def get_cookie(self): + endpoint = '/api/login' + pre_cookies = {'username': self.server.username, 'password': self.server.password, 'remember': True} + req = self.session.prepare_request(Request('POST', self.server.url + endpoint, json=pre_cookies)) + post = connection_handler(self.session, req, self.server.verify_ssl, as_is_reply=True) + + if not post.cookies.get('unifises'): + return + + cookies = {'unifises': post.cookies.get('unifises')} + self.session.cookies.update(cookies) + + def get_usg_stats(self): + now = datetime.now(timezone.utc).astimezone().isoformat() + endpoint = f'/api/s/{self.server.site}/stat/device' + req = self.session.prepare_request(Request('GET', self.server.url + endpoint)) + get = connection_handler(self.session, req, self.server.verify_ssl) + + if not get: + return + + devices = {device['name']: device for device in get['data']} + if devices.get(self.server.usg_name): + device = devices[self.server.usg_name] + else: + self.logger.error("Could not find a USG named %s from your UniFi Controller", self.server.usg_name) + return + + influx_payload = [ + { + "measurement": "UniFi", + "tags": { + "model": device['model'], + "name": device['name'] + }, + "time": now, + "fields": { + "bytes_current": device['wan1']['bytes-r'], + "rx_bytes_total": device['wan1']['rx_bytes'], + "rx_bytes_current": device['wan1']['rx_bytes-r'], + "tx_bytes_total": device['wan1']['tx_bytes'], + "tx_bytes_current": device['wan1']['tx_bytes-r'], + "speedtest_latency": device['speedtest-status']['latency'], + "speedtest_download": device['speedtest-status']['xput_download'], + "speedtest_upload": device['speedtest-status']['xput_upload'], + "cpu_loadavg_1": device['sys_stats']['loadavg_1'], + "cpu_loadavg_5": device['sys_stats']['loadavg_5'], + "cpu_loadavg_15": device['sys_stats']['loadavg_15'], + "cpu_util": device['system-stats']['cpu'], + "mem_util": device['system-stats']['mem'], + } + } + ] + self.dbmanager.write_points(influx_payload)