Merge pull request #110 from lostsnow/feature/code-clean

code clean

dongtai_agent_python/assess/common_hook.py

@@ -1,49 +1,39 @@
 import sys
 
-from dongtai_agent_python import CONTEXT_TRACKER
 from dongtai_agent_python.policy.deal_data import wrap_data
-from dongtai_agent_python.setting import Setting, const
-from dongtai_agent_python.utils import scope, utils
+from dongtai_agent_python.setting import const
+from dongtai_agent_python.utils import scope
 
 
 # 普通方法 hook
 class InstallFcnHook(object):
-    def __init__(self, old_cls, fcn, signature=None, node_type=None):
+    def __init__(self, origin_cls, origin_func, signature=None, node_type=None):
         self.signature = signature
-        self._fcn = fcn
-        self.__name__ = fcn.__name__
+        self.origin_func = origin_func
+        self.__name__ = origin_func.__name__
 
-        self.old_cls = old_cls
+        self.origin_cls = origin_cls
         self.node_type = node_type
 
     def __call__(self, *args, **kwargs):
         if self.node_type == const.NODE_TYPE_FILTER:
             with scope.scope(scope.SCOPE_AGENT):
-                ret_val = self._fcn(*args, **kwargs)
+                result = self.origin_func(*args, **kwargs)
         else:
-            ret_val = self._fcn(*args, **kwargs)
+            result = self.origin_func(*args, **kwargs)
 
         if scope.in_scope(scope.SCOPE_AGENT):
-            return ret_val
-
-        context = CONTEXT_TRACKER.current()
-        if not utils.needs_propagation(context, self.node_type):
-            return ret_val
-
-        with scope.scope(scope.SCOPE_AGENT):
-            setting = Setting()
-        if setting.is_agent_paused():
-            return ret_val
+            return result
 
         wrap_data(
-            ret_val, self.old_cls.__name__, self._fcn.__name__,
+            result, self.origin_cls.__name__, self.origin_func.__name__,
             signature=self.signature, node_type=self.node_type,
             come_args=args, come_kwargs=kwargs)
 
