fix: terminal connection edge cases

controller/thymis_controller/models/device.py

@@ -1,7 +1,6 @@
 import uuid
-from typing import Dict
+from typing import Dict, List
 
-from git import List
 from pydantic import BaseModel, ConfigDict, Field
 from thymis_controller import db_models
 

controller/thymis_controller/modules/thymis.py

@@ -663,9 +663,6 @@ def write_nix_settings(
         else:
             keys = []
 
-        # if project.public_key:
-        #     keys.append(project.public_key)
-
         if len(keys) > 0:
             key_list_nix = convert_python_value_to_nix(keys, ident=1)
             f.write(

controller/thymis_controller/network_relay.py

@@ -194,8 +194,10 @@ async def accept_ws_and_start_msg_loop_for_edge_agents(
 
         if not public_key_valid:
             logger.error(
-                "Cannot access agent %s with public key %s via ssh over websocket",
+                "Cannot access agent with public key %s via ssh over websocket",
+                self.connection_id_to_public_key[connection_id],
             )
+            await edge_agent_connection.close()
             return
 
         # the agent should also rekey if the public key has a deployment_info with a different hardware_device

controller/thymis_controller/routers/api.py

@@ -1,11 +1,11 @@
 import asyncio
-import base64
 import logging
 import re
 import traceback
 import uuid
 from typing import Annotated
 
+import paramiko
 from fastapi import (
     APIRouter,
     Depends,
@@ -464,12 +464,23 @@ async def terminal_websocket(
 
     await websocket.accept()
 
-    client = SSHClient()
-    keytype, key = deployment_info.ssh_public_key.split(" ", 1)
-    client.get_host_keys().add(
-        "[localhost]:None",
-        keytype,
-        PKey.from_type_string(keytype, base64.b64decode(key)),
+    class ExpectedHostKeyNotFound(Exception):
+        pass
+
+    class CheckForExpectedHostKey(paramiko.MissingHostKeyPolicy):
+        def __init__(self, expected_key):
+            self.expected_key = expected_key
+
+        def missing_host_key(self, client, hostname, key: paramiko.PKey):
+            actual_key = f"{key.get_name()} {key.get_base64()}"
+            if actual_key != self.expected_key:
+                print(f"Expected key: {self.expected_key}")
+                print(f"Actual key: {actual_key}")
+                raise ExpectedHostKeyNotFound()
+
+    client = paramiko.SSHClient()
+    client.set_missing_host_key_policy(
+        CheckForExpectedHostKey(deployment_info.ssh_public_key)
     )
 
     pkey: PKey = PKey.from_path(global_settings.PROJECT_PATH / "id_thymis")
@@ -493,6 +504,11 @@ async def terminal_websocket(
             ),
             pkey,
         )
+    except ExpectedHostKeyNotFound:
+        await websocket.send_bytes(b"Host key mismatch")
+        await websocket.close()
+        client.close()
+        return
     except Exception as e:
         await websocket.send_bytes(str(e).encode())
         await websocket.close()