Refactor error handling and logging in '''get_extractor''' - main.py

capa/helpers.py

@@ -9,6 +9,7 @@
 import logging
 import contextlib
 import importlib.util
+import functools
 from typing import NoReturn
 from pathlib import Path
 
@@ -137,22 +138,43 @@ def log_unsupported_format_error():
     logger.error("-" * 80)
 
 
-def log_unsupported_os_error():
-    logger.error("-" * 80)
-    logger.error(" Input file does not appear to target a supported OS.")
-    logger.error(" ")
-    logger.error(
-        " capa currently only supports analyzing executables for some operating systems (including Windows and Linux)."
-    )
-    logger.error("-" * 80)
-
-
-def log_unsupported_arch_error():
-    logger.error("-" * 80)
-    logger.error(" Input file does not appear to target a supported architecture.")
-    logger.error(" ")
-    logger.error(" capa currently only supports analyzing x86 (32- and 64-bit).")
-    logger.error("-" * 80)
+def catch_log_return_errors(func):
+    error_list, return_values, message_list = \
+                            [(UnsupportedFormatError, E_INVALID_FILE_TYPE, 
+                            (" Input file does not appear to be a PE or ELF file.",
+                            " capa currently only supports analyzing PE and ELF files (or shellcode, when using --format sc32|sc64).",
+                            " If you don't know the input file type, you can try using the `file` utility to guess it.")),
+
+                             (UnsupportedArchError, E_INVALID_FILE_ARCH, 
+                              (" Input file does not appear to target a supported architecture.",
+                             " capa currently only supports analyzing x86 (32- and 64-bit).")),
+
+                             (UnsupportedOSError, E_INVALID_FILE_OS, 
+                              (" Input file does not appear to target a supported OS.",
+                             " capa currently only supports analyzing executables for some operating systems (including Windows and Linux)."))]
+
+    @functools.wraps(func)
+    def logging_wrapper(exception):
+        assert(exception in error_list)
+        error_messages = message_list[error_list.index(exception)]
+        error_return_value = return_values[error_list.index(exception)]
+
+        logger.error("-" * 80)
+        logger.error(f"{error_messages[0]}")
+        logger.error(" ")
+
+        for i in error_messages[1:]:
+            logger.error(i)
+
+        logger.error("-" * 80)
+
+        return error_return_value
+
+    if type(func(*args, **kwargs)) = ValueError:
+        return logging_wrapper(func(*args, **kwargs))
+
+    else:
+        return func(*args, **kwargs)
 
 
 def log_unsupported_runtime_error():

capa/main.py

@@ -57,9 +57,8 @@
     get_format,
     get_file_taste,
     get_auto_format,
-    log_unsupported_os_error,
     redirecting_print_to_tqdm,
-    log_unsupported_arch_error,
+    catch_log_return_errors,
     log_unsupported_format_error,
 )
 from capa.exceptions import UnsupportedOSError, UnsupportedArchError, UnsupportedFormatError, UnsupportedRuntimeError
@@ -517,6 +516,7 @@ def get_workspace(path: Path, format_: str, sigpaths: List[Path]):
     return vw
 
 
+@catch_log_return_errors
 def get_extractor(
     path: Path,
     format_: str,
@@ -1257,25 +1257,16 @@ def main(argv: Optional[List[str]] = None):
 
             should_save_workspace = os.environ.get("CAPA_SAVE_WORKSPACE") not in ("0", "no", "NO", "n", None)
 
-            try:
-                extractor = get_extractor(
-                    args.sample,
-                    format_,
-                    args.os,
-                    args.backend,
-                    sig_paths,
-                    should_save_workspace,
-                    disable_progress=args.quiet or args.debug,
-                )
-            except UnsupportedFormatError:
-                log_unsupported_format_error()
-                return E_INVALID_FILE_TYPE
-            except UnsupportedArchError:
-                log_unsupported_arch_error()
-                return E_INVALID_FILE_ARCH
-            except UnsupportedOSError:
-                log_unsupported_os_error()
-                return E_INVALID_FILE_OS
+            # Error checking and logging is performed in the get_extractor call
+            extractor = get_extractor(
+                args.sample,
+                format_,
+                args.os,
+                args.backend,
+                sig_paths,
+                should_save_workspace,
+                disable_progress=args.quiet or args.debug,
+            )
 
         meta = collect_metadata(argv, args.sample, args.format, args.os, args.rules, extractor)