build: T8194: use an exception to simplify exception handling

1 files changed, 105 insertions, 111 deletions

diff --git a/scripts/image-build/build-vyos-image b/scripts/image-build/build-vyos-image
index 6919b003..bb92ce4e 100755
--- a/scripts/image-build/build-vyos-image
+++ b/scripts/image-build/build-vyos-image
@@ -32,96 +32,9 @@ import datetime
 import functools
 import string
 
-## Check if the script is running wirh root permissions
-## Since live-build requires privileged calls such as chroot(),
-## there's no real way around it.
-if os.getuid() != 0:
-    print("E: this script requires root privileges")
-    sys.exit(1)
-
-# Import third-party modules
-try:
-    import tomli
-    import jinja2
-    import git
-except ModuleNotFoundError as e:
-    print(f"E: Cannot load required library {e}")
-    print("E: Please make sure the following Python3 modules are installed: tomli jinja2 git")
-    sys.exit(1)
-
-# Import local defaults
-import defaults
-
-## Load the file with default build configuration options
-try:
-    with open(defaults.DEFAULTS_FILE, 'rb') as f:
-        build_defaults = tomli.load(f)
-except Exception as e:
-    print("E: Failed to open the defaults file {0}: {1}".format(defaults.DEFAULTS_FILE, e))
-    sys.exit(1)
-
-# Checkout vyos-1x under build directory
-try:
-    branch_name = build_defaults['vyos_branch']
-    url_vyos_1x = os.getenv('VYOS1X_REPO_URL', default='https://github.com/vyos/vyos-1x')
-    path_vyos_1x = os.path.join(defaults.BUILD_DIR, 'vyos-1x')
-    try:
-        repo_vyos_1x = git.Repo.clone_from(url_vyos_1x, path_vyos_1x, no_checkout=True)
-    except git.GitCommandError:
-        if os.path.exists(path_vyos_1x):
-            try:
-                repo_vyos_1x = git.Repo(path_vyos_1x)
-            except git.GitError:
-                print(f'E: Corrupted vyos-1x git repo: {path_vyos_1x}; remove')
-                sys.exit(1)
-        else:
-            raise
-    # alternatively, pass commit hash or tag as arg:
-    repo_vyos_1x.git.checkout(branch_name)
-except Exception as e:
-    print(f'E: Could not retrieve vyos-1x from branch {branch_name}: {repr(e)}')
-    sys.exit(1)
-
-# Add the vyos-1x directory to the Python path so that
-# we can import modules from it.
-VYOS1X_DIR = os.path.join(os.getcwd(), defaults.BUILD_DIR, 'vyos-1x/python')
-if not os.path.exists(VYOS1X_DIR):
-    print("E: vyos-1x subdirectory does not exist, did git checkout fail?")
-    sys.exit(1)
-else:
-    sys.path.append(VYOS1X_DIR)
-
-# Import local modules from scripts/image-build
-# They rely on modules from vyos-1x
-import utils
-import raw_image
-
-from utils import cmd, rc_cmd
-
-## Check if there are missing build dependencies
-deps = {
-    'packages': [
-       'sudo',
-       'make',
-       'live-build',
-       'pbuilder',
-       'devscripts',
-       'python3-pystache',
-       'python3-git',
-       'qemu-utils',
-       'gdisk',
-       'kpartx',
-       'dosfstools'
-    ],
-   'binaries': []
-}
-
-print("I: Checking if packages required for VyOS image build are installed")
-try:
-    utils.check_system_dependencies(deps)
-except OSError as e:
-    print(f"E: {e}")
-    sys.exit(1)
+class ImageBuildError(Exception):
+    pass
+
 
 # argparse converts hyphens to underscores,
 # so for lookups in the original options hash we have to convert them back
@@ -167,7 +80,91 @@ def make_toml_path(dir, file_basename):
     return os.path.join(dir, file_basename + ".toml")
 
 
-if __name__ == "__main__":
+def build():
+    ## Check if the script is running wirh root permissions
+    ## Since live-build requires privileged calls such as chroot(),
+    ## there's no real way around it.
+    if os.getuid() != 0:
+        raise ImageBuildError('This script requires root privileges')
+
+    # Import third-party modules
+    try:
+        import tomli
+        import jinja2
+        import git
+    except ModuleNotFoundError as e:
+        raise ImageBuildError(f'Cannot load required library {e}. ' \
+                'Please make sure the following Python3 modules are installed: tomli jinja2 git')
+
+    # Import local defaults
+    import defaults
+
+    ## Load the file with default build configuration options
+    try:
+        with open(defaults.DEFAULTS_FILE, 'rb') as f:
+            build_defaults = tomli.load(f)
+    except Exception as e:
+        raise ImageBuildError(f'Failed to open the defaults file {defaults.DEFAULTS_FILE}: {e}')
+
+    # Checkout vyos-1x under build directory
+    try:
+        branch_name = build_defaults['vyos_branch']
+        url_vyos_1x = os.getenv('VYOS1X_REPO_URL', default='https://github.com/vyos/vyos-1x')
+        path_vyos_1x = os.path.join(defaults.BUILD_DIR, 'vyos-1x')
+        try:
+            repo_vyos_1x = git.Repo.clone_from(url_vyos_1x, path_vyos_1x, no_checkout=True)
+        except git.GitCommandError:
+            if os.path.exists(path_vyos_1x):
+                try:
+                    repo_vyos_1x = git.Repo(path_vyos_1x)
+                except git.GitError:
+                    raise ImageBuildError(f'Corrupted vyos-1x git repo: {path_vyos_1x}; remove')
+            else:
+                raise
+        # alternatively, pass commit hash or tag as arg:
+        repo_vyos_1x.git.checkout(branch_name)
+    except Exception as e:
+        raise ImageBuildError(f'Could not retrieve vyos-1x from branch {branch_name}: {repr(e)}')
+
+    # Add the vyos-1x directory to the Python path so that
+    # we can import modules from it.
+    VYOS1X_DIR = os.path.join(os.getcwd(), defaults.BUILD_DIR, 'vyos-1x/python')
+    if not os.path.exists(VYOS1X_DIR):
+        raise ImageBuildError('vyos-1x subdirectory does not exist, did git checkout fail?')
+    else:
+        sys.path.append(VYOS1X_DIR)
+
+    # Import local modules from scripts/image-build
+    # They rely on modules from vyos-1x
+    import utils
+    import raw_image
+
+    from utils import cmd, rc_cmd
+
+    ## Check if there are missing build dependencies
+    deps = {
+        'packages': [
+           'sudo',
+           'make',
+           'live-build',
+           'pbuilder',
+           'devscripts',
+           'python3-pystache',
+           'python3-git',
+           'qemu-utils',
+           'gdisk',
+           'kpartx',
+           'dosfstools'
+        ],
+       'binaries': []
+    }
+
+    print("I: Checking if packages required for VyOS image build are installed")
+    try:
+        utils.check_system_dependencies(deps)
+    except OSError as e:
+        raise ImageBuildError(f'{e}')
+
     ## Get a list of available build flavors
     flavor_dir_env = os.getenv("VYOS_BUILD_FLAVORS_DIR")
     if flavor_dir_env:
