From d9b6e825f3d0812cca8f9dcc6d487135dcc3c619 Mon Sep 17 00:00:00 2001
From: "Nicholas St. Germain" <nick@cajun.pro>
Date: Tue, 23 Apr 2019 23:19:06 -0500
Subject: [PATCH] get sites and map for unifi even if its a description (alias)

---
 varken/unifi.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/varken/unifi.py b/varken/unifi.py
index a32cd52..eb24097 100644
--- a/varken/unifi.py
+++ b/varken/unifi.py
@@ -14,6 +14,7 @@ class UniFiAPI(object):
         self.logger = getLogger()
 
         self.get_cookie()
+        self.get_site()
 
     def __repr__(self):
         return f"<unifi-{self.server.id}>"
@@ -25,11 +26,27 @@ class UniFiAPI(object):
         post = connection_handler(self.session, req, self.server.verify_ssl, as_is_reply=True)
 
         if not post or not post.cookies.get('unifises'):
+            self.logger.error(f"Could not retrieve session cookie from UniFi Controller")
             return
 
         cookies = {'unifises': post.cookies.get('unifises')}
         self.session.cookies.update(cookies)
 
+    def get_site(self):
+        endpoint = '/api/self/sites'
+        req = self.session.prepare_request(Request('GET', self.server.url + endpoint))
+        get = connection_handler(self.session, req, self.server.verify_ssl)
+
+        if not get:
+            self.logger.error(f"Could not get list of sites from UniFi Controller")
+            return
+        site = [site['name'] for site in get['data'] if site['name'].lower() == self.server.site.lower()
+                or site['desc'].lower() == self.server.site.lower()]
+        if site:
+            self.server.site = site[0]
+        else:
+            self.logger.error(f"Could not map site {self.server.site} to a site id/alias")
+
     def get_usg_stats(self):
         now = datetime.now(timezone.utc).astimezone().isoformat()
         endpoint = f'/api/s/{self.server.site}/stat/device'