Merge addr.v1.set_public_key with ansible_roles

addr_agent/agent_feature/mapping.py

@@ -6,7 +6,6 @@ from addr_agent.agent_feature.methods.addr_v1_hwreport import RPCAddrV1HWReportM
 from addr_agent.agent_feature.methods.addr_v1_nebula_create_ca import RPCAddrV1NebulaCreateCaMethod
 from addr_agent.agent_feature.methods.addr_v1_nebula_create_host import RPCAddrV1NebulaCreateHostMethod
 from addr_agent.agent_feature.methods.addr_v1_ping import RPCAddrV1PingMethod
-from addr_agent.agent_feature.methods.addr_v1_set_public_key import RPCAddrV1SetPublicKeyMethod
 
 
 class RPCMethod(Enum):
@@ -14,7 +13,6 @@ class RPCMethod(Enum):
     ADDR_V1_HWREPORT = "addr.v1.hwreport"
     ADDR_VI_AGENT_UPDATE = "addr.v1.agent_update"
     ADDR_VI_ANSIBLE_ROLES = "addr.v1.ansible_roles"
-    ADDR_VI_SET_PUBLIC_KEY = "addr.v1.set_public_key"
 
     ADDR_V1_NEBULA_CREATE_CA = "addr.v1.nebula_create_ca"
     ADDR_V1_NEBULA_CREATE_HOST = "addr.v1.nebula_create_host"
@@ -25,7 +23,6 @@ RPC_METHOD_MAPPING = {
     RPCMethod.ADDR_V1_HWREPORT: RPCAddrV1HWReportMethod,
     RPCMethod.ADDR_VI_AGENT_UPDATE: RPCAddrV1AgentUpdateMethod,
     RPCMethod.ADDR_VI_ANSIBLE_ROLES: RPCAddrV1AnsibleRolesMethod,
-    RPCMethod.ADDR_VI_SET_PUBLIC_KEY: RPCAddrV1SetPublicKeyMethod,
 
     RPCMethod.ADDR_V1_NEBULA_CREATE_CA: RPCAddrV1NebulaCreateCaMethod,
     RPCMethod.ADDR_V1_NEBULA_CREATE_HOST: RPCAddrV1NebulaCreateHostMethod,

addr_agent/agent_feature/methods/addr_v1_ansible_roles.py

@@ -13,6 +13,7 @@ from addr_agent.agent_feature.models import RPCMethodImpl, RPCTaskResult, RPCTas
 class RPCAddrV1AnsibleRolesMethod(RPCMethodImpl):
     async def run(self) -> Generator[RPCTaskResult, None, None]:
         request = json.loads(self.request)
+        host_public_key = request.get("host_public_key")
         addr_host = request["meta"]["addr_host"]
         ansible_host = addr_host["ssh_host"]
         ansible_port = addr_host["ssh_port"]
@@ -27,6 +28,24 @@ class RPCAddrV1AnsibleRolesMethod(RPCMethodImpl):
         repos_path = ansible_path / "repos"
         repos_path.mkdir(exist_ok=True)
 
+        if host_public_key:
+            yield RPCTaskResult(
+                status=RPCTaskStatus.PROGRESS,
+                response=f"Update host SSH public key",
+            )
+            authorized_keys = Path("/host/root/.ssh/authorized_keys")
+            keys = set()
+
+            if authorized_keys.exists():
+                existing_keys = authorized_keys.read_text()
+                keys.update(existing_keys.split("\n"))
+            else:
+                authorized_keys.parent.mkdir(parents=True, exist_ok=True)
+
+            keys.add(self.request)
+            keys = [x for x in keys if x]
+            authorized_keys.write_text("\n".join(keys) + "\n")
+
         roles_path = []
         for repo in ansible_repos:
             yield RPCTaskResult(

addr_agent/agent_feature/methods/addr_v1_set_public_key.py

-from pathlib import Path
-from typing import Generator
-
-from addr_agent.agent_feature.models import RPCMethodImpl, RPCTaskResult, RPCTaskStatus
-
-
-class RPCAddrV1SetPublicKeyMethod(RPCMethodImpl):
-    async def run(self, json=None) -> Generator[RPCTaskResult, None, None]:
-        authorized_keys = Path("/host/root/.ssh/authorized_keys")
-        keys = set()
-
-        if authorized_keys.exists():
-            existing_keys = authorized_keys.read_text()
-            keys.update(existing_keys.split("\n"))
-        else:
-            authorized_keys.parent.mkdir(parents=True, exist_ok=True)
-
-        keys.add(self.request)
-        keys = [x for x in keys if x]
-        authorized_keys.write_text("\n".join(keys) + "\n")
-
-        yield RPCTaskResult(status=RPCTaskStatus.COMPLETED)