Add get free port util

exploitdb/ctfsetup/qemu/setup.py

@@ -10,7 +10,7 @@ import logging
 from shutil import copyfile
 from pathlib import Path
 from datetime import datetime
-from utils import get_gcc_version_by_kernel_version
+from utils import get_gcc_version_by_kernel_version, get_available_port
 
 import buildroot_setup
 
@@ -33,7 +33,7 @@ def verbose_log(*msgs):
         result = str(msgs[0])
         for msg in msgs[1:]:
             result += " " + str(msg)
-        print("[{}] - {}".format("Setup", result))
+        print("[{} (verbose)] - {}".format("Setup", result))
 
 
 def main(*args):
@@ -67,11 +67,22 @@ def setup(**kwargs):
     compile_kernel()
     copy_kernel_to_vm_directory(vm_directory)
 
+    verbose_log("Booting VM image on QEMU")
+    # Booting up need an available port
+    # Race condition after getting an available port. Need loop.
+
+    for i in range(5):
+        try:
             verbose_log("  Get a port for SSH connection")
             port = get_available_port()
 
-    verbose_log("Running with QEMU", port)
-    boot_with_qemu(vm_directory)
+            verbose_log("  Booting with QEMU", port)
+            boot_with_qemu(vm_directory, port)
+            break
+        except subprocess.CalledProcessError as e:
+            verbose_log("Failed to boot QEMU, trying again with another port...")
+
+    verbose_log("Booted VM image on QEMU")
 
 
 def parse_args(args):
@@ -172,11 +183,6 @@ def copy_kernel_to_vm_directory(vm_directory):
     copyfile(BUILDROOT_OUTPUT_KERNEL_PATH, str(vm_directory) + "/bzImage")
 
 
-def get_available_port():
-    # TODO
-    return 56789
-
-
 def boot_with_qemu(vm_directory, port):
     command = f"cd {str(vm_directory)} && "
     command += f"bash ../../qemu_background_boot.sh -p {port}"

exploitdb/ctfsetup/qemu/utils.py

+import socket
 import subprocess
 import sys
 
@@ -25,6 +26,15 @@ def get_available_linux_dists():
     pass
 
 
+def get_available_port():
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.bind(("", 0))
+    s.listen(1)
+    port = s.getsockname()[1]
+    s.close()
+    return port
+
+
 def run_bash_command(command, verbose=True):
     if verbose:
         subprocess.check_call(command.split(' '), stdout=sys.stdout, stderr=subprocess.STDOUT)