LoaderPE: Add `guess_arch` and link it to ContainerPE.arch

2 files changed, 14 insertions, 2 deletions

diff --git a/miasm2/analysis/binary.py b/miasm2/analysis/binary.py
index 98dd7b6a..c71c5e9b 100644
--- a/miasm2/analysis/binary.py
+++ b/miasm2/analysis/binary.py
@@ -122,9 +122,8 @@ class Container(object):
 class ContainerPE(Container):
     "Container abstraction for PE"
 
-
     def parse(self, data, vm=None):
-        from miasm2.jitter.loader.pe import vm_load_pe, preload_pe
+        from miasm2.jitter.loader.pe import vm_load_pe, preload_pe, guess_arch
         from elfesteem import pe_init
 
         # Parse signature
@@ -145,6 +144,9 @@ class ContainerPE(Container):
                 self._executable.NTsig.signature_value != 0x4550:
             raise ContainerSignatureException()
 
+        # Guess the architecture
+        self._arch = guess_arch(self._executable)
+
         # Build the bin_stream instance and set the entry point
         try:
             self._bin_stream = bin_stream_pe(self._executable.virt)
diff --git a/miasm2/jitter/loader/pe.py b/miasm2/jitter/loader/pe.py
index a3834d03..7c11b9c5 100644
--- a/miasm2/jitter/loader/pe.py
+++ b/miasm2/jitter/loader/pe.py
@@ -406,3 +406,13 @@ class libimp_pe(libimp):
                 all_ads = all_ads[i + 1:]
 
         return new_lib
+
+# machine -> arch
+PE_machine = {0x14c: "x86_32",
+              0x8664: "x86_64",
+              }
+
+def guess_arch(pe):
+    """Return the architecture specified by the PE container @pe.
+    If unknown, return None"""
+    return PE_machine.get(pe.Coffhdr.machine, None)