-        return ret_val
+        return result
 
 
-def build_exec_eval_patch(orig_cls, orig_func, signature, node_type):
+def build_exec_eval_patch(origin_cls, origin_func, signature, node_type):
     def exec_eval_patch(code, globs=None, locs=None):
         """
         Code ported from six module
@@ -62,27 +52,18 @@ def exec_eval_patch(code, globs=None, locs=None):
 
         try:
             with scope.scope(scope.SCOPE_AGENT):
-                ret_val = orig_func(code, globs, locs)
+                result = origin_func(code, globs, locs)
         except Exception:
-            ret_val = None
             raise
 
         if scope.in_scope(scope.SCOPE_AGENT):
-            return ret_val
-
-        context = CONTEXT_TRACKER.current()
-        if not utils.needs_propagation(context, node_type):
-            return ret_val
-
-        with scope.scope(scope.SCOPE_AGENT):
-            setting = Setting()
-        if setting.is_agent_paused():
-            return ret_val
+            return result
 
         wrap_data(
-            ret_val, orig_cls.__name__, orig_func.__name__,
+            result, origin_cls.__name__, origin_func.__name__,
             signature=signature, node_type=node_type,
             come_args=[code])
-        return ret_val
+
+        return result
 
     return exec_eval_patch

dongtai_agent_python/assess/ctypes_hook.py

@@ -2,10 +2,8 @@
 import os
 import sys
 
-from dongtai_agent_python import CONTEXT_TRACKER
 from dongtai_agent_python.policy.deal_data import wrap_data
-from dongtai_agent_python.utils import utils
-from dongtai_agent_python.setting import const, Setting
+from dongtai_agent_python.setting import const
 from dongtai_agent_python.utils import scope
 
 
@@ -34,7 +32,7 @@ def new_func(origin_cls, method_name, signature=None, node_type=None, *args, **k
     copy_new_class = type(origin_cls.__name__, origin_cls.__bases__, dict(origin_cls.__dict__))
     if method_name not in copy_new_class.__dict__:
         return None
-    origin_fcn = getattr(origin_cls, method_name)
+    origin_func = getattr(origin_cls, method_name)
 
     def child_func(*args, **kwargs):
         if node_type == const.NODE_TYPE_FILTER:
@@ -45,34 +43,10 @@ def child_func(*args, **kwargs):
         if scope.in_scope(scope.SCOPE_AGENT):
             return result
 
-        context = CONTEXT_TRACKER.current()
-        if not utils.needs_propagation(context, node_type):
-            return result
-
-        setting = Setting()
-        if setting.is_agent_paused():
-            return result
-
-        if ((args == ([], '*.mo') or args == (['*.mo'], '**')) and method_name == "append") or (
-                args == ('**/*.mo', '/') and method_name == "split"):
-            return result
-
-        if signature in const.FIRST_RETURN and len(args) > 0:
-            real_result = args[0]
-        else:
-            real_result = result
-
-        with scope.scope(scope.SCOPE_AGENT):
-            come_args = []
-            for k, v in enumerate(args):
-                if signature in const.FIRST_RETURN and k == 0:
-                    continue
-                come_args.append(v)
-
         wrap_data(
-            real_result, origin_cls.__name__, origin_fcn.__name__,
+            result, origin_cls.__name__, origin_func.__name__,
             signature=signature, node_type=node_type,
-            come_args=come_args, come_kwargs=kwargs)
+            come_args=args, come_kwargs=kwargs)
 
         return result
 

dongtai_agent_python/assess/patch.py

@@ -46,19 +46,14 @@ def enable_patches(policies):
                 continue
 
             if policy == 'builtins.eval' or policy == 'builtins.exec':
-                imp_arr = copy.deepcopy(policy_arr)
-                method_name = imp_arr[-1]
-                setting.policy[method_name] = rules['type']
-                # 存储到全局变量
-                del imp_arr[-1]
-                policy_str = ".".join(imp_arr)
-                old_module = HookLazyImport(policy_str, [method_name])
-                old_func = getattr(old_module, method_name)
-                old_cls = old_module.origin_module()
-
-                new_fn = build_exec_eval_patch(old_cls, old_func, policy, rules['type'])
-                after_cls = magic_get_dict(old_cls)
-                after_cls[method_name] = new_fn
+                method_name = policy_arr[-1]
+                new_module = HookLazyImport('builtins', [method_name])
+                origin_func = getattr(new_module, method_name)
+                origin_cls = new_module.origin_module()
+
+                new_fn = build_exec_eval_patch(origin_cls, origin_func, policy, rules['type'])
+                origin_cls_ptr = magic_get_dict(origin_cls)
+                origin_cls_ptr[method_name] = new_fn
                 logger.debug("------exec_eval_patch---------- " + "[" + str(rules['type']) + "]" + policy)
                 has_patched[policy] = True
                 continue
@@ -69,22 +64,20 @@ def enable_patches(policies):
             try:
                 imp_arr = copy.deepcopy(policy_arr)
                 method_name = imp_arr[-1]
-                setting.policy[method_name] = rules['type']
-                # 存储到全局变量
                 del imp_arr[-1]
                 policy_str = ".".join(imp_arr)
-                old_module = HookLazyImport(policy_str, [method_name])
-                old_func = getattr(old_module, method_name)
-                old_cls = old_module.origin_module()
-                if old_cls is None:
+                new_module = HookLazyImport(policy_str, [method_name])
+                origin_func = getattr(new_module, method_name)
+                origin_cls = new_module.origin_module()
+                if origin_cls is None:
                     imp_arr = copy.deepcopy(policy_arr)
                     method_name = imp_arr[-1]
                     class_name = imp_arr[-2]
                     del imp_arr[-1]
                     del imp_arr[-1]
                     policy_str = ".".join(imp_arr)
-                    old_module = HookLazyImport(policy_str, [class_name])
-                    old_cls = getattr(old_module, class_name)
+                    new_module = HookLazyImport(policy_str, [class_name])
+                    origin_cls = getattr(new_module, class_name)
             except Exception as e:
                 imp_arr = copy.deepcopy(policy_arr)
                 if imp_arr[0] not in sys.modules:
@@ -98,16 +91,16 @@ def enable_patches(policies):
                 policy_str = ".".join(imp_arr)
 
                 try:
-                    old_module = HookLazyImport(policy_str, [class_name])
-                    old_cls = getattr(old_module, class_name)
+                    new_module = HookLazyImport(policy_str, [class_name])
+                    origin_cls = getattr(new_module, class_name)
                 except Exception as e:
                     continue
 
-            after_cls = magic_get_dict(old_cls)
+            after_cls = magic_get_dict(origin_cls)
 
-            if isinstance(old_cls, type):
+            if isinstance(origin_cls, type):
                 hooked = new_func(
-                    old_cls,
+                    origin_cls,
                     method_name,
                     policy,
                     rules['type']
@@ -118,7 +111,7 @@ def enable_patches(policies):
                 after_cls[method_name] = hooked
             else:
                 logger.debug("------origin_cls_function------ " + "[" + str(rules['type']) + "]" + policy)
-                after_cls[method_name] = InstallFcnHook(old_cls, old_func, policy, rules['type'])
+                after_cls[method_name] = InstallFcnHook(origin_cls, origin_func, policy, rules['type'])
 
             has_patched[policy] = True
 

dongtai_agent_python/policy/deal_data.py

@@ -1,23 +1,27 @@
 from dongtai_agent_python import CONTEXT_TRACKER
 from dongtai_agent_python.policy.tracking import Tracking
-from dongtai_agent_python.setting import const
+from dongtai_agent_python.setting import Setting, const
 from dongtai_agent_python.utils import scope, utils
 
 
 @scope.with_scope(scope.SCOPE_AGENT)
-def wrap_data(target, origin_cls, origin_func, signature=None, node_type=None, come_args=None, come_kwargs=None):
+def wrap_data(result, class_name, func_name, signature=None, node_type=None, come_args=None, come_kwargs=None):
+    setting = Setting()
+    if setting.is_agent_paused():
+        return
+
     context = CONTEXT_TRACKER.current()
     if not utils.needs_propagation(context, node_type):
         return
 
-    if not filter_result(target, node_type):
+    if not filter_result(result, node_type):
         return
 
     if node_type == const.NODE_TYPE_SOURCE:
         context.has_source = True
 
-    tracking = Tracking(signature, node_type, origin_cls, origin_func)
-    tracking.apply(come_args, come_kwargs, target)
+    tracking = Tracking(signature, node_type, class_name, func_name)
+    tracking.apply(come_args, come_kwargs, result)
 
 
 def filter_result(result, node_type=None):

dongtai_agent_python/policy/tracking.py

@@ -11,7 +11,7 @@
 
 
 class Tracking(object):
-    def __init__(self, signature, node_type, origin_cls, func_name):
+    def __init__(self, signature, node_type, class_name, func_name):
         scope.enter_scope(scope.SCOPE_AGENT)
 
         self.context = CONTEXT_TRACKER.current()
@@ -26,7 +26,7 @@ def __init__(self, signature, node_type, origin_cls, func_name):
         self.signature = signature
         self.node_type = node_type
 
-        self.class_name = origin_cls
+        self.class_name = class_name
         if signature.endswith("." + func_name):
             self.class_name = signature[:-len(func_name) - 1]
 

dongtai_agent_python/setting/const.py

@@ -16,11 +16,6 @@
 TAINT_SOURCE = 1
 TAINT_TARGET = 2
 
-FIRST_RETURN = [
-    'builtins.list.append',
-    'builtins.list.insert',
-]
-
 C_API_PATCHES = [
     'builtins.str.fstring',
     'builtins.str.cformat',