Merge pull request #334 from mrphrazer/fix_smt2translator

Fix smt2translator

2 files changed, 7 insertions, 1 deletions

diff --git a/miasm2/ir/translators/__init__.py b/miasm2/ir/translators/__init__.py
index 22e00d98..d3678ffc 100644
--- a/miasm2/ir/translators/__init__.py
+++ b/miasm2/ir/translators/__init__.py
@@ -3,6 +3,7 @@ from miasm2.ir.translators.translator import Translator
 import miasm2.ir.translators.C
 import miasm2.ir.translators.python
 import miasm2.ir.translators.miasm
+import miasm2.ir.translators.smt2
 try:
     import miasm2.ir.translators.z3_ir
 except ImportError:
diff --git a/miasm2/ir/translators/smt2.py b/miasm2/ir/translators/smt2.py
index 21f40e54..5bffd7f2 100644
--- a/miasm2/ir/translators/smt2.py
+++ b/miasm2/ir/translators/smt2.py
@@ -284,13 +284,14 @@ class TranslatorSMT2(Translator):
         dst = self.from_expr(expr.dst)
         return smt2_assert(smt2_eq(src, dst))
 
-    def to_smt2(self, exprs, logic="QF_ABV"):
+    def to_smt2(self, exprs, logic="QF_ABV", model=False):
         """
         Converts a valid SMT2 file for a given list of
         SMT2 expressions.
 
         :param exprs: list of SMT2 expressions
         :param logic: SMT2 logic
+        :param model: model generation flag
         :return: String of the SMT2 file
         """
         ret = ""
@@ -313,6 +314,10 @@ class TranslatorSMT2(Translator):
         # define action
         ret += "(check-sat)\n"
 
+        # enable model generation
+        if model:
+            ret += "(get-model)\n"
+
         return ret