Merge pull request #116 from commial/container-arch

Container arch

5 files changed, 68 insertions, 13 deletions

diff --git a/example/disasm/full.py b/example/disasm/full.py
index 14829755..9e953122 100644
--- a/example/disasm/full.py
+++ b/example/disasm/full.py
@@ -21,11 +21,11 @@ if filename and os.path.isfile(filename):
 
 
 parser = ArgumentParser("Disassemble a binary")
-parser.add_argument('architecture', help="architecture: " + \
-                        ",".join(Machine.available_machine()))
 parser.add_argument('filename', help="File to disassemble")
 parser.add_argument('address', help="Starting address for disassembly engine",
-                    nargs="+")
+                    nargs="*")
+parser.add_argument('-m', '--architecture', help="architecture: " + \
+                        ",".join(Machine.available_machine()))
 parser.add_argument('-f', "--followcall", action="store_true",
                     help="Follow call instructions")
 parser.add_argument('-b', "--blockwatchdog", default=None, type=int,
@@ -55,12 +55,6 @@ args = parser.parse_args()
 if args.verbose:
     log_asmbloc.setLevel(logging.DEBUG)
 
-log.info("import machine...")
-machine = Machine(args.architecture)
-mn, dis_engine = machine.mn, machine.dis_engine
-ira, ir = machine.ira, machine.ir
-log.info('ok')
-
 log.info('Load binary')
 with open(args.filename) as fdesc:
     cont = Container.from_stream(fdesc, addr=args.shiftoffset)
@@ -68,8 +62,21 @@ with open(args.filename) as fdesc:
 default_addr = cont.entry_point
 bs = cont.bin_stream
 e = cont.executable
+log.info('ok')
 
+log.info("import machine...")
+# Use the guessed architecture or the specified one
+arch = args.architecture if args.architecture else cont.arch
+if not arch:
+    print "Architecture recognition fail. Please specify it in arguments"
+    exit(-1)
+
+# Instance the arch-dependent machine
+machine = Machine(arch)
+mn, dis_engine = machine.mn, machine.dis_engine
+ira, ir = machine.ira, machine.ir
 log.info('ok')
+
 mdis = dis_engine(bs)
 # configure disasm engine
 mdis.dontdis_retcall = args.dontdis_retcall
diff --git a/miasm2/analysis/binary.py b/miasm2/analysis/binary.py
index f5cecc87..c71c5e9b 100644
--- a/miasm2/analysis/binary.py
+++ b/miasm2/analysis/binary.py
@@ -88,6 +88,13 @@ class Container(object):
 
     def __init__(self, *args, **kwargs):
         "Alias for 'parse'"
+        # Init attributes
+        self._executable = None
+        self._bin_stream = None
+        self._entry_point = None
+        self._arch = None
+
+        # Launch parsing
         self.parse(*args, **kwargs)
 
     @property
@@ -105,14 +112,18 @@ class Container(object):
         "Return the detected entry_point"
         return self._entry_point
 
+    @property
+    def arch(self):
+        "Return the guessed architecture"
+        return self._arch
+
 
 ## Format dependent classes
 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
@@ -133,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)
@@ -146,7 +160,8 @@ class ContainerELF(Container):
     "Container abstraction for ELF"
 
     def parse(self, data, vm=None):
-        from miasm2.jitter.loader.elf import vm_load_elf, preload_elf
+        from miasm2.jitter.loader.elf import \
+            vm_load_elf, preload_elf, guess_arch
         from elfesteem import elf_init
 
         # Parse signature
@@ -162,6 +177,9 @@ class ContainerELF(Container):
         except Exception, error:
             raise ContainerParsingException('Cannot read ELF: %s' % error)
 
+        # 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_elf(self._executable.virt)
diff --git a/miasm2/jitter/loader/elf.py b/miasm2/jitter/loader/elf.py
index 916b37c4..c0427e79 100644
--- a/miasm2/jitter/loader/elf.py
+++ b/miasm2/jitter/loader/elf.py
@@ -3,6 +3,8 @@ from collections import defaultdict
 
 from elfesteem import cstruct
 from elfesteem import *
+import elfesteem.elf as elf_csts
+
 from miasm2.jitter.csts import *
 from miasm2.jitter.loader.utils import canon_libname_libfunc, libimp
 from miasm2.core.interval import interval
@@ -80,3 +82,19 @@ def vm_load_elf(vm, fdata, **kargs):
 
 class libimp_elf(libimp):
     pass
+
+
+# machine, size, sex -> arch_name
+ELF_machine = {(elf_csts.EM_ARM, 32, elf_csts.ELFDATA2LSB): "arml",
+               (elf_csts.EM_ARM, 32, elf_csts.ELFDATA2MSB): "armb",
+               (elf_csts.EM_MIPS, 32, elf_csts.ELFDATA2MSB): "mips32b",
+               (elf_csts.EM_MIPS, 32, elf_csts.ELFDATA2LSB): "mips32l",
+               (elf_csts.EM_386, 32, elf_csts.ELFDATA2LSB): "x86_32",
+               (elf_csts.EM_X86_64, 64, elf_csts.ELFDATA2LSB): "x86_64",
+               (elf_csts.EM_SH, 32, elf_csts.ELFDATA2LSB): "sh4",
+               }
+
+def guess_arch(elf):
+    """Return the architecture specified by the ELF container @elf.
+    If unknown, return None"""
+    return ELF_machine.get((elf.Ehdr.machine, elf.size, elf.sex), None)
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)
diff --git a/test/test_all.py b/test/test_all.py
index 8d759053..ecf9a63b 100644
--- a/test/test_all.py
+++ b/test/test_all.py
@@ -236,6 +236,8 @@ for script, prods in [(["single_instr.py"], []),
                       (["function.py"], ["graph.txt"]),
                       (["file.py", Example.get_sample("box_upx.exe"),
                         "0x410f90"], ["graph.txt"]),
+                      (["full.py", Example.get_sample("box_upx.exe")],
+                       ["graph_execflow.txt", "lines.txt"]),
                       ]:
     testset += ExampleDisassembler(script, products=prods)
 
@@ -249,7 +251,7 @@ class ExampleDisasmFull(ExampleDisassembler):
 
     def __init__(self, *args, **kwargs):
         super(ExampleDisasmFull, self).__init__(*args, **kwargs)
-        self.command_line = ["full.py", "-g", "-s"] + self.command_line
+        self.command_line = ["full.py", "-g", "-s", "-m"] + self.command_line
         self.products += ["graph_execflow.txt", "graph_irflow.txt", "lines.txt"]