From fe87838dbe996feb4211bc16f27c800491e749db Mon Sep 17 00:00:00 2001
From: Michael Hunhoff <michael.hunhoff@mandiant.com>
Date: Mon, 21 Mar 2022 13:23:38 -0600
Subject: [PATCH] fixes #911

---
 CHANGELOG.md                         |  1 +
 capa/features/extractors/ida/file.py | 13 +++++++------
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 63ac459a..418842bd 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,7 @@
 ### Bug Fixes
 
 ### capa explorer IDA Pro plugin
+- improve file format extraction #918 @mike-hunhoff
 
 ### Development
 
diff --git a/capa/features/extractors/ida/file.py b/capa/features/extractors/ida/file.py
index 61a47d2d..d938e393 100644
--- a/capa/features/extractors/ida/file.py
+++ b/capa/features/extractors/ida/file.py
@@ -155,16 +155,17 @@ def extract_file_function_names():
 
 
 def extract_file_format():
-    format_name = ida_loader.get_file_type_name()
+    file_info = idaapi.get_inf_structure()
 
-    if "PE" in format_name:
+    if file_info.filetype == idaapi.f_PE:
         yield Format(FORMAT_PE), 0x0
-    elif "ELF64" in format_name:
-        yield Format(FORMAT_ELF), 0x0
-    elif "ELF32" in format_name:
+    elif file_info.filetype == idaapi.f_ELF:
         yield Format(FORMAT_ELF), 0x0
+    elif file_info.filetype == idaapi.f_BIN:
+        # no file type to return when processing a binary file, but we want to continue processing
+        return
     else:
-        raise NotImplementedError("file format: %s", format_name)
+        raise NotImplementedError("file format: %d" % file_info.filetype)
 
 
 def extract_features():