From 0b9e851c804182475e6bea4de666fa0c241bc664 Mon Sep 17 00:00:00 2001
From: Mohaned AbdElMonsef <mo0han3d@gmail.com>
Date: Tue, 1 Oct 2024 14:56:46 +0300
Subject: [PATCH 1/2] Refactor docker actions in miner.py to pass public_key as
 an argument for every docker operation

---
 neurons/miner.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/neurons/miner.py b/neurons/miner.py
index dc114a5e..944de9ec 100644
--- a/neurons/miner.py
+++ b/neurons/miner.py
@@ -447,22 +447,22 @@ def allocate(self, synapse: Allocate) -> Allocate:
                 if docker_action["action"] == "exchange_key":
                     public_key = synapse.public_key
                     new_ssh_key = docker_action["ssh_key"]
-                    result = exchange_key_container(new_ssh_key)
+                    result = exchange_key_container(public_key=public_key, new_ssh_key=new_ssh_key)
                     synapse.output = result
                 elif docker_action["action"] == "restart":
                     public_key = synapse.public_key
-                    result = restart_container()
+                    result = restart_container(public_key=public_key)
                     synapse.output = result
                 elif docker_action["action"] == "pause":
                     public_key = synapse.public_key
-                    result = pause_container()
+                    result = pause_container(public_key=public_key)
                     synapse.output = result
                 elif (
                     docker_action["action"] == "unpause"
                     or docker_action["action"] == "resume"
                 ):
                     public_key = synapse.public_key
-                    result = unpause_container()
+                    result = unpause_container(public_key=public_key)
                     synapse.output = result
                 else:
                     bt.logging.info(f"Unknown action: {docker_action['action']}")

From 9c549a4d872ed438a354f72c2ec8974b5d96cf8d Mon Sep 17 00:00:00 2001
From: Mohaned AbdElMonsef <mo0han3d@gmail.com>
Date: Tue, 1 Oct 2024 14:57:14 +0300
Subject: [PATCH 2/2] Refactor container.py to pass public_key as an argument
 for restart, pause, unpause, and exchange_key_container functions

---
 neurons/Miner/container.py | 36 +++++++++++++++++++++++++++++++-----
 1 file changed, 31 insertions(+), 5 deletions(-)

diff --git a/neurons/Miner/container.py b/neurons/Miner/container.py
index 2d9ed9ad..e94657b0 100644
--- a/neurons/Miner/container.py
+++ b/neurons/Miner/container.py
@@ -42,6 +42,16 @@
 ssh_port = 4444  # Port to map SSH service on the host
 
 
+def get_stored_ssh_key():
+    try:
+        file_path = 'allocation_key'
+        with open(file_path, 'r') as file:
+            allocation_key = file.read()
+        allocation_key = base64.b64decode(allocation_key).decode("utf-8")
+        return allocation_key
+    except Exception as e:
+        bt.logging.info(f"Error getting stored SSH key {e}")
+        return None
 # Initialize Docker client
 def get_docker():
     client = docker.from_env()
@@ -280,8 +290,12 @@ def build_sample_container():
         return {"status": False}
 
 
-def restart_container():
+def restart_container(public_key):
     try:
+        stored_ssh_key = get_stored_ssh_key()
+        if stored_ssh_key.strip() != public_key.strip():
+            bt.logging.info("Error restart the container: Permission denied.")
+            return {"status": False}
         client, containers = get_docker()
         running_container = None
         for container in containers:
@@ -302,8 +316,12 @@ def restart_container():
         bt.logging.info(f"Error restart container {e}")
         return {"status": False}
 
-def pause_container():
+def pause_container(public_key):
     try:
+        stored_ssh_key = get_stored_ssh_key()
+        if stored_ssh_key.strip() != public_key.strip():
+            bt.logging.info("Error pausing the container: Permission denied.")
+            return {"status": False}
         client, containers = get_docker()
         running_container = None
         for container in containers:
@@ -320,8 +338,12 @@ def pause_container():
         bt.logging.info(f"Error pausing container {e}")
         return {"status": False}
 
-def unpause_container():
+def unpause_container(public_key):
     try:
+        stored_ssh_key = get_stored_ssh_key()
+        if stored_ssh_key.strip() != public_key.strip():
+            bt.logging.info("Error unpausing the container: Permission denied.")
+            return {"status": False}
         client, containers = get_docker()
         running_container = None
         for container in containers:
@@ -338,8 +360,12 @@ def unpause_container():
         bt.logging.info(f"Error unpausing container {e}")
         return {"status": False}
 
-def exchange_key_container(new_ssh_key: str):
+def exchange_key_container(public_key,new_ssh_key: str):
     try:
+        stored_ssh_key = get_stored_ssh_key()
+        if stored_ssh_key.strip() != public_key.strip():
+            bt.logging.info("Error exchanging the container key: Permission denied.")
+            return {"status": False}
         client, containers = get_docker()
         running_container = None
         for container in containers:
@@ -359,4 +385,4 @@ def exchange_key_container(new_ssh_key: str):
             return {"status": False}
     except Exception as e:
         bt.logging.info(f"Error changing SSH key on container {e}")
-        return {"status": False}
\ No newline at end of file
+        return {"status": False}