Enabled the script to switch off the slave TrueNAS

4 files changed, 173 insertions, 16 deletions

diff --git a/autoBackup/.env.example b/autoBackup/.env.example
new file mode 100644
index 0000000..a7c0436
--- /dev/null
+++ b/autoBackup/.env.example
@@ -0,0 +1,14 @@
+MASTER_HOST=192.168.69.2
+MASTER_KEY=*****************************************************************
+MASTER_REPLICATION_TASKS=replicateSpinningRust,autoReplicateTheVault
+
+SLAVE_HOST=192.168.69.4
+SLAVE_KEY==*****************************************************************
+SLAVE_REPLICATION_TASKS=localVMs/localVMs - fivehundred/localVMs,ReplicateDatabaseBackups
+
+POLLING_RATE=300
+
+MQTT_HOST=192.168.69.5
+MQTT_USER=eden
+MQTT_PASSWORD=***********
+SLAVE_PLUG_FRIENDLYNAME=TasmotaBackup
diff --git a/autoBackup/TasmotaCLI b/autoBackup/TasmotaCLI
new file mode 160000
+Subproject dd7790dab8d3fbea8f2b58eb4d5aaffc36b3cb0
diff --git a/autoBackup/autoBackup.py b/autoBackup/autoBackup.py
index fc5c6a7..9b5dee5 100644
--- a/autoBackup/autoBackup.py
+++ b/autoBackup/autoBackup.py
@@ -1,20 +1,48 @@
-from truenas_api_client import Client
 import requests
+import logging
 import dotenv
 import json
+import time
+import sys
 import os
 
+sys.path.insert(1, os.path.join(os.path.dirname(__file__), "TasmotaCLI"))
+import tasmotaMQTTClient
+
 env_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env")
 if os.path.exists(env_path):
     dotenv.load_dotenv(dotenv_path = env_path)
 
+logging.basicConfig( 
+    format = "[%(asctime)s]\t%(message)s", 
+    level = logging.INFO,
+    handlers=[
+        logging.FileHandler(os.path.join(os.path.dirname(__file__), "logs", "backup.log")),
+        logging.StreamHandler()
+    ]
+)
+
 class TrueNASAPIClient:
-    def __init__(self, host, api_key):
-        self.base_url = base_url = "http://%s/api/v2.0" % host
+    def __init__(self, host, api_key, replication_task_names = None):
+        self.host = host
+        self.base_url = "http://%s/api/v2.0" % host
         self.headers = {
             "Authorization": "Bearer " + api_key
         }
+        if replication_task_names is None:
+            self.replication_task_names = []
+        else:
+            self.replication_task_names = replication_task_names
 
+        self.running_replication_jobs = {}
+
+    @staticmethod
+    def filter_running_jobs(jobs):
+        return list(filter(
+            lambda i: i["method"] == "replication.run" and i["progress"]["percent"] != 100 and not i["state"] == "FAILED",
+            jobs
+        ))
+        
     def base_get(self, endpoint, payload = None):
         if payload is None:
             payload = {}
@@ -30,24 +58,138 @@ class TrueNASAPIClient:
     def get_websocket_connections(self):
         return self.base_get("/core/sessions")
     
-    def get_replication_naming_schemas(self):
-        return self.base_get("/replication/list_naming_schemas")
-    
     def get_jobs(self):
         return self.base_get("/core/get_jobs")
     
-    def get_replication_jobs(self):
-        return [i for i in self.get_jobs() if i["method"] == "replication.run"]
-    
     def get_running_replication_jobs(self):
         return [i for i in self.get_jobs() if i["method"] == "replication.run" and i["progress"]["percent"] != 100 and not i["state"] == "FAILED"]
     
-    def get_running_jobs(self):
-        return [i for i in self.get_jobs() if i["progress"]["percent"] != 100]
-    
     def get_replication_tasks(self):
