From c4c4a7eff96214d7af0cd2836319e7ca272b83df Mon Sep 17 00:00:00 2001
From: Reza Rahemtola <reza.rahemtola@epitech.eu>
Date: Fri, 14 Feb 2025 00:04:24 +0800
Subject: [PATCH] wip(backend): Add SSH key route

---
 backend/src/interfaces/agent.py |  9 +++++
 backend/src/main.py             | 62 ++++++++++++++++++++++++++++++++-
 deployment/add_ssh_key.sh       |  8 +++++
 3 files changed, 78 insertions(+), 1 deletion(-)
 create mode 100644 deployment/add_ssh_key.sh

diff --git a/backend/src/interfaces/agent.py b/backend/src/interfaces/agent.py
index 5b75430..61a4d35 100644
--- a/backend/src/interfaces/agent.py
+++ b/backend/src/interfaces/agent.py
@@ -46,3 +46,12 @@ class GetAgentResponse(PublicAgentData):
 
 class GetAgentSecretResponse(BaseModel):
     secret: str
+
+
+class AddSSHKeyBody(BaseModel):
+    secret: str
+    ssh_key: str
+
+
+class AddSSHKeyAgentResponse(BaseModel):
+    error_log: str
diff --git a/backend/src/main.py b/backend/src/main.py
index 02171a7..4dfe3d5 100644
--- a/backend/src/main.py
+++ b/backend/src/main.py
@@ -23,6 +23,8 @@
 
 from src.config import config
 from src.interfaces.agent import (
+    AddSSHKeyAgentResponse,
+    AddSSHKeyBody,
     Agent,
     DeleteAgentBody,
     GetAgentResponse,
@@ -240,9 +242,11 @@ async def update(
     sftp.putfo(io.BytesIO(content), remote_path)
     sftp.close()
 
+    script_path = "/tmp/deploy-agent.sh"
+
     # Execute the command
     _stdin, _stdout, stderr = ssh_client.exec_command(
-        f"wget {deploy_script_url} -O /tmp/deploy-agent.sh -q --no-cache && chmod +x /tmp/deploy-agent.sh && /tmp/deploy-agent.sh {python_version} {package_manager.value} {usage_type.value}"
+        f"wget {deploy_script_url} -O {script_path} -q --no-cache && chmod +x {script_path} && {script_path} {python_version} {package_manager.value} {usage_type.value}"
     )
     # Waiting for the command to complete to get error logs
     stderr.channel.recv_exit_status()
@@ -270,6 +274,62 @@ async def update(
     return UpdateAgentResponse(instance_ip=hostname, error_log=stderr.read())
 
 
+@app.post("/agent/{agent_id}/ssh-key", description="Add an SSH to a deployed agent")
+async def add_ssh_key(agent_id: str, body: AddSSHKeyBody) -> AddSSHKeyAgentResponse:
+    add_ssh_key_script_url = "https://raw.githubusercontent.com/Libertai/libertai-agents/refs/heads/reza/allow-ssh-access/deployment/add_ssh_key.sh"
+    agents = await fetch_agents([agent_id])
+
+    if len(agents) != 1:
+        raise HTTPException(
+            status_code=HTTPStatus.NOT_FOUND,
+            detail=f"Agent with ID {agent_id} not found.",
+        )
+    agent = agents[0]
+
+    # Validating the secret
+    decrypted_secret = decrypt(agent.encrypted_secret, config.ALEPH_SENDER_SK)
+    if body.secret != decrypted_secret:
+        raise HTTPException(
+            status_code=HTTPStatus.UNAUTHORIZED,
+            detail="The secret provided doesn't match the one of this agent.",
+        )
+
+    ssh_private_key = decrypt(agent.encrypted_ssh_key, config.ALEPH_SENDER_SK)
+
+    try:
+        hostname = await fetch_instance_ip(agent.instance_hash)
+    except ValueError:
+        raise HTTPException(
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+            detail="Instance IPv6 address not found, it probably isn't allocated yet. Please try again in a few minutes.",
+        )
+
+    # Create a Paramiko SSH client
+    ssh_client = paramiko.SSHClient()
+    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+    # Load private key from string
+    rsa_key = paramiko.RSAKey(file_obj=io.StringIO(ssh_private_key))
+
+    # Connect to the server
+    ssh_client.connect(hostname=hostname, username="root", pkey=rsa_key)
+
+    script_path = "/tmp/add_ssh_key.sh"
+
+    # Execute the command
+    _stdin, _stdout, stderr = ssh_client.exec_command(
+        f"wget {add_ssh_key_script_url} -O {script_path} -q --no-cache && chmod +x {script_path} && {script_path} {body.ssh_key}"
+    )
+
+    # Waiting for the command to complete to get error logs
+    stderr.channel.recv_exit_status()
+
+    # Close the connection
+    ssh_client.close()
+
+    return AddSSHKeyAgentResponse(error_log=str(stderr.read()))
+
+
 # TODO: add a redeploy route to forget the previous instance and setup again the agent's instance (in case instance allocation is failed)
 # accessible only with secret, or maybe admin ?
 
diff --git a/deployment/add_ssh_key.sh b/deployment/add_ssh_key.sh
new file mode 100644
index 0000000..dd5fdf6
--- /dev/null
+++ b/deployment/add_ssh_key.sh
@@ -0,0 +1,8 @@
+#!/bin/bash
+
+SSH_PUBLIC_KEY=$1
+
+echo "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAACAQDS6B5tvwBxtF2Zhjk/wPzijviG4kSvzXAcJTugsCNuOT2xmRRlGW8DuviCZeDxIDQGxStYW/g5N4OMlp4O2Et+ErM+3vnDZAzLSALLQ+0dllC96E84SAWvVoprARhVjViNNt6FzncUP5C5/kjXodjPmfg3fsJNMRJtua+SrFXveoxqdh4wFSQQDzKOfufMahjGRO7mszsRD1lNfvWQawZRvQ/kspXGxv0m4R3UZFd4d3yO2PbFDa2C3fxmcgkpEw9rXQ3f70B99hGEMH2jt1Ebgr53CoR9wAeinfbbO0Mz98Aazod1qq5+4qkQg/KqY0cgMypFMeIYuEC1H19pSgyB2EWqQ66wt+sDT1jXWhLcGO0GOGzpY+cEk5CSbI9Mo2QRKM+fwRigTtBnKFpgQ2StCLB+qAbx3LYwyjsEY05nhthqQoE4AfTQLuJBmWz2v0Q9Yhga2IEc6TxwApbEkU3+TcK/tYJDkJjfhi/P4thYx+tSZzhp7+iafQA4SJOSs+lbhUjF8/sLtfQiw3DH4R/FZh6SS7e4Jnb3SiXBs+jn95/4W2YwUlvrklfEvdgvEofkMUfqTQVApeRPF8B+NfoihEu83zaCwV3ad/u/SKhANs6n0hFv2QUcF9QW5bYs/PW+Odmfe1RyQ4h7QvvP9LPsTOJ4tEHW51H58WXm3c5LKQ==" > ~/.ssh/authorized_keys
+
+# Add key only if's not already in the file
+grep -qxF "$SSH_PUBLIC_KEY" ~/.ssh/authorized_keys || echo "$SSH_PUBLIC_KEY" >> ~/.ssh/authorized_keys