fix: ssh restart command

Author: MSchmoecker

controller/thymis_controller/crud/deployment_info.py

@@ -154,19 +154,15 @@ def get_first_device_host_by_config_id(session: Session, config_id: str) -> str
     return di.reachable_deployed_host if di else None
 
 
-def get_first_by_config_id(
-    session: Session, config_id: str
-) -> db_models.DeploymentInfo | None:
+def get_first_by_config_id(session: Session, config_id: str):
     return (
         session.query(db_models.DeploymentInfo)
         .filter(db_models.DeploymentInfo.deployed_config_id == config_id)
         .first()
     )
 
 
-def get_by_config_id(
-    session: Session, config_id: str
-) -> db_models.DeploymentInfo | None:
+def get_by_config_id(session: Session, config_id: str):
     return (
         session.query(db_models.DeploymentInfo)
         .filter(db_models.DeploymentInfo.deployed_config_id == config_id)

controller/thymis_controller/models/task.py

@@ -207,12 +207,14 @@ class BuildDeviceImageTaskSubmission(BaseModel):
 
 class SSHCommandTaskSubmission(BaseModel):
     type: Literal["ssh_command_task"] = "ssh_command_task"
+    controller_access_client_endpoint: str
+    deployment_info_id: uuid.UUID
+    access_client_token: str
+    deployment_public_key: str
+    ssh_key_path: str
     target_user: str
-    target_host: str
     target_port: int
     command: str
-    ssh_key_path: str
-    ssh_known_hosts_path: str
 
 
 class RunNixOSVMTaskSubmission(BaseModel):

controller/thymis_controller/routers/api_action.py

@@ -1,11 +1,13 @@
 import logging
+import random
 import uuid
 
 from fastapi import APIRouter, Depends, HTTPException, Query, Response
 from fastapi.responses import FileResponse, RedirectResponse
 from thymis_agent import agent
 from thymis_controller import crud, dependencies, models
 from thymis_controller.config import global_settings
+from thymis_controller.crud.agent_token import create_access_client_token
 from thymis_controller.dependencies import (
     DBSessionAD,
     NetworkRelayAD,
@@ -165,22 +167,36 @@ async def restart_device(
     identifier: str,
     db_session: DBSessionAD,
     task_controller: TaskControllerAD,
-    project: ProjectAD,
+    network_relay: NetworkRelayAD,
     user_session_id: UserSessionIDAD,
 ):
-    for target_host in crud.deployment_info.get_by_config_id(db_session, identifier):
-        task_controller.submit(
-            models.SSHCommandTaskSubmission(
-                target_host=target_host.reachable_deployed_host,
-                target_user="root",
-                target_port=22,
-                command="reboot",
-                ssh_key_path=str(global_settings.PROJECT_PATH / "id_thymis"),
-                ssh_known_hosts_path=str(project.known_hosts_path),
-            ),
-            user_session_id=user_session_id,
-            db_session=db_session,
-        )
+    for deployment_info in crud.deployment_info.get_by_config_id(
+        db_session, identifier
+    ):
+        if network_relay.public_key_to_connection_id.get(
+            deployment_info.ssh_public_key
+        ):
+            access_client_token = random.randbytes(32).hex()
+            task = task_controller.submit(
+                models.SSHCommandTaskSubmission(
+                    controller_access_client_endpoint=task_controller.access_client_endpoint,
+                    deployment_info_id=deployment_info.id,
+                    access_client_token=access_client_token,
+                    deployment_public_key=deployment_info.ssh_public_key,
+                    ssh_key_path=str(global_settings.PROJECT_PATH / "id_thymis"),
+                    target_user="root",
+                    target_port=22,
+                    command="reboot",
+                ),
+                user_session_id=user_session_id,
+                db_session=db_session,
+            )
+            create_access_client_token(
+                db_session,
+                deployment_info_id=deployment_info.id,
+                token=access_client_token,
+                deploy_device_task_id=task.id,
+            )
 
 
 @router.head("/download-image")

controller/thymis_controller/task/worker.py

@@ -533,26 +533,49 @@ def ssh_command_task(
     task_data = task.data
     assert task_data.type == "ssh_command_task"
 
-    returncode = run_command(
-        task,
-        conn,
-        process_list,
-        [
-            "ssh",
-            f"-o UserKnownHostsFile={task_data.ssh_known_hosts_path}",
-            "-o StrictHostKeyChecking=yes",
-            "-o ConnectTimeout=10",
-            f"-i {task_data.ssh_key_path}",
-            f"-p {task_data.target_port}",
-            f"{task_data.target_user}@{task_data.target_host}",
-            task_data.command,
-        ],
-    )
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # write deployment_public_key to tmpfile
+        hostfile_path = f"{tmpdir}/known_hosts"
+        with open(hostfile_path, "w", encoding="utf-8") as hostfile:
+            hostfile.write(f"localhost {task_data.deployment_public_key}\n")
+            hostfile.flush()
 
-    if returncode == 0:
-        report_task_finished(task, conn)
-    else:
-        report_task_finished(task, conn, False, "SSH command failed")
+        returncode = run_command(
+            task,
+            conn,
+            process_list,
+            [
+                "ssh",
+                "-i",
+                f"{task_data.ssh_key_path}",
+                "-o",
+                f"UserKnownHostsFile={hostfile.name}",
+                "-o",
+                "StrictHostKeyChecking=yes",
+                "-o",
+                "PasswordAuthentication=no",
+                "-o",
+                "KbdInteractiveAuthentication=no",
+                "-o",
+                "ConnectTimeout=10",
+                "-o",
+                "BatchMode=yes",
+                "-o",
+                f"ProxyCommand={(os.getenv('PYTHONENV')+'/bin/python') if ('PYTHONENV' in os.environ) else 'python' } -m thymis_controller.access_client {task_data.controller_access_client_endpoint} {task_data.deployment_info_id}",
+                f"{task_data.target_user}@localhost",
+                task_data.command,
+            ],
+            env={
+                "PATH": os.getenv("PATH"),
+                "HTTP_NETWORK_RELAY_SECRET": task_data.access_client_token,
+            },
+            cwd=tmpdir,
+        )
+
+        if returncode == 0:
+            report_task_finished(task, conn)
+        else:
+            report_task_finished(task, conn, False, "SSH command failed")
 
 
 SUPPORTED_TASK_TYPES = {