Feature added: Unifi. #79

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")

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

varken/__init__.py

@@ -1,2 +1,2 @@
-VERSION = 1.7
+VERSION = "1.6.1"
 BRANCH = 'pre-nightly'

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']

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
 
 

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()

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

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:

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"<unifi-{self.server.id}>"
+
+    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)