From 78025ad2bcb0802722dde75c878790a35c9fe810 Mon Sep 17 00:00:00 2001
From: BlBana <635373043@qq.com>
Date: Thu, 7 Sep 2017 16:07:48 +0800
Subject: [PATCH] FPC-Mode vuls repair, fixed #559

---
 cobra/engine.py |  5 +++--
 cobra/parser.py | 12 +++++++++---
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/cobra/engine.py b/cobra/engine.py
index 4aae64e7..84caa107 100644
--- a/cobra/engine.py
+++ b/cobra/engine.py
@@ -621,12 +621,13 @@ def scan(self):
                 try:
                     ast = CAST(self.rule_match, self.target_directory, self.file_path, self.line_number, self.code_content)
                     if self.rule_match_mode == const.mm_function_param_controllable:
-                        rule_match = self.rule_match.strip('()').split('|')
+                        rule_match = self.rule_match.strip('()').split('|')  # 漏洞规则整理为列表
+                        rule_repair = self.rule_repair.strip('()').split('|')  # 修复规则整理为列表
                         logger.debug('[RULE_MATCH] {r}'.format(r=rule_match))
                         try:
                             with open(self.file_path, 'r') as fi:
                                 code_contents = fi.read()
-                                result = scan_parser(code_contents, rule_match, self.line_number)
+                                result = scan_parser(code_contents, rule_match, self.line_number, rule_repair)
                                 logger.debug('[AST] [RET] {c}'.format(c=result))
                                 if len(result) > 0:
                                     if result[0]['code'] == 1:  # 函数参数可控
diff --git a/cobra/parser.py b/cobra/parser.py
index f3dc3db3..39c0f04a 100644
--- a/cobra/parser.py
+++ b/cobra/parser.py
@@ -18,6 +18,7 @@
 
 with_line = True
 scan_results = []  # 结果存放列表初始化
+repairs = []  # 用于存放修复函数
 
 
 def export(items):
@@ -213,8 +214,10 @@ def is_repair(expr):
     :return:
     """
     is_re = False  # 是否修复，默认值是未修复
-    if expr == 'escapeshellcmd':
-        is_re = True
+    for repair in repairs:
+        if expr == repair:
+            is_re = True
+            return is_re
     return is_re
 
 
@@ -661,16 +664,19 @@ def analysis(nodes, vul_function, back_node, vul_lineo, function_params=None):
         back_node.append(node)
 
 
-def scan_parser(code_content, sensitive_func, vul_lineno):
+def scan_parser(code_content, sensitive_func, vul_lineno, repair):
     """
     开始检测函数
     :param code_content: 要检测的文件内容
     :param sensitive_func: 要检测的敏感函数,传入的为函数列表
     :param vul_lineno: 漏洞函数所在行号
+    :param repair: 对应漏洞的修复函数列表
     :return:
     """
     try:
+        global repairs
         global scan_results
+        repairs = repair
         scan_results = []
         parser = make_parser()
         all_nodes = parser.parse(code_content, debug=False, lexer=lexer.clone(), tracking=with_line)