Merge pull request #273 from AceAttorney/feature/add-sp_helptext

添加可执行命令白名单及放行sp_helptext命令。

sql/engines/mssql.py

@@ -81,6 +81,9 @@ def query_check(self, db_name=None, sql=''):
                            "string_escape", "string_split", "stuff", "substring", "trim", "unicode"]
         keyword_warning = ''
         star_patter = r"(^|,| )\*( |\(|$)"
+        sql_whitelist = ['select', 'sp_helptext']
+        # 根据白名单list拼接pattern语句
+        whitelist_pattern = "^" + "|^".join(sql_whitelist)
         # 删除注释语句，进行语法判断，执行第一条有效sql
         try:
             sql = sql.format(sql, strip_comments=True)
@@ -91,9 +94,9 @@ def query_check(self, db_name=None, sql=''):
             result['has_star'] = True
             result['msg'] = '没有有效的SQL语句'
             return result
-        if re.match(r"^select", sql_lower) is None:
+        if re.match(whitelist_pattern, sql_lower) is None:
             result['bad_query'] = True
-            result['msg'] = '仅支持^select语法!'
+            result['msg'] = '仅支持{}语法!'.format(','.join(sql_whitelist))
             return result
         if re.search(star_patter, sql_lower) is not None:
             keyword_warning += '禁止使用 * 关键词\n'

sql/engines/tests.py

@@ -162,6 +162,10 @@ def testQueryCheck(self):
         banned_sql = 'select phone from user_table where phone=concat(phone,1)'
         check_result = new_engine.query_check(db_name='some_db', sql=banned_sql)
         self.assertTrue(check_result.get('bad_query'))
+        sp_sql = "sp_helptext '[SomeName].[SomeAction]'"
+        check_result = new_engine.query_check(db_name='some_db', sql=sp_sql)
+        self.assertFalse(check_result.get('bad_query'))
+        self.assertEqual(check_result.get('filtered_sql'), sp_sql)
 
     def test_filter_sql(self):
         new_engine = MssqlEngine(instance=self.ins1)