Refactor setup

exploitdb/ctfsetup/qemu/setup.py

@@ -9,7 +9,7 @@ import sys
 from shutil import copyfile
 from pathlib import Path
 from datetime import datetime
-from utils import get_gcc_version_table
+from utils import get_gcc_version_by_kernel_version
 
 import buildroot_setup
 
@@ -35,20 +35,9 @@ def verbose_log(*msgs):
 def main(*args):
     kwargs = parse_args(args)
 
-    verbose_log("Validating input kwargs")
+    verbose_log("Validating and fixing input kwargs")
 
-    if 'arch' not in kwargs:
-        kwargs['arch'] = DEFAULT_ARCH
-        print("Architecture set to default:", DEFAULT_ARCH)
-
-    if 'kernel_version' not in kwargs:
-        kwargs['kernel_version'] = DEFAULT_KERNEL_VERSION
-        print("Kernel set to default:", kwargs['kernel_version'])
-
-    if 'gcc_version' not in kwargs:
-        get_gcc_version_table()
-        kwargs['gcc_version'] = "gcc 4.9.x"
-        print("GCC set to default:", kwargs['gcc_version'])
+    kwargs = derive_input_kwargs(**kwargs)
 
     verbose_log("Checking dependencies")
     check_dependencies()
@@ -57,11 +46,10 @@ def main(*args):
     vm_directory = create_vm_directory()
 
     verbose_log("Setup VM image")
-    if setup_vm_disk_image(vm_directory, DEFAULT_ARCH, DEFAULT_OS) != 0:
-        print("Something went wrong")
+    setup_vm_disk_image(vm_directory, DEFAULT_ARCH, DEFAULT_OS)
 
     verbose_log("Configuring kernel")
-    configure_kernel(kernel_version=DEFAULT_KERNEL_VERSION, arch=DEFAULT_ARCH, gcc_version="gcc 4.9.x")
+    configure_kernel(**kwargs)
     copy_config_to_buildroot_src(BUILDROOT_CUSTOM_CONFIG_PATH)
 
     verbose_log("Compiling kernel")
@@ -69,7 +57,7 @@ def main(*args):
     copy_kernel_to_vm_directory(vm_directory)
 
     verbose_log("Running with QEMU")
-    boot_with_qemu()
+    boot_with_qemu(vm_directory)
     sys.exit()
 
 
@@ -81,6 +69,25 @@ def parse_args(args):
     return args_dict
 
 
+def derive_input_kwargs(**kwargs):
+    if 'arch' not in kwargs:
+        kwargs['arch'] = DEFAULT_ARCH
+        verbose_log("Architecture set to default:", DEFAULT_ARCH)
+
+    if 'kernel_version' not in kwargs:
+        kwargs['kernel_version'] = DEFAULT_KERNEL_VERSION
+        verbose_log("kernel_version set to default:", kwargs['kernel_version'])
+
+    if 'gcc_version' not in kwargs:
+        gcc_version = get_gcc_version_by_kernel_version(kwargs['kernel_version'])
+        if not gcc_version:
+            raise Exception("GCC version cannot be derived from kernel version")
+        kwargs['gcc_version'] = gcc_version
+        verbose_log("gcc_version derived from kernel version:", kwargs['gcc_version'])
+
+    return kwargs
+
+
 def simple_call_shell_command(command):
     """
     Executes a command on shell. Use with caution.
@@ -108,8 +115,7 @@ def setup_vm_disk_image(parent_dir, arch=DEFAULT_ARCH, _os=DEFAULT_OS):
     try:
         simple_call_shell_command("bash debian_image_setup.sh " + str(parent_dir))
     except subprocess.CalledProcessError as e:
-        print("Error in download_vm_disk_image()\n", e.output)
-        return None
+        raise Exception("Error in download_vm_disk_image()")
 
 
 def configure_kernel(**kwargs):
@@ -141,8 +147,9 @@ def copy_kernel_to_vm_directory(vm_directory):
     copyfile(BUILDROOT_OUTPUT_KERNEL_PATH, str(vm_directory) + "/bzImage")
 
 
-def boot_with_qemu():
-    command = """qemu-system-x86_64 \
+def boot_with_qemu(vm_directory):
+    command = f"cd {QEMU_VMS_PATH}/{str(vm_directory)}"
+    command += " && " + """qemu-system-x86_64 \
       -kernel bzImage \
       -append "root=/dev/sda rw biosdevname=0 net.ifnames=0 console=tty0 console=ttyS0,115200 earlyprintk=ttyS0,115200 consoleblank=0" \
       -nographic \

exploitdb/ctfsetup/qemu/utils.py

@@ -39,7 +39,6 @@ def get_gcc_version_by_kernel_version(kernel_version):
 
     formatted_kv = format_kernel_version_for_table_lookup(kernel_version)
     kernel_date = None
-    print(formatted_kv)
     try:
         for kernel in kernel_version_table:
             if kernel[0] == formatted_kv:
@@ -63,9 +62,9 @@ def format_gcc_version_for_buildroot(gcc_version):
     result = ''
     gcc_version_splitted = gcc_version.split('.')
     if int(gcc_version_splitted[0]) <= 4:
-        result = "gcc {}.{}.x".format(gcc_version_splitted[0], gcc_version_splitted[1])
+        result = f"gcc {gcc_version_splitted[0]}.{gcc_version_splitted[1]}.x"
     else:
-        result = "gcc {}.x".format(gcc_version_splitted[0])
+        result = f"gcc {gcc_version_splitted[0]}.x"
     return result