Refactor, handle exceptions

Quit when a step fails

exploitdb/ctfsetup/qemu/buildroot_menuconfig.py

@@ -197,9 +197,9 @@ BR2_HOST_GCC_AT_LEAST_8=y
         """
         start_line = file_lines.index(config_menu_name) + 1
         end_line = file_lines.index("\n", start_line)
-        print(start_line, end_line)
 
         new_file_lines = file_lines[:start_line + 1] + value + file_lines[end_line:]
+        print(f"Patched config from line {start_line} to line {end_line}")
         return new_file_lines
 
     def preconfig_replace(self, config_lines, preconfig_lines):

exploitdb/ctfsetup/qemu/buildroot_menuconfig_navigator.py

@@ -44,8 +44,7 @@ class MenuConfigNavigator:
         if child_node:
             self.current_node = child_node
         else:
-            self.log("Cannot find the desired node. Not moving into any child nodes.")
-            # TODO: quit this script, because this essentially mean that the compilation will fail?
+            raise Exception("Cannot find the desired node. Not selecting any child nodes.")
 
     def select_sibling_node(self, node_prompt_name: str, node_type=None):
         """
@@ -57,8 +56,7 @@ class MenuConfigNavigator:
         if sibling_node:
             self.current_node = sibling_node
         else:
-            self.log("Cannot find the desired node. Not moving into any sibling nodes.")
-            # TODO: quit this script, because this essentially mean that the compilation will fail?
+            raise Exception("Cannot find the desired node. Not selecting any sibling nodes.")
 
     def move_up_one_level(self):
         """
@@ -67,7 +65,7 @@ class MenuConfigNavigator:
         if self.current_node.parent:
             self.current_node = self.current_node.parent
         else:
-            self.log("Cannot find the parent node. Not moving up a level")
+            raise Exception("Cannot find the parent node. Not moving up a level")
 
     def move_up_to_top_level(self):
         """
@@ -76,7 +74,7 @@ class MenuConfigNavigator:
         if self.current_node.parent:
             self.current_node = self.buildroot_config.config_tree.top_node
         else:
-            self.log("Cannot find the top node. Not moving")
+            raise Exception("Cannot find the top node. Not moving")
 
     def input_value(self, value):
         self.current_node.item.set_value(value)

exploitdb/ctfsetup/qemu/buildroot_setup.py

@@ -18,8 +18,9 @@ def main(*args):
     kwargs = parse_args(args)
     try:
         setup(**DEFAULT_CONFIG)
-    except Exception:
-        raise Exception("Something went wrong at setup")
+    except Exception as e:
+        print("Something went wrong at setup")
+        raise e
 
 
 def setup(**kwargs):
@@ -29,8 +30,9 @@ def setup(**kwargs):
     try:
         validate_kwargs(kwargs)
         configure_kernel(kwargs)
-    except Exception:
-        raise Exception("Something went wrong on configuring kernel with Buildroot")
+    except Exception as e:
+        print("Something went wrong on configuring kernel with Buildroot")
+        raise e
 
 
 def needed_kwargs():

exploitdb/ctfsetup/qemu/debian_image_setup.sh

@@ -83,7 +83,6 @@ EOF
   #  mount --bind /proc /mnt/temp/proc
   #  mount --bind /sys /mnt/temp/sys
 
-  # TODO: Need code to check if the script in the image fails.
   cp "${WORKING_DIR}/debian_image_configuration_chroot.sh" mnt/custom_config.sh
   chmod +x mnt/custom_config.sh
 

exploitdb/ctfsetup/qemu/setup.py

@@ -6,6 +6,7 @@ and runs in Qemu
 import os
 import subprocess
 import sys
+import logging
 from shutil import copyfile
 from pathlib import Path
 from datetime import datetime
@@ -34,7 +35,14 @@ def verbose_log(*msgs):
 
 def main(*args):
     kwargs = parse_args(args)
+    try:
+        setup(**kwargs)
+    except Exception:
+        logging.exception("VM setup failed")
+    sys.exit()
 
+
+def setup(**kwargs):
     verbose_log("Validating and fixing input kwargs")
 
     kwargs = derive_input_kwargs(**kwargs)
@@ -58,7 +66,6 @@ def main(*args):
 
     verbose_log("Running with QEMU")
     boot_with_qemu(vm_directory)
-    sys.exit()
 
 
 def parse_args(args):
@@ -98,9 +105,9 @@ def simple_call_shell_command(command):
 def check_dependencies():
     try:
         simple_call_shell_command("bash setup_dependencies.sh")
-    except subprocess.CalledProcessError as e:
-        print("Error in checking dependencies\n", e.output)
-        return None
+    except (FileNotFoundError, subprocess.CalledProcessError) as e:
+        print("Error in checking dependencies")
+        raise e
 
 
 def create_vm_directory():
@@ -115,8 +122,9 @@ def setup_vm_disk_image(vm_directory, **kwargs):
     args = f"{str(vm_directory)} {kwargs['arch']}"
     try:
         simple_call_shell_command("bash debian_image_setup.sh " + args)
-    except subprocess.CalledProcessError as e:
-        raise Exception("Error in download_vm_disk_image()")
+    except (FileNotFoundError, subprocess.CalledProcessError) as e:
+        print("Error in creating VM disk image")
+        raise e
 
 
 def configure_kernel(**kwargs):
@@ -128,7 +136,11 @@ def configure_kernel(**kwargs):
 
     We have to navigate to that dir and call buildroot_setup script
     """
+    try:
         buildroot_setup.setup(**kwargs)
+    except Exception as e:
+        print("Error in kernel configuration")
+        raise e
 
 
 def copy_config_to_buildroot_src(config):
@@ -140,8 +152,8 @@ def compile_kernel():
     try:
         simple_call_shell_command("bash buildroot_compile.sh")
     except subprocess.CalledProcessError as e:
-        print("Error in kernel compilation\n", e.output)
-        return None
+        print("Error in kernel compilation")
+        raise e
 
 
 def copy_kernel_to_vm_directory(vm_directory):

exploitdb/ctfsetup/qemu/utils.py

@@ -51,10 +51,10 @@ def get_gcc_version_by_kernel_version(kernel_version):
                 a = format_gcc_version_for_buildroot(gcc[0])
                 return a
 
-        raise Exception()
+        raise Exception("Could not find gcc version by kernel version")
 
-    except Exception:
-        print("Could not find gcc version by kernel version")
+    except Exception as e:
+        print(e)
         return None
 
 
@@ -94,7 +94,7 @@ def get_gcc_version_table():
 
         return releases
     except Exception as e:
-        print(e.with_traceback())
+        print(e)
         return None
 
 
@@ -112,7 +112,7 @@ def get_kernel_version_table():
 
         return releases
     except Exception as e:
-        print(e.with_traceback())
+        print(e)
         return None