[tasks] clean up known hosts through ssh-keygen

DataDog · Jan 7, 2025 · 6454af4 · 6454af4
1 parent d8a0dff
commit 6454af4
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 7 deletions.
diff --git a/tasks/aws/vm.py b/tasks/aws/vm.py
@@ -145,7 +145,7 @@ def destroy_vm(
         stack=stack_name,
     )
     if clean_known_hosts:
-        clean_known_hosts_func(host)
+        clean_known_hosts_func(ctx, host)
 
 
 def _get_os_family(os_family: Optional[str]) -> str:

diff --git a/tasks/azure/vm.py b/tasks/azure/vm.py
@@ -139,7 +139,7 @@ def destroy_vm(
     )
     if clean_known_hosts:
         host = get_host(ctx, remote_hostname, scenario_name, stack_name)
-        clean_known_hosts_func(host)
+        clean_known_hosts_func(ctx, host)
 
 
 def _get_os_information(os_family: Optional[str], arch: Optional[str]) -> Tuple[str, Optional[str]]:

diff --git a/tasks/gcp/vm.py b/tasks/gcp/vm.py
@@ -146,7 +146,7 @@ def destroy_vm(
     )
     if clean_known_hosts:
         host = get_host(ctx, remote_hostname, scenario_name, stack_name)
-        clean_known_hosts_func(host)
+        clean_known_hosts_func(ctx, host)
 
 
 def _get_os_information(os_family: Optional[str], arch: Optional[str]) -> Tuple[str, Optional[str]]:

diff --git a/tasks/localpodman/vm.py b/tasks/localpodman/vm.py
@@ -104,4 +104,4 @@ def destroy_vm(
     )
     if clean_known_hosts:
         host = get_host(ctx, remote_hostname, scenario_name, stack_name)
-        clean_known_hosts_func(host)
+        clean_known_hosts_func(ctx, host)
diff --git a/tasks/tool.py b/tasks/tool.py
@@ -267,20 +267,22 @@ def add_known_host(ctx: Context, host: str) -> None:
     Add the host to the known_hosts file.
     """
     # remove the host if it already exists
-    clean_known_hosts(host)
+    clean_known_hosts(ctx, host)
     home = os.environ.get("HOME", f"/Users/{getpass.getuser()}")
     ctx.run(f"ssh-keyscan -H {host} >> {home}/.ssh/known_hosts")
 
 
-def clean_known_hosts(host: str) -> None:
+def clean_known_hosts(ctx: Context, host: str) -> None:
     """
     Remove the host from the known_hosts file.
     """
+    ctx.run(f"ssh-keygen -R {host}")
     home = os.environ.get("HOME", f"/Users/{getpass.getuser()}")
     with open(f"{home}/.ssh/known_hosts") as f:
         lines = f.readlines()
 
-    filtered_lines = [line for line in lines if not line.startswith(host)]
+    host_comment = f"# {host}"
+    filtered_lines = [line for line in lines if not line.startswith(host_comment)]
     with open(f"{home}/.ssh/known_hosts", "w") as f:
         f.writelines(filtered_lines)