Renew ROC api key

Change-Id: Ieaf4ffd83a3cf6ea47071578fd48bbfd5da308be
diff --git a/person_detection/roc.py b/person_detection/roc.py
index 7488af9..003d121 100644
--- a/person_detection/roc.py
+++ b/person_detection/roc.py
@@ -6,46 +6,77 @@
 import requests
 import json
 from requests.structures import CaseInsensitiveDict
+import logging as log
 
 
 URL = "https://roc.menlo.aetherproject.org/aether-roc-api/aether/v4.0.0/connectivity-service-v4/"
 #URL = "https://roc.staging.aether.onlab.us/aether-roc-api/aether/v4.0.0/connectivity-service-v4/"
 
 
-def headers(key):
-    h = CaseInsensitiveDict()
-    h["Content-Type"] = "application/json"
-    h["Authorization"] = "Bearer " + key
-    return h
+class Roc(object):
 
-def get_mbr(key, device_group):
-    url = URL + "vcs/vcs/vcs-{}/slice/mbr".format(device_group)
-    response = requests.get(url, headers=headers(key))
-    if response.status_code != 200:
-        print("Failed to get mbr, status_code: {}".format(response.status_code))
-        return None
-    mbr = json.loads(response.text)
-    return mbr
+    def __init__(self, user, password):
+        self.user = user
+        self.password = password
+        self.key = self.get_key()
 
-def set_mbr(key, device_group, mbr):
-    print("Set {} mbr to {}".format(device_group, mbr))
-    m = {'uplink' : mbr}
-    url = URL + "vcs/vcs/vcs-{}/slice/mbr".format(device_group)
-    response = requests.post(url, headers=headers(key), json=m)
-    assert response.status_code == 201, "Failed to set mbr"
-    if response.status_code != 201:
-        print("Failed to set mbr, device_group:{}, mbr:{}, status_code: {}".format(device_group, mbr, response.status_code))
+    def headers(self):
+        h = CaseInsensitiveDict()
+        h["Content-Type"] = "application/json"
+        h["Authorization"] = "Bearer " + self.key
+        return h
+
+    def get_mbr(self, device_group):
+        url = URL + "vcs/vcs/vcs-{}/slice/mbr".format(device_group)
+        response = requests.get(url, headers=self.headers())
+        if response.status_code != 200:
+            log.error("Failed to get mbr, status_code: {}".format(response.status_code))
+            return None
+        mbr = json.loads(response.text)
+        return mbr
+
+    def set_mbr(self, device_group, mbr):
+        log.info("Set {} mbr to {}".format(device_group, mbr))
+        m = {'uplink' : mbr}
+        url = URL + "vcs/vcs/vcs-{}/slice/mbr".format(device_group)
+        response = requests.post(url, headers=self.headers(), json=m)
+
+        # If error, renew api key and try again
+        if response.status_code != 201:
+            log.info("Renew ROC api key")
+            self.key = self.get_key()
+            response = requests.post(url, headers=self.headers(), json=m)
+            if response.status_code != 201:
+                log.error("Failed to set mbr, device_group:{}, mbr:{}, status_code: {}".format(device_group, mbr, response.status_code))
+                sys.exit()
+
+    def get_key(self):
+        url = "https://keycloak.opennetworking.org/auth/realms/master/protocol/openid-connect/token"
+        headers = CaseInsensitiveDict()
+        headers["Content-Type"] = "application/x-www-form-urlencoded"
+        data = {
+                'grant_type': 'password',
+                'client_id': 'aether-roc-gui',
+                'username': self.user,
+                'password': self.password,
+                'scope': 'openid profile email groups'
+                }
+        response = requests.post(url, data, headers)
+        key = json.loads(response.text)['access_token']
+        return key
 
 if __name__ == '__main__':
-    key = ""
-    cameras = "menlo-4g-cameras"
-    #cameras = "cameras-4g"
 
-    mbr = get_mbr(key, cameras)
+    # use valid keycloak user/password
+    roc = Roc("user", "password")
+
+    cameras = "menlo-4g-cameras"
+
+    mbr = roc.get_mbr(cameras)
     print("uplink mbr:{}, downlink mbr: {}".format(mbr["uplink"], mbr["downlink"]))
-    set_mbr(key, cameras, 5000000)
-    mbr = get_mbr(key, cameras)
+    roc.set_mbr(cameras, 5000000)
+    mbr = roc.get_mbr(cameras)
     print("uplink mbr:{}, downlink mbr: {}".format(mbr["uplink"], mbr["downlink"]))
-    set_mbr(key, cameras, 10000000)
-    mbr = get_mbr(key, cameras)
+    roc.set_mbr(cameras, 10000000)
+    mbr = roc.get_mbr(cameras)
     print("uplink mbr:{}, downlink mbr: {}".format(mbr["uplink"], mbr["downlink"]))