core: improve code quality

src/fuzz_introspector/cli.py

@@ -23,10 +23,12 @@
 sys.setrecursionlimit(10000)
 
 logger = logging.getLogger(name=__name__)
-LOG_FMT = '%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s'
+LOG_FMT = ('%(asctime)s.%(msecs)03d %(levelname)s '
+           '%(module)s - %(funcName)s: %(message)s')
 
 
 def get_cmdline_parser() -> argparse.ArgumentParser:
+    """Parse the commandline"""
     parser = argparse.ArgumentParser()
 
     subparsers = parser.add_subparsers(dest='command')
@@ -49,7 +51,7 @@ def get_cmdline_parser() -> argparse.ArgumentParser:
     full_parser.add_argument('--language',
                              type=str,
                              help='Programming of the source code to analyse.',
-                             choices=constants.LANGUAGES)
+                             choices=constants.LANGUAGES_SUPPORTED)
     full_parser.add_argument('--out-dir',
                              default='',
                              type=str,
@@ -136,20 +138,13 @@ def get_cmdline_parser() -> argparse.ArgumentParser:
 
 
 def set_logging_level() -> None:
-    if os.environ.get("FUZZ_LOGLEVEL"):
-        level = os.environ.get("FUZZ_LOGLEVEL")
-        if level == "debug":
-            logging.basicConfig(
-                level=logging.DEBUG,
-                format=LOG_FMT,
-                datefmt='%Y-%m-%d %H:%M:%S',
-            )
-        else:
-            logging.basicConfig(
-                level=logging.INFO,
-                format=LOG_FMT,
-                datefmt='%Y-%m-%d %H:%M:%S',
-            )
+    """Sets logging level."""
+    if os.environ.get('FUZZ_LOGLEVEL', 'info') == 'debug':
+        logging.basicConfig(
+            level=logging.DEBUG,
+            format=LOG_FMT,
+            datefmt='%Y-%m-%d %H:%M:%S',
+        )
     else:
         logging.basicConfig(
             level=logging.INFO,
@@ -160,6 +155,7 @@ def set_logging_level() -> None:
 
 
 def main() -> int:
+    """Main CLI entrypoint."""
     set_logging_level()
 
     parser = get_cmdline_parser()

src/fuzz_introspector/commands.py

@@ -54,7 +54,7 @@ def end_to_end(args) -> int:
     else:
         out_dir = os.getcwd()
 
-    if args.language == 'jvm':
+    if args.language == constants.LANGUAGES.JAVA:
         entrypoint = 'fuzzerTestOneInput'
     else:
         entrypoint = 'LLVMFuzzerTestOneInput'

src/fuzz_introspector/constants.py

@@ -31,7 +31,27 @@
 
 SAVED_SOURCE_FOLDER = 'source-code'
 
-LANGUAGES = ['c', 'c++', 'jvm', 'go', 'rust']
+
+class LANGUAGES:
+    C = 'c'
+    CPP = 'c++'
+    JAVA = 'jvm'
+    GO = 'go'
+    RUST = 'rust'
+
+
+LANGUAGES_SUPPORTED = [
+    LANGUAGES.C, LANGUAGES.CPP, LANGUAGES.JAVA, LANGUAGES.GO, LANGUAGES.RUST
+]
+
+LANGUAGE_EXTENSIONS = {
+    LANGUAGES.C: ['.c', '.h'],
+    LANGUAGES.CPP:
+    ['.c', '.cpp', '.cc', '.c++', '.cxx', '.h', '.hpp', '.hh', '.hxx'],
+    LANGUAGES.JAVA: ['.java'],
+    LANGUAGES.RUST: ['.rs'],
+    LANGUAGES.GO: ['.go', '.cgo'],
+}
 
 # Holds data about all functions in javascript, to ease loading of static
 # website.

src/fuzz_introspector/frontends/frontend_cpp.py

@@ -26,7 +26,6 @@
 from fuzz_introspector.frontends.datatypes import Project
 
 logger = logging.getLogger(name=__name__)
-LOG_FMT = '%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s'
 
 
 class SourceCodeFile():

src/fuzz_introspector/frontends/frontend_rust.py

@@ -26,7 +26,6 @@
 from fuzz_introspector.frontends.datatypes import Project
 
 logger = logging.getLogger(name=__name__)
-LOG_FMT = '%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s'
 
 
 class SourceCodeFile():

src/fuzz_introspector/frontends/oss_fuzz.py

@@ -15,7 +15,6 @@
 ################################################################################
 
 import os
-import argparse
 import pathlib
 import logging
 
@@ -26,55 +25,22 @@
 from fuzz_introspector.frontends import frontend_rust
 from fuzz_introspector.frontends.datatypes import Project
 
+from fuzz_introspector import constants
+
 logger = logging.getLogger(name=__name__)
-LOG_FMT = '%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s'
-
-LANGUAGE_EXTENSION_MAP = {
-    'c': ['.c', '.h'],
-    'c++':
-    ['.c', '.cpp', '.cc', '.c++', '.cxx', '.h', '.hpp', '.hh', '.hxx', '.inl'],
-    'cpp':
-    ['.c', '.cpp', '.cc', '.c++', '.cxx', '.h', '.hpp', '.hh', '.hxx', '.inl'],
-    'go': ['.go', '.cgo'],
-    'jvm': ['.java'],
-    'rust': ['.rs'],
-}
 
 EXCLUDE_DIRECTORIES = [
     'node_modules', 'aflplusplus', 'honggfuzz', 'inspector', 'libfuzzer',
     'fuzztest', 'target', 'build'
 ]
 
 
-def setup_logging():
-    """Initializes logging"""
-    logging.basicConfig(
-        level=logging.INFO,
-        format=LOG_FMT,
-        datefmt='%Y-%m-%d %H:%M:%S',
-    )
-
-
-def parse_args():
-    """Parse command line arguments."""
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument('--target-dir',
-                        help='Directory of which do analysis',
-                        required=True)
-    parser.add_argument('--entrypoint', help='Entrypoint for the calltree')
-    parser.add_argument('--language',
-                        help='Language of target project',
-                        required=True)
-
-    return parser.parse_args()
-
-
 def capture_source_files_in_tree(directory_tree: str,
                                  language: str) -> list[str]:
     """Captures source code files in a given directory."""
     language_files = []
-    language_extensions = LANGUAGE_EXTENSION_MAP.get(language.lower(), [])
+    language_extensions = constants.LANGUAGE_EXTENSIONS.get(
+        language.lower(), [])
 
     for dirpath, _, filenames in os.walk(directory_tree):
         # Skip some non project directories
@@ -242,7 +208,7 @@ def analyse_folder(language: str = '',
     # Extract source files for target language
     source_files = capture_source_files_in_tree(directory, language)
 
-    if language == 'c':
+    if language == constants.LANGUAGES.C:
         project = process_c_project(directory,
                                     entrypoint,
                                     out,
@@ -251,26 +217,26 @@ def analyse_folder(language: str = '',
                                     dump_output=dump_output)
     else:
         # Process for different language
-        if language.lower() in ['cpp', 'c++']:
+        if language == constants.LANGUAGES.CPP:
             project = process_cpp_project(entrypoint,
                                           out,
                                           source_files,
                                           dump_output=dump_output)
-        elif language == 'go':
+        elif language == constants.LANGUAGES.GO:
             project = process_go_project(out,
                                          source_files,
                                          dump_output=dump_output)
-        elif language == 'jvm':
+        elif language == constants.LANGUAGES.JAVA:
             project = process_jvm_project(entrypoint,
                                           out,
                                           source_files,
                                           dump_output=dump_output)
-        elif language == 'rust':
+        elif language == constants.LANGUAGES.RUST:
             project = process_rust_project(out,
                                            source_files,
                                            dump_output=dump_output)
         else:
-            logger.error('Unsupported language: %s' % language)
+            logger.error('Unsupported language: %s', language)
             return Project([])
 
         # Process calltree and method data
@@ -303,16 +269,3 @@ def analyse_folder(language: str = '',
                     f.write(f'Call tree\n{calltree}')
 
     return project
-
-
-def main():
-    """Main"""
-
-    setup_logging()
-    args = parse_args()
-
-    analyse_folder(args.language, args.target_dir, args.entrypoint)
-
-
-if __name__ == "__main__":
-    main()

src/fuzz_introspector/utils.py

@@ -575,23 +575,17 @@ def locate_rust_fuzz_item(funcname: str, item_list: List[str]) -> str:
 
 def detect_language(directory) -> str:
     """Given a folder finds the likely programming language of the project"""
-    language_extensions = {
-        'c': ['.c', '.h'],
-        'cpp': ['.cpp', '.cc', '.c++', '.h', '.hpp'],
-        'jvm': ['.java'],
-        'rust': ['.rs']
-    }
+
     paths_to_avoid = [
         '/src/aflplusplus', '/src/honggfuzz', '/src/libfuzzer', '/src/fuzztest'
     ]
 
     language_counts: Dict[str, int] = {}
-
     for dirpath, _, filenames in os.walk(directory):
         if any([x for x in paths_to_avoid if dirpath.startswith(x)]):
             continue
         for filename in filenames:
-            for language, extensions in language_extensions.items():
+            for language, extensions in constants.LANGUAGE_EXTENSIONS.items():
                 if pathlib.Path(filename).suffix in extensions:
                     curr_count = language_counts.get(language, 0)
                     language_counts[language] = curr_count + 1