-        return self.base_get("/replication")
+        return list(filter(lambda a: a["name"] in self.replication_task_names, self.base_get("/replication")))
+    
+    def run_replication_task(self, task_id):
+        req = requests.post(self.base_url + "/replication/id/%d/run" % task_id, headers = self.headers)
+        if not req.status_code == 200:
+            raise ConnectionError("API call failed (%d): '%s'" % (req.status_code, req.content.decode()))
+        return req.json()
+    
+    def is_ready(self):
+        return self.base_get("/system/ready")
+    
+    def shutdown(self):
+        req = requests.post(self.base_url + "/system/shutdown", headers = self.headers)
+        if not req.status_code == 200:
+            raise ConnectionError("API call failed (%d): '%s'" % (req.status_code, req.content.decode()))
+        return req.json()
+
+    def run_all_replication_tasks(self):
+        for task in self.get_replication_tasks():
+            job_id = self.run_replication_task(task["id"])
+            self.running_replication_jobs[job_id] = task["name"]
+            logging.info("Started replication task '%s' on '%s' with job id %d" % (task["name"], self.host, job_id))
+
+    def get_state_of_replication_jobs(self):
+        all_complete = True
+        for job in self.get_jobs():
+            if job["id"] in self.running_replication_jobs.keys():
+                if job["state"] == "RUNNING":
+                    all_complete = False
+                logging.info("Replication job '%s' on '%s' is currently '%s' (%d%%)" % (
+                    self.running_replication_jobs[job["id"]], self.host, job["state"], job["progress"]["percent"]
+                ))
+
+        if all_complete:
+            self.running_replication_jobs = {}
+            logging.info("No more running replication jobs on '%s'" % self.host)
+        return all_complete
+
+def check_if_all_complete(truenasclients):
+    logging.info("Slave plug '%s' is using %dw of power" % (os.environ["SLAVE_PLUG_FRIENDLYNAME"], get_mqtt().switch_energy['Power']))
+    all_complete = True
+    for truenas in truenasclients:
+        if not truenas.get_state_of_replication_jobs():
+            all_complete = False
+    return all_complete
+
+def get_mqtt(message = None):
+    return tasmotaMQTTClient.MQTTClient(
+        host = os.environ["MQTT_HOST"],
+        username = os.environ["MQTT_USER"],
+        password = os.environ["MQTT_PASSWORD"],
+        friendlyname = os.environ["SLAVE_PLUG_FRIENDLYNAME"],
+        message = message
+    )
+
+def wait_for_slave(slave):
+    while True:
+        time.sleep(int(os.environ["POLLING_RATE"]))
+        try:
+            logging.info("Slave is ready: " + str(slave.is_ready()))
+        except requests.exceptions.ConnectionError:
+            logging.info("'%s' hasn't booted, waiting for %d more seconds" % (slave.host, int(os.environ["POLLING_RATE"])))
+        else:
+            break
+    logging.info("Slave TrueNAS has booted and is ready for API requests")
+
+def wait_till_idle_power():
+    while True:
+        p = get_mqtt().switch_energy['Power']
+        logging.info("'%s' plug is using %dw of power" % (os.environ["SLAVE_PLUG_FRIENDLYNAME"], p))
+        if p == 0:
+            break
+
+def main():
+    if os.environ["MASTER_REPLICATION_TASKS"] != "":
+        tasks = os.environ["MASTER_REPLICATION_TASKS"].split(",")
+    else:
+        tasks = []
+    master = TrueNASAPIClient(
+        host = os.environ["MASTER_HOST"], 
+        api_key = os.environ["MASTER_KEY"], 
+        replication_task_names = tasks
+    )
+    if os.environ["SLAVE_REPLICATION_TASKS"] != "":
+        tasks = os.environ["SLAVE_REPLICATION_TASKS"].split(",")
+    else:
+        tasks = []
+    slave = TrueNASAPIClient(
+        host = os.environ["SLAVE_HOST"], 
+        api_key = os.environ["SLAVE_KEY"], 
+        replication_task_names = tasks
+    )
+
+    logging.info("Began autoBackup procedure")
+    m = get_mqtt()
+    logging.info("Slave plug '%s' is currently %s" % (m.friendlyname, m.switch_power))
+    if m.switch_power == "ON":
+        was_already_on = True
+    else:
+        was_already_on = False
+        get_mqtt("ON")
+        logging.info("Turned on the slave plug. Now waiting for it to boot")
+        wait_for_slave(slave)
+
+    master.run_all_replication_tasks()
+    slave.run_all_replication_tasks()
+    # while (not master.get_state_of_replication_jobs()) or (not slave.get_state_of_replication_jobs()):
+    while not check_if_all_complete([master, slave]):
+        time.sleep(int(os.environ["POLLING_RATE"]))
+    logging.info("All replication jobs on all hosts complete")
+
+    if was_already_on:
+        logging.info("The slave TrueNAS was turned on not by us, so stopping here")
+    else:
+        logging.info("The slave TrueNAS was turned on my us, so starting the shutdown procedure")
+        logging.info(json.dumps(slave.shutdown(), indent = 4))
+
+        # wait until the slave TrueNAS is using 0w of power, which implies it has finished shutting down,
+        # then turn off the power to it
+        wait_till_idle_power()
+        get_mqtt("OFF")
+        logging.info("Turned off the slave's plug")
+
+    logging.info("autoBackup procedure completed\n\n")
 
 if __name__ == "__main__":
-    truenas = TrueNASAPIClient(host = os.environ["SLAVE_HOST"], api_key = os.environ["SLAVE_KEY"])
-    print(json.dumps(truenas.get_replication_tasks(), indent = 4))
\ No newline at end of file
+    main()
+    
\ No newline at end of file
diff --git a/autoBackup/requirements.txt b/autoBackup/requirements.txt
index 323ef8c..2e7e54b 100644
--- a/autoBackup/requirements.txt
+++ b/autoBackup/requirements.txt
@@ -1,2 +1,3 @@
 python-dotenv
-requests
\ No newline at end of file
+requests
+docker