Fix ssh not reconnenting in kata challenges

sshd/enter.py

@@ -6,7 +6,7 @@
 import shlex
 import sys
 import time
-
+import signal
 import docker
 import redis
 
@@ -86,8 +86,8 @@ def print(*args, **kwargs):
 
         attempts = 0
         print("\r", " " * 80, "\rConnected!")
-
-        if not os.fork():
+        child_pid = os.fork();
+        if not child_pid:
             ssh_entrypoint = "/run/dojo/bin/ssh-entrypoint"
             if is_mac:
                 cmd = f"/bin/bash -c {shlex.quote(original_command)}" if original_command  else "zsh -i"
@@ -114,7 +114,22 @@ def print(*args, **kwargs):
                 )
 
         else:
-            _, status = os.wait()
+            runtime = (container.attrs or {}).get("HostConfig",{}).get("Runtime")
+            is_kata = runtime == "io.containerd.run.kata.v2"
+            if is_kata:
+                if not os.fork():
+                    try:
+                        container.wait(condition="not-running")
+                        try:
+                            os.kill(child_pid, signal.SIGTERM)
+                            time.sleep(0.5)
+                        except ProcessLockupError:
+                            pass
+                    except:
+                        pass
+                    finally:
+                        os._exit(0)  
+            _, status = os.waitpid(child_pid,0)
             if simple or status == 0:
                 break
             print()