Merge pull request #243 from thomas-chu123/feature/CD-1513-ssh-key-en…

…dpoint-re-design

Feature/cd-1513-ssh-key-endpoint-re-design

compute/protocol.py

@@ -87,6 +87,7 @@ class Allocate(bt.Synapse):
     docker_action: dict = {
         "action": "",
         "ssh_key": "",
+        "key_type": "",
     }
 
     def deserialize(self) -> dict:

neurons/Miner/container.py

@@ -432,7 +432,7 @@ 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(new_ssh_key: str, key_type: str = "user"):
     try:
         client, containers = get_docker()
         running_container = None
@@ -443,7 +443,22 @@ def exchange_key_container(new_ssh_key: str):
         if running_container:
             # stop and remove the container by using the SIGTERM signal to PID 1 (init) process in the container
             if running_container.status == "running":
-                running_container.exec_run(cmd=f"bash -c \"echo '{new_ssh_key}' > /root/.ssh/authorized_keys & sync & sleep 1\"")
+                exist_key = running_container.exec_run(cmd="cat /root/.ssh/authorized_keys")
+                exist_key = exist_key.output.decode("utf-8").split("\n")
+                user_key = exist_key[0]
+                terminal_key = ""
+                if len(exist_key) > 1:
+                    terminal_key = exist_key[1]
+                if key_type == "terminal":
+                    terminal_key = new_ssh_key
+                elif key_type == "user":
+                    user_key = new_ssh_key
+                else:
+                    bt.logging.debug("Invalid key type to swap the SSH key")
+                    return {"status": False}
+                key_list = user_key + "\n" + terminal_key
+                # bt.logging.debug(f"New SSH key: {key_list}")
+                running_container.exec_run(cmd=f"bash -c \"echo '{key_list}' > /root/.ssh/authorized_keys & sync & sleep 1\"")
                 running_container.exec_run(cmd="kill -15 1")
                 running_container.wait()
                 running_container.restart()

neurons/miner.py

@@ -438,7 +438,8 @@ 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)
+                    key_type = docker_action["key_type"]
+                    result = exchange_key_container(new_ssh_key, key_type)
                     synapse.output = result
                 elif docker_action["action"] == "restart":
                     public_key = synapse.public_key

neurons/register_api.py

@@ -524,12 +524,12 @@ async def allocate_hotkey(hotkey: str, ssh_key: Optional[str] = None,
                 uuid_key = str(uuid.uuid1())
 
                 private_key, public_key = rsa.generate_key_pair()
-                if ssh_key:
-                    if docker_requirement is None:
-                        docker_requirement = DockerRequirement()
-                        docker_requirement.ssh_key = ssh_key
-                    else:
-                        docker_requirement.ssh_key = ssh_key
+                if docker_requirement is None:
+                    docker_requirement = DockerRequirement()
+                if ssh_key is None:
+                    docker_requirement.ssh_key = ""
+                else:
+                    docker_requirement.ssh_key = ssh_key
 
                 run_start = time.time()
                 result = await run_in_threadpool(self._allocate_container_hotkey, requirements, hotkey,
@@ -1142,7 +1142,7 @@ async def unpause_docker(hotkey: str, uuid_key: str) -> JSONResponse:
                                "description": "An error occurred while exchanging docker key.",
                            },
                        })
-        async def exchange_docker_key(hotkey: str, uuid_key: str, ssh_key: str) -> JSONResponse:
+        async def exchange_docker_key(hotkey: str, uuid_key: str, ssh_key: str, key_type: str = "user") -> JSONResponse:
             # Instantiate the connection to the db
             db = ComputeDb()
             cursor = db.get_cursor()
@@ -1170,13 +1170,14 @@ async def exchange_docker_key(hotkey: str, uuid_key: str, ssh_key: str) -> JSONR
                     docker_action = {
                         "action": "exchange_key",
                         "ssh_key": ssh_key,
+                        "key_type": key_type,
                     }
 
                     if uuid_key_db == uuid_key:
                         index = self.metagraph.hotkeys.index(hotkey)
                         axon = self.metagraph.axons[index]
                         run_start = time.time()
-                        allocate_class = Allocate(timeline=0, device_requirement={}, checking=False, public_key=regkey,
+                        allocate_class = Allocate(timeline=1, device_requirement={}, checking=False, public_key=regkey,
                                                   docker_change=True, docker_action=docker_action)
                         response = await run_in_threadpool(
                             self.dendrite.query, axon, allocate_class, timeout=60