@@ -230,15 +227,12 @@ if __name__ == "__main__":
         func = get_validator(options, k)
         if func is not None:
             if not func(v):
-                print("E: {v} is not a valid value for --{o} option".format(o=key, v=v))
-                sys.exit(1)
+                raise ImageBuildError(f'{v} is not a valid value for --{key} option')
 
     if not args["build_flavor"]:
-        print("E: Build flavor is not specified!")
-        print("E: For example, to build the generic ISO, run {} iso".format(sys.argv[0]))
-        print("Available build flavors:\n")
-        print("\n".join(build_flavors))
-        sys.exit(1)
+        raise ImageBuildError("Build flavor is not specified!\n" \
+                              'Available build flavors: ' +
+                              ', '.join(build_flavors))
 
     ## Try to get correct architecture and build type from build flavor and CLI arguments
     pre_build_config = copy.deepcopy(build_defaults)
@@ -251,11 +245,9 @@ if __name__ == "__main__":
             flavor_config = tomli.load(f)
             pre_build_config = merge_defaults(flavor_config, defaults=pre_build_config)
     except FileNotFoundError:
-        print(f"E: Flavor '{build_flavor}' does not exist")
-        sys.exit(1)
+        raise ImageBuildError(f'Build flavor "{build_flavor}" does not exist!')
     except tomli.TOMLDecodeError as e:
-        print(f"E: Failed to parse TOML file for flavor '{build_flavor}': {e}")
-        sys.exit(1)
+        raise ImageBuildError(f'Failed to parse TOML file for flavor "{build_flavor}": {e}')
 
     ## Combine configs args > flavor > defaults
     pre_build_config = merge_defaults(args, defaults=pre_build_config, skip_none=True)
@@ -265,8 +257,7 @@ if __name__ == "__main__":
     # but --pbuilder-debian-mirror or --debian-security-mirror are not,
     # use the --debian-mirror value for those
     if pre_build_config['debian_mirror'] is None:
-        print("E: debian_mirror must be specified")
-        sys.exit(1)
+        raise ImageBuildError('debian_mirror must be specified')
 
     if pre_build_config['pbuilder_debian_mirror'] is None:
         pre_build_config['pbuilder_debian_mirror'] = pre_build_config['debian_mirror']
@@ -278,8 +269,7 @@ if __name__ == "__main__":
     if args.get('version'):
         allowed = string.ascii_letters + string.digits + '.' + '-' + '+'
         if not set(args['version']) <= set(allowed):
-            print(f'Version string contains illegal character(s), allowed: {allowed}')
-            sys.exit(1)
+            raise ImageBuildError(f'Version string contains illegal character(s), allowed: {allowed}')
 
     ## Inject some useful hardcoded options
     args['build_dir'] = defaults.BUILD_DIR
@@ -326,8 +316,7 @@ if __name__ == "__main__":
     ## Check if image format is specified,
     ## else we have no idea what we are actually supposed to build.
     if not has_nonempty_key(build_config, "image_format"):
-        print("E: image format is not specified in the build flavor file")
-        sys.exit(1)
+        raise ImageBuildError('Image format is not specified in the build flavor file')
 
     ## Override bootloaders if specified
     if args['bootloaders'] is not None:
@@ -357,7 +346,6 @@ if __name__ == "__main__":
 
     ## Dump the complete config if the user enabled debug mode
     if debug:
-        import json
         print("D: Effective build config:\n")
         print(json.dumps(build_config, indent=4))
 
@@ -538,8 +526,7 @@ DOCUMENTATION_URL="{build_config['documentation_url']}"
                 components = repo_data.get('components', 'main')
 
                 if not url:
-                    print(f'E: repository {r} does not specify URL')
-                    sys.exit(1)
+                    raise ImageBuildError(f'Repository {r} does not specify URL')
 
                 if arch:
                     arch_string = f'[arch={arch}]'
@@ -744,3 +731,10 @@ Pin-Priority: 600
 
     with open('manifest.json', 'w') as f:
         f.write(json.dumps(manifest))
+
+if __name__ == '__main__':
+    try:
+        build()
+    except ImageBuildError as e:
+        print(f"E: {e}")
+        sys.exit(1)