From 4e1c8ba7eeb534cbd59c0c6b2aad623e49dcb06c Mon Sep 17 00:00:00 2001 From: Adrian Block Date: Thu, 9 Jan 2025 20:29:37 +0100 Subject: [PATCH] implemented backend migration of hostkeys to deployment_info --- .../routes/api/test_hostkey_interfaces.py | 55 ------- ...79a_added_deployment_info_and_hardware_.py | 55 ++++++- controller/thymis_controller/crud/__init__.py | 2 +- .../thymis_controller/crud/deployment_info.py | 23 +++ controller/thymis_controller/crud/hostkey.py | 135 ------------------ .../thymis_controller/db_models/__init__.py | 1 - .../db_models/deployment_info.py | 4 +- .../thymis_controller/db_models/hostkey.py | 14 -- controller/thymis_controller/models/device.py | 11 +- controller/thymis_controller/project.py | 7 +- controller/thymis_controller/routers/agent.py | 31 ---- controller/thymis_controller/routers/api.py | 59 ++------ 12 files changed, 103 insertions(+), 294 deletions(-) delete mode 100644 controller/tests/routes/api/test_hostkey_interfaces.py delete mode 100644 controller/thymis_controller/crud/hostkey.py delete mode 100644 controller/thymis_controller/db_models/hostkey.py diff --git a/controller/tests/routes/api/test_hostkey_interfaces.py b/controller/tests/routes/api/test_hostkey_interfaces.py deleted file mode 100644 index 71c32372..00000000 --- a/controller/tests/routes/api/test_hostkey_interfaces.py +++ /dev/null @@ -1,55 +0,0 @@ -from fastapi.testclient import TestClient -from thymis_controller import crud, models -from thymis_controller.dependencies import get_db_session - -PATH_PREFIX = "/api/hostkey" - - -def test_create_hostkey(test_client): - hostkey = models.CreateHostkeyRequest( - public_key="test_public_key", - device_host="test", - ) - response = test_client.put(f"{PATH_PREFIX}/test_id", json=hostkey.model_dump()) - assert response.status_code == 200 - assert response.json()["public_key"] == "test_public_key" - - hostkey = models.CreateHostkeyRequest( - public_key="test_public_key2", - device_host="test2", - ) - - response = test_client.put(f"{PATH_PREFIX}/test_id", json=hostkey.model_dump()) - assert response.status_code == 200 - assert response.json()["public_key"] == "test_public_key2" - - -def test_get_hostkey_not_found(test_client): - response = test_client.get(f"{PATH_PREFIX}/test_id") - assert response.status_code == 404 - - -def test_get_hostkey(test_client, project): - db_session_func = test_client.app.dependency_overrides.get(get_db_session) - db_session = next(db_session_func()) - - crud.hostkey.create(db_session, "test_id", None, "test_public_key", "test", project) - - response = test_client.get(f"{PATH_PREFIX}/test_id") - assert response.status_code == 200 - assert response.json()["public_key"] == "test_public_key" - - -def test_delete_hostkey(test_client, project): - db_session_func = test_client.app.dependency_overrides.get(get_db_session) - db_session = next(db_session_func()) - - crud.hostkey.create(db_session, "test_id", None, "test_public_key", "test", project) - - response = test_client.delete(f"{PATH_PREFIX}/test_id") - assert response.status_code == 200 - - -def test_delete_hostkey_not_found(test_client): - response = test_client.get(f"{PATH_PREFIX}/test_id") - assert response.status_code == 404 diff --git a/controller/thymis_controller/alembic/versions/b41f88f5679a_added_deployment_info_and_hardware_.py b/controller/thymis_controller/alembic/versions/b41f88f5679a_added_deployment_info_and_hardware_.py index 283a77c8..4dc23293 100644 --- a/controller/thymis_controller/alembic/versions/b41f88f5679a_added_deployment_info_and_hardware_.py +++ b/controller/thymis_controller/alembic/versions/b41f88f5679a_added_deployment_info_and_hardware_.py @@ -5,6 +5,8 @@ Create Date: 2024-12-18 02:05:04.287386 """ +import uuid + import sqlalchemy as sa from alembic import op from sqlalchemy.dialects import sqlite @@ -18,12 +20,12 @@ def upgrade(): # ### commands auto generated by Alembic - please adjust! ### - op.create_table( + deployment_info = op.create_table( "deployment_info", sa.Column("id", sa.Uuid(), nullable=False), sa.Column("ssh_public_key", sa.String(), nullable=False), - sa.Column("deployed_config_commit", sa.String(), nullable=False), - sa.Column("deployed_config_id", sa.String(), nullable=False), + sa.Column("deployed_config_commit", sa.String(), nullable=True), + sa.Column("deployed_config_id", sa.String(), nullable=True), sa.Column("reachable_deployed_host", sa.String(), nullable=True), sa.PrimaryKeyConstraint("id"), sa.UniqueConstraint("ssh_public_key"), @@ -49,8 +51,55 @@ def upgrade(): op.drop_table("images") # ### end Alembic commands ### + conn = op.get_bind() + res = conn.execute( + sa.text("SELECT identifier, public_key, device_host FROM hostkeys") + ) + results = res.fetchall() + hostkey_data = [ + { + "id": uuid.uuid4(), + "deployed_config_id": identifier, + "ssh_public_key": public_key, + "reachable_deployed_host": device_host, + } + for identifier, public_key, device_host in results + ] + op.bulk_insert(deployment_info, hostkey_data) + op.drop_index("ix_hostkeys_identifier", table_name="hostkeys") + op.drop_table("hostkeys") + def downgrade(): + hostkeys = op.create_table( + "hostkeys", + sa.Column("identifier", sa.String(), nullable=False), + sa.Column("build_hash", sa.String(), nullable=True), + sa.Column("public_key", sa.String(), nullable=True), + sa.Column("device_host", sa.String(), nullable=True), + sa.Column("created_at", sa.DateTime(), nullable=True), + sa.PrimaryKeyConstraint("identifier"), + ) + op.create_index( + op.f("ix_hostkeys_identifier"), "hostkeys", ["identifier"], unique=False + ) + + conn = op.get_bind() + res = conn.execute( + sa.text( + "SELECT ssh_public_key, deployed_config_id, reachable_deployed_host FROM deployment_info" + ) + ) + results = res.fetchall() + deployment_info_data = [ + { + "identifier": deployed_config_id, + "public_key": ssh_public_key, + "device_host": reachable_deployed_host, + } + for ssh_public_key, deployed_config_id, reachable_deployed_host in results + ] + op.bulk_insert(hostkeys, deployment_info_data) # ### commands auto generated by Alembic - please adjust! ### op.create_table( "images", diff --git a/controller/thymis_controller/crud/__init__.py b/controller/thymis_controller/crud/__init__.py index 1163361f..de25b7ae 100644 --- a/controller/thymis_controller/crud/__init__.py +++ b/controller/thymis_controller/crud/__init__.py @@ -1 +1 @@ -from . import deployment_info, hardware_device, hostkey, task, web_session +from . import deployment_info, hardware_device, task, web_session diff --git a/controller/thymis_controller/crud/deployment_info.py b/controller/thymis_controller/crud/deployment_info.py index b07b4dae..11b5e0de 100644 --- a/controller/thymis_controller/crud/deployment_info.py +++ b/controller/thymis_controller/crud/deployment_info.py @@ -36,3 +36,26 @@ def check_if_ssh_public_key_exists(session: Session, ssh_public_key: str) -> boo .first() is not None ) + + +def get_all(session: Session): + return session.query(db_models.DeploymentInfo).all() + + +def get_device_host_by_config_id(session: Session, config_id: str) -> str | None: + di = ( + session.query(db_models.DeploymentInfo) + .filter(db_models.DeploymentInfo.deployed_config_id == config_id) + .first() + ) + return di.reachable_deployed_host if di else None + + +def get_by_config_id( + session: Session, config_id: str +) -> db_models.DeploymentInfo | None: + return ( + session.query(db_models.DeploymentInfo) + .filter(db_models.DeploymentInfo.deployed_config_id == config_id) + .first() + ) diff --git a/controller/thymis_controller/crud/hostkey.py b/controller/thymis_controller/crud/hostkey.py deleted file mode 100644 index 67f18741..00000000 --- a/controller/thymis_controller/crud/hostkey.py +++ /dev/null @@ -1,135 +0,0 @@ -import datetime -from typing import TYPE_CHECKING - -from sqlalchemy.orm import Session -from thymis_controller import db_models - -if TYPE_CHECKING: - from thymis_controller.project import Project - - -def create( - db_session: Session, - identifier: str, - build_hash: str, - public_key: str, - device_host: str, - project: "Project", -): - host_key = db_models.HostKey( - identifier=identifier, - build_hash=build_hash, - public_key=public_key, - device_host=device_host, - created_at=datetime.datetime.now(datetime.timezone.utc), - ) - db_session.add(host_key) - db_session.commit() - - # update known hosts - project.update_known_hosts(db_session) - - return host_key - - -def build_hash_is_registered(db_session: Session, build_hash: str): - return ( - db_session.query(db_models.HostKey) - .where(db_models.HostKey.build_hash == build_hash) - .first() - ) is not None - - -def register_device( - db_session: Session, - project: "Project", - build_hash: str, - public_key: str, - device_host: str, -) -> bool: - device = ( - db_session.query(db_models.HostKey) - .where(db_models.HostKey.build_hash == build_hash) - .first() - ) - - if not device: - return False - - device.public_key = public_key - device.device_host = device_host - db_session.commit() - - # update known hosts - project.update_known_hosts(db_session) - - return True - - -def get_all(db_session: Session): - # order by created_at as a workaround, paramiko only matches the first key (unlike OpenSSH) - return ( - db_session.query(db_models.HostKey) - .order_by(db_models.HostKey.created_at.desc()) - .all() - ) - - -def get_device_host(db_session: Session, identifier: str): - device = ( - db_session.query(db_models.HostKey) - .where(db_models.HostKey.identifier == identifier) - .first() - ) - return device and device.device_host or None - - -def get_by_public_key(db_session: Session, public_key: str): - return ( - db_session.query(db_models.HostKey) - .where(db_models.HostKey.public_key == public_key) - .first() - ) - - -def get_by_build_hash(db_session: Session, build_hash: str): - return ( - db_session.query(db_models.HostKey) - .where(db_models.HostKey.build_hash == build_hash) - .order_by(db_models.HostKey.created_at.desc()) - .first() - ) - - -def rename_device(db_session: Session, old_identifier: str, new_identifier: str): - device = ( - db_session.query(db_models.HostKey) - .where(db_models.HostKey.identifier == old_identifier) - .first() - ) - device.identifier = new_identifier - db_session.commit() - return device - - -def has_device(db_session: Session, identifier: str): - return ( - db_session.query(db_models.HostKey) - .where(db_models.HostKey.identifier == identifier) - .first() - ) is not None - - -def get_by_identifier(db_session: Session, identifier: str): - return ( - db_session.query(db_models.HostKey) - .where(db_models.HostKey.identifier == identifier) - .first() - ) - - -def delete(db_session: Session, identifier: str): - db_session.query(db_models.HostKey).where( - db_models.HostKey.identifier == identifier - ).delete() - db_session.commit() diff --git a/controller/thymis_controller/db_models/__init__.py b/controller/thymis_controller/db_models/__init__.py index 010a3c15..1269488e 100644 --- a/controller/thymis_controller/db_models/__init__.py +++ b/controller/thymis_controller/db_models/__init__.py @@ -1,5 +1,4 @@ from .deployment_info import DeploymentInfo from .hardware_device import HardwareDevice -from .hostkey import HostKey from .task import Task from .web_session import WebSession diff --git a/controller/thymis_controller/db_models/deployment_info.py b/controller/thymis_controller/db_models/deployment_info.py index 471af7ef..6f47291a 100644 --- a/controller/thymis_controller/db_models/deployment_info.py +++ b/controller/thymis_controller/db_models/deployment_info.py @@ -17,8 +17,8 @@ class DeploymentInfo(Base): ) ssh_public_key: Mapped[str] = mapped_column(nullable=False, unique=True) - deployed_config_commit: Mapped[str] - deployed_config_id: Mapped[str] + deployed_config_commit: Mapped[str] = mapped_column(nullable=True) + deployed_config_id: Mapped[str] = mapped_column(nullable=True) reachable_deployed_host: Mapped[str | None] diff --git a/controller/thymis_controller/db_models/hostkey.py b/controller/thymis_controller/db_models/hostkey.py deleted file mode 100644 index 49fb1d6c..00000000 --- a/controller/thymis_controller/db_models/hostkey.py +++ /dev/null @@ -1,14 +0,0 @@ -import datetime - -from sqlalchemy import Column, DateTime, Integer, String -from thymis_controller.database.base import Base - - -class HostKey(Base): - __tablename__ = "hostkeys" - - identifier = Column(String, primary_key=True, index=True) - build_hash = Column(String) - public_key = Column(String) - device_host = Column(String) # working ip address - created_at = Column(DateTime, default=datetime.datetime.now(datetime.timezone.utc)) diff --git a/controller/thymis_controller/models/device.py b/controller/thymis_controller/models/device.py index ea6ab90a..80370b0f 100644 --- a/controller/thymis_controller/models/device.py +++ b/controller/thymis_controller/models/device.py @@ -2,11 +2,12 @@ from git import List from pydantic import BaseModel, Field +from thymis_controller import db_models class DeviceNotifyRequest(BaseModel): commit_hash: str | None - config_id: str + config_id: str | None hardware_ids: Dict[str, str | None] public_key: str ip_addresses: List[str] @@ -22,6 +23,14 @@ class Hostkey(BaseModel): public_key: str device_host: str + @staticmethod + def from_deployment_info(deployment_info: db_models.DeploymentInfo) -> "Hostkey": + return Hostkey( + identifier=deployment_info.deployed_config_id, + public_key=deployment_info.ssh_public_key, + device_host=deployment_info.reachable_deployed_host, + ) + class CreateHostkeyRequest(BaseModel): public_key: str diff --git a/controller/thymis_controller/project.py b/controller/thymis_controller/project.py index dc6320be..760ffbde 100644 --- a/controller/thymis_controller/project.py +++ b/controller/thymis_controller/project.py @@ -283,10 +283,11 @@ def update_known_hosts(self, db_session: sqlalchemy.orm.Session): tempfile.NamedTemporaryFile(delete=False).name ) - hostkeys = crud.hostkey.get_all(db_session) + deployment_infos = crud.deployment_info.get_all(db_session) with open(self.known_hosts_path, "w", encoding="utf-8") as f: - for hostkey in hostkeys: - f.write(f"{hostkey.device_host} {hostkey.public_key}\n") + for di in deployment_infos: + if di.reachable_deployed_host and di.ssh_public_key: + f.write(f"{di.reachable_deployed_host} {di.ssh_public_key}\n") logger.debug("Updated known_hosts file at %s", self.known_hosts_path) diff --git a/controller/thymis_controller/routers/agent.py b/controller/thymis_controller/routers/agent.py index 7da11201..d4f667e0 100644 --- a/controller/thymis_controller/routers/agent.py +++ b/controller/thymis_controller/routers/agent.py @@ -71,34 +71,3 @@ def device_notify( # check if device is registered in hardware_device table crud.hardware_device.create_or_update(db_session, hardware_id, deployment_info.id) - - -@router.post("/heartbeat") -def heartbeat( - heartbeat: models.DeviceHeartbeatRequest, - db_session: SessionAD, - request: Request, -): - # check if device is registered - device = crud.hostkey.get_by_public_key(db_session, heartbeat.public_key) - if not device: - logging.info(f"Device with public key {heartbeat.public_key} is not registered") - raise HTTPException(status_code=404, detail="Your device is not registered") - - logging.debug(f"Device with identifier {device.identifier} sends heartbeat") - # check for reachable device - device_host = determine_first_host_with_key( - hosts=[request.client.host, *heartbeat.ip_addresses], - public_key=heartbeat.public_key, - ) - - if not device_host: - logging.error(f"Device with identifier {device.identifier} is not reachable") - raise HTTPException(status_code=400, detail="Your device is not reachable") - - if device.device_host != device_host: - logging.info( - f"Device with identifier {device.identifier} has new host {device_host}" - ) - device.device_host = device_host - db_session.commit() diff --git a/controller/thymis_controller/routers/api.py b/controller/thymis_controller/routers/api.py index 2d9e2812..4e6b574b 100644 --- a/controller/thymis_controller/routers/api.py +++ b/controller/thymis_controller/routers/api.py @@ -75,10 +75,8 @@ async def deploy( project.commit(summary) registered_devices = [] - for device in crud.hostkey.get_all(session): - registered_devices.append( - models.Hostkey.model_validate(device, from_attributes=True) - ) + for device in crud.deployment_info.get_all(session): + registered_devices.append(models.Hostkey.from_deployment_info(device)) # runs a nix command to deploy the flake await project.create_deploy_project_task(registered_devices) @@ -111,8 +109,6 @@ async def device_and_build_download_image_for_clone( x += 1 new_identifier = device_name(x) - if crud.hostkey.has_device(db_session, identifier): - crud.hostkey.rename_device(db_session, identifier, new_identifier) project.clone_state_device(identifier, new_identifier, lambda n: f"{n}-{x}") await project.create_build_device_image_task(new_identifier, db_session) @@ -125,7 +121,9 @@ async def restart_device( state: State = Depends(dependencies.get_state), ): device = next(device for device in state.devices if device.identifier == identifier) - target_host = crud.hostkey.get_device_host(db_session, identifier) + target_host = crud.deployment_info.get_device_host_by_config_id( + db_session, identifier + ) await project.create_restart_device_task(device, target_host) @@ -209,7 +207,9 @@ async def vnc_websocket( state: State = Depends(dependencies.get_state), ): device = next(device for device in state.devices if device.identifier == identifier) - target_host = crud.hostkey.get_device_host(db_session, identifier) + target_host = crud.deployment_info.get_device_host_by_config_id( + db_session, identifier + ) if device is None or target_host is None: await websocket.close() @@ -233,45 +233,6 @@ async def vnc_websocket( ws_to_tcp_task.cancel() -@router.get("/hostkey/{identifier}", response_model=models.Hostkey) -def get_hostkey(db_session: SessionAD, identifier: str): - """ - Get the hostkey for a device - """ - hostkey = crud.hostkey.get_by_identifier(db_session, identifier) - if not hostkey: - raise HTTPException(status_code=404, detail="Hostkey not found") - return hostkey - - -@router.put("/hostkey/{identifier}", response_model=models.Hostkey) -def create_hostkey( - identifier: str, - hostkey: models.CreateHostkeyRequest, - db_session: SessionAD, - project: ProjectAD, -): - """ - Create a hostkey for a device - """ - if crud.hostkey.get_by_identifier(db_session, identifier): - crud.hostkey.delete(db_session, identifier) - - return crud.hostkey.create( - db_session, identifier, None, hostkey.public_key, hostkey.device_host, project - ) - - -@router.delete("/hostkey/{identifier}") -def delete_hostkey(db_session: SessionAD, identifier: str): - """ - Delete the hostkey for a device - """ - if not crud.hostkey.get_by_identifier(db_session, identifier): - raise HTTPException(status_code=404, detail="Hostkey not found") - crud.hostkey.delete(db_session, identifier) - - HOST_PATTERN = r"^([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-]{0,61}[a-zA-Z0-9])(\.([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-]{0,61}[a-zA-Z0-9]))*$" @@ -307,7 +268,9 @@ async def terminal_websocket( state: State = Depends(dependencies.get_state), ): device = next(device for device in state.devices if device.identifier == identifier) - target_host = crud.hostkey.get_device_host(db_session, identifier) + target_host = crud.deployment_info.get_device_host_by_config_id( + db_session, identifier + ) if device is None or target_host is None: await websocket.close()