From ad03ea5dd128234f161d6376698202df886d14dc Mon Sep 17 00:00:00 2001
From: c-urly <bhattrishi8@gmail.com>
Date: Thu, 14 Mar 2024 11:42:11 -0400
Subject: [PATCH] ELF support

---
 floss/const.py |  2 +-
 floss/main.py  | 50 +++++++++++++++++++++++++++++++++++---------------
 2 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/floss/const.py b/floss/const.py
index 336988083..653ce056c 100644
--- a/floss/const.py
+++ b/floss/const.py
@@ -3,7 +3,7 @@
 KILOBYTE = 1024
 MEGABYTE = 1024 * KILOBYTE
 MAX_FILE_SIZE = 16 * MEGABYTE
-SUPPORTED_FILE_MAGIC = {b"MZ"}
+SUPPORTED_FILE_MAGIC = {b"MZ", b"\x7fELF"}
 MIN_STRING_LENGTH = 4
 MAX_STRING_LENGTH = 2048
 
diff --git a/floss/main.py b/floss/main.py
index a49edea73..43feb546e 100644
--- a/floss/main.py
+++ b/floss/main.py
@@ -70,6 +70,11 @@ class StringType(str, Enum):
     DECODED = "decoded"
 
 
+class FileType:
+    PE = False
+    ELF = False
+
+
 class WorkspaceLoadError(ValueError):
     pass
 
@@ -200,9 +205,11 @@ def make_parser(argv):
         type=str,
         choices=[l.value for l in Language if l != Language.UNKNOWN],
         default=Language.UNKNOWN.value,
-        help="use language-specific string extraction, auto-detect language by default, disable using 'none'"
-        if show_all_options
-        else argparse.SUPPRESS,
+        help=(
+            "use language-specific string extraction, auto-detect language by default, disable using 'none'"
+            if show_all_options
+            else argparse.SUPPRESS
+        ),
     )
     advanced_group.add_argument(
         "-l",
@@ -215,9 +222,11 @@ def make_parser(argv):
         type=lambda x: int(x, 0x10),
         default=None,
         nargs="+",
-        help="only analyze the specified functions, hex-encoded like 0x401000, space-separate multiple functions"
-        if show_all_options
-        else argparse.SUPPRESS,
+        help=(
+            "only analyze the specified functions, hex-encoded like 0x401000, space-separate multiple functions"
+            if show_all_options
+            else argparse.SUPPRESS
+        ),
     )
     advanced_group.add_argument(
         "--disable-progress",
@@ -228,17 +237,21 @@ def make_parser(argv):
         "--signatures",
         type=str,
         default=SIGNATURES_PATH_DEFAULT_STRING,
-        help="path to .sig/.pat file or directory used to identify library functions, use embedded signatures by default"
-        if show_all_options
-        else argparse.SUPPRESS,
+        help=(
+            "path to .sig/.pat file or directory used to identify library functions, use embedded signatures by default"
+            if show_all_options
+            else argparse.SUPPRESS
+        ),
     )
     advanced_group.add_argument(
         "-L",
         "--large-file",
         action="store_true",
-        help="allow processing files larger than {} MB".format(int(MAX_FILE_SIZE / MEGABYTE))
-        if show_all_options
-        else argparse.SUPPRESS,
+        help=(
+            "allow processing files larger than {} MB".format(int(MAX_FILE_SIZE / MEGABYTE))
+            if show_all_options
+            else argparse.SUPPRESS
+        ),
     )
     advanced_group.add_argument(
         "--version",
@@ -356,9 +369,13 @@ def is_supported_file_type(sample_file_path: Path):
     :return: True if file type is supported, False otherwise
     """
     with sample_file_path.open("rb") as f:
-        magic = f.read(2)
+        magic = f.read(4)
 
     if magic in SUPPORTED_FILE_MAGIC:
+        FileType.ELF = True
+        return True
+    elif magic[:2] in SUPPORTED_FILE_MAGIC:
+        FileType.PE = True
         return True
     else:
         return False
@@ -390,7 +407,8 @@ def load_vw(
     else:
         vw = viv_utils.getWorkspace(str(sample_path), analyze=False, should_save=False)
 
-    viv_utils.flirt.register_flirt_signature_analyzers(vw, list(map(str, sigpaths)))
+    if not FileType.ELF:
+        viv_utils.flirt.register_flirt_signature_analyzers(vw, list(map(str, sigpaths)))
 
     vw.analyze()
 
@@ -547,6 +565,8 @@ def main(argv=None) -> int:
         return 0
 
     static_runtime = get_runtime_diff(interim)
+    if not is_supported_file_type(sample):
+        logger.error("FileType not Supported")
 
     # set language configurations
     selected_lang = Language(args.language)
@@ -554,7 +574,7 @@ def main(argv=None) -> int:
         results.metadata.language = ""
         results.metadata.language_version = ""
         results.metadata.language_selected = ""
-    else:
+    elif FileType.PE:
         lang_id, lang_version = identify_language_and_version(sample, static_strings)
 
         if selected_lang == Language.UNKNOWN: