From d2fa2a7616d145868e311b0cd0a6ddff58538ab0 Mon Sep 17 00:00:00 2001
From: a1ext <a13x4nd3r.t@gmail.com>
Date: Sat, 28 Oct 2017 18:13:37 +0300
Subject: [PATCH] Don't rename function if there are jumps to outside.

---
 auto_re.py | 35 +++++++++++++++++++++++++++++++----
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/auto_re.py b/auto_re.py
index a05399e..a04597e 100644
--- a/auto_re.py
+++ b/auto_re.py
@@ -4,7 +4,7 @@
 from collections import defaultdict
 import idaapi
 from idautils import FuncItems, CodeRefsTo
-from idaapi import o_reg, o_imm, o_far, o_near, o_mem
+from idaapi import o_reg, o_imm, o_far, o_near, o_mem, o_displ
 import os
 import sys
 import traceback
@@ -22,7 +22,7 @@
 # enable to allow PyCharm remote debug
 RDEBUG = False
 # adjust this value to be a full path to a debug egg
-RDEBUG_EGG = r'c:\Program Files (x86)\JetBrains\PyCharm 2016.3\debug-eggs\pycharm-debug.egg'
+RDEBUG_EGG = r'c:\Program Files\JetBrains\PyCharm 2017.1.4\debug-eggs\pycharm-debug.egg'
 RDEBUG_HOST = 'localhost'
 RDEBUG_PORT = 12321
 
@@ -209,6 +209,8 @@ class auto_re_t(idaapi.plugin_t):
     }
     _DEFAULT_CALLEE_NODE_NAME = '$ vmm functions'
 
+    _JMP_TYPES = {idaapi.NN_jmp, idaapi.NN_jmpni, idaapi.NN_jmpfi, idaapi.NN_jmpshort}
+
     def __init__(self):
         super(auto_re_t, self).__init__()
         self._data = None
@@ -269,6 +271,10 @@ def _handle_calls(self, fn, fn_an):
         if idaapi.has_dummy_name(idaapi.getFlags(ea)):
             return
 
+        # TODO: check is there jmp, push+retn then don't rename the func
+        if fn_an['strange_flow']:
+            return
+
         possible_name = idaapi.get_ea_name(ea)
         if not possible_name or possible_name in blacklist:
             return
@@ -286,7 +292,7 @@ def _handle_calls(self, fn, fn_an):
     # noinspection PyMethodMayBeStatic
     def _check_is_jmp_wrapper(self, dis):
         # checks instructions like `jmp API`
-        if dis.itype not in (idaapi.NN_jmp, idaapi.NN_jmpni, idaapi.NN_jmpfi, idaapi.NN_jmpshort):
+        if dis.itype not in self._JMP_TYPES:
             return
 
         # handle call wrappers like jmp GetProcAddress
@@ -459,8 +465,16 @@ def _analysis_handle_call_insn(cls, dis, rv):
 
     @classmethod
     def analyze_func(cls, fn):
-        rv = {'fn': fn, 'calls': [], 'math': [], 'has_bads': False, 'tags': defaultdict(list)}
+        rv = {
+            'fn': fn,
+            'calls': [],
+            'math': [],
+            'has_bads': False,
+            'strange_flow': False,
+            'tags': defaultdict(list)
+        }
         items = cls.disasm_func(fn)
+        items_set = set(map(lambda x: x['ea'], items))
 
         for item in items:
             dis = item['dis']
@@ -478,6 +492,19 @@ def analyze_func(cls, fn):
                                idaapi.NN_rol, idaapi.NN_rcl, idaapi.NN_rcl):
                 # TODO
                 rv['math'].append(dis)
+            elif dis.itype in cls._JMP_TYPES:
+                if dis.Op1.type not in (o_far, o_near, o_mem, o_displ):
+                    continue
+
+                if dis.Op1.type == o_displ:
+                    rv['strange_flow'] = True
+                    continue
+
+                ea = dis.Op1.value
+                if not ea and dis.Op1.addr:
+                    ea = dis.Op1.addr
+                if ea not in items_set:
+                    rv['strange_flow'] = True
 
         return rv