From 0c6a7472bddb95d3aa3739c5879b47c1558a0e87 Mon Sep 17 00:00:00 2001
From: deajan <ozy@netpower.fr>
Date: Tue, 14 Jan 2025 22:38:58 +0100
Subject: [PATCH] upgrade_server: Implement server disabling

---
 upgrade_server/upgrade_server/api.py  | 21 ++++++++++++++++++---
 upgrade_server/upgrade_server/crud.py |  3 ++-
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/upgrade_server/upgrade_server/api.py b/upgrade_server/upgrade_server/api.py
index b9cb335..3677df5 100644
--- a/upgrade_server/upgrade_server/api.py
+++ b/upgrade_server/upgrade_server/api.py
@@ -119,6 +119,10 @@ async def current_version(
         crud.store_host_info(config_dict["upgrades"]["statistics_file"], host_id=data)
     except KeyError:
         logger.error("No statistics file set.")
+
+    if not crud.is_enabled():
+        return CurrentVersion(version="0.00-disabled")
+
     try:
         result = crud.get_current_version()
         if not result:
@@ -188,8 +192,12 @@ async def upgrades(
     except KeyError:
         logger.error("No statistics file set.")
 
-    # TODO:
-    # This can be amended by adding specific rules for host identity or groups or installed
+    if not crud.is_enabled():
+        raise HTTPException(
+            status_code=503,
+            detail="Service is currently disabled for maintenance"
+        )
+    
     file = FileGet(platform=platform, arch=arch, build_type=build_type, auto_upgrade_host_identity=auto_upgrade_host_identity, installed_version=installed_version, group=group)
     try:
         result = crud.get_file(file)
@@ -252,11 +260,18 @@ async def download(
         "arch": arch.value,
         "build_type": build_type.value,
     }
-
+    
     try:
         crud.store_host_info(config_dict["upgrades"]["statistics_file"], host_id=data)
     except KeyError:
         logger.error("No statistics file set.")
+
+    if not crud.is_enabled():
+        raise HTTPException(
+            status_code=503,
+            detail="Service is currently disabled for maintenance"
+        )
+    
     file = FileGet(platform=platform, arch=arch, build_type=build_type, auto_upgrade_host_identity=auto_upgrade_host_identity, installed_version=installed_version, group=group)
     try:
         result = crud.get_file(file, content=True)
diff --git a/upgrade_server/upgrade_server/crud.py b/upgrade_server/upgrade_server/crud.py
index b270331..b7b0a55 100644
--- a/upgrade_server/upgrade_server/crud.py
+++ b/upgrade_server/upgrade_server/crud.py
@@ -53,7 +53,8 @@ def sha256sum_data(data):
 
 
 def is_enabled() -> bool:
-    return not os.path.isfile("DISABLED")
+    path = os.path.join(config_dict["upgrades"]["data_root"], "DISABLED")
+    return not os.path.isfile(path)
 
 
 def store_host_info(destination: str, host_id: dict) -> None: