tprelog · January 21, 2024 18:41 · Aug 3, 2021 · Jul 15, 2020
diff --git a/unifi.py b/unifi.py
@@ -1,7 +1,13 @@
+import argparse
 import requests
 import json
 import urllib3
 
+import unifi_secrets
+
+HOST = unifi_secrets.unifi_host
+USER = unifi_secrets.unifi_username
+PASS = unifi_secrets.unifi_password
 
 UNIFI_LOGIN_PATH = '/api/auth/login'
 
@@ -42,4 +48,27 @@ def request(self, path, data={}, method='POST'):
         except KeyError:
             pass
 
-        return r
+        return r
+
+
+    def reboot(self):
+        r = self.request("/api/system/reboot", method='POST')
+        return r.ok
+
+    def shutdown(self):
+        r = self.request("/api/system/poweroff", method='POST')
+        return r.ok
+
+
+if __name__ == "__main__":
+    udm = Unifi(HOST,USER,PASS)
+    udm.login()
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('command', help='reboot | shutdown')
+    args = parser.parse_args()
+
+    if args.command == 'reboot':
+        udm.reboot()
+    elif args.command == 'shutdown':
+        udm.shutdown()
diff --git a/unifi_secrets.py b/unifi_secrets.py
@@ -0,0 +1,3 @@
+unifi_host = ''
+unifi_username = ''
+unifi_password = ''
diff --git a/unifi.py b/unifi.py
@@ -0,0 +1,45 @@
+import requests
+import json
+import urllib3
+
+
+UNIFI_LOGIN_PATH = '/api/auth/login'
+
+# Disable SSL verification warnings
+urllib3.disable_warnings()
+
+
+class Unifi(object):
+    def __init__(self, host, username, password):
+        self.host = host
+        self.username = username
+        self.password = password
+        self.session = requests.Session()
+        self.csrf = ""
+
+    def login(self):
+        payload = {
+            "username": self.username,
+            "password": self.password,
+        }
+
+        r = self.request(UNIFI_LOGIN_PATH, payload)
+        return r.ok
+
+    def request(self, path, data={}, method='POST'):
+        uri = 'https://{}{}'.format(self.host, path)
+        headers = {
+            "Accept": "application/json",
+            "Content-Type": "application/json; charset=utf-8",
+        }
+
+        if self.csrf:
+            headers["X-CSRF-Token"] = self.csrf
+
+        r = getattr(self.session, method.lower())(uri, json=data, verify=False, headers=headers)
+        try:
+            self.csrf = r.headers['X-CSRF-Token']
+        except KeyError:
+            pass
+
+        return r