From 590dc4b65e1a86936a7b67d1b55fea47df84769f Mon Sep 17 00:00:00 2001
From: clore <support@clore.ai>
Date: Wed, 4 Sep 2024 00:51:52 +0000
Subject: [PATCH] add optional whitelist for outside images

---
 clore_hosting/main.py |  4 +++-
 lib/config.py         |  3 ++-
 lib/utils.py          | 23 ++++++++++++++++++++++-
 3 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/clore_hosting/main.py b/clore_hosting/main.py
index 7359ff4..5779abd 100644
--- a/clore_hosting/main.py
+++ b/clore_hosting/main.py
@@ -121,6 +121,8 @@ class CloreClient:
 
         nvml.init(allow_hive_binaries=not self.dont_use_hive_binaries)
 
+        self.extra_allowed_images = utils.get_extra_allowed_images()
+
         self.gpu_oc_specs = nvml.get_gpu_oc_specs()
         self.last_oc_service_submit = 0
         self.last_applied_oc = {}
@@ -417,7 +419,7 @@ class CloreClient:
                         if type(result)==types.ServerConfig:
                             if result.success:
                                 self.last_checked_ws_peers = utils.unix_timestamp()
-                                self.allowed_images=result.allowed_images
+                                self.allowed_images=result.allowed_images+self.extra_allowed_images
                                 if not config.debug_ws_peer:
                                     for pure_ws_peer in result.ws_peers:
                                         self.ws_peers[pure_ws_peer]={
diff --git a/lib/config.py b/lib/config.py
index 2fae529..f163b67 100644
--- a/lib/config.py
+++ b/lib/config.py
@@ -48,7 +48,8 @@ parser.add_argument('--startup-scripts-folder', type=str, default='/opt/clore-ho
 parser.add_argument('--wireguard-config-folder', type=str, default='/opt/clore-hosting/wireguard/configs', help='Folder with wireguard configs')
 parser.add_argument('--entrypoints-folder', type=str, default='/opt/clore-hosting/entrypoints', help='Folder with custom entrypoints')
 parser.add_argument('--debug-ws-peer', type=str, help="Specific ws peer to connect to (for debugging only)")
-parser.add_argument('--gpu-specs-file', type=str, default='/opt/clore-hosting/client/gpu_specs.json' ,help="Cache with specs of GPU possible OC/Power limit changes")
+parser.add_argument('--gpu-specs-file', type=str, default='/opt/clore-hosting/client/gpu_specs.json', help="Cache with specs of GPU possible OC/Power limit changes")
+parser.add_argument('--extra-allowed-images-file', type=str, default="/opt/clore-hosting/extra_allowed_images.json", help="Docker image whitelist, that are allowed by clore.ai hosting software")
 
 # Parse arguments, ignoring any non-defined arguments
 args, _ = parser.parse_known_args()
diff --git a/lib/utils.py b/lib/utils.py
index c6b9aa4..a34347c 100644
--- a/lib/utils.py
+++ b/lib/utils.py
@@ -109,4 +109,25 @@ def hive_set_miner_status(enabled=False):
     if miner_screen_running and not enabled:
         run_command(f"/bin/bash -c \"PATH={HIVE_PATH} && sudo /hive/bin/miner stop\"")
     elif enabled and not miner_screen_running:
-        run_command(f"/bin/bash -c \"PATH={HIVE_PATH} && sudo /hive/sbin/nvidia-oc && sudo /hive/bin/miner start\"")
\ No newline at end of file
+        run_command(f"/bin/bash -c \"PATH={HIVE_PATH} && sudo /hive/sbin/nvidia-oc && sudo /hive/bin/miner start\"")
+
+def get_extra_allowed_images():
+    if os.path.exists(config.extra_allowed_images_file):
+        try:
+            with open(config.extra_allowed_images_file, 'r') as file:
+                content = file.read()
+            
+            data = json.loads(content)
+            
+            if isinstance(data, list):
+                if all(isinstance(item, dict) and set(item.keys()) == {'repository', 'allowed_tags'} and isinstance(item['repository'], str) and isinstance(item['allowed_tags'], list) and all(isinstance(tag, str) for tag in item['allowed_tags']) for item in data):
+                    return data
+                else:
+                    return []
+            else:
+                return []
+        except Exception as e:
+            log.error(f"get_extra_allowed_images() | error: {e}")
+            return []
+    else:
+        return []
\ No newline at end of file