Skip to content

Commit

Permalink
new asyncio task add
Browse files Browse the repository at this point in the history
  • Loading branch information
LoRexxar committed Nov 4, 2019
1 parent 7df20c7 commit ed4fafa
Show file tree
Hide file tree
Showing 6 changed files with 44 additions and 14 deletions.
2 changes: 0 additions & 2 deletions .travis.yml
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
language: python
python:
- "3.5"
- "3.6"
- "3.7"
- "nightly"
install:
Expand Down
6 changes: 3 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
<big>**写在最前,Cobra-W仍然处于测试开发阶段,未发布正式版本,谨慎应用...**</big>
<big>**写在最前,Cobra-W就像手中的一把剑,这把剑好不好用是Cobra-W的事,如何使用是你的事,希望能有更多的人参与到Cobra-W的变化中来...**</big>

**请使用python3运行该工具,已停止维护python2.7环境**
**请使用python3.7+运行该工具,已停止维护python2.7环境**

# Cobra-W
[![GitHub (pre-)release](https://img.shields.io/github/release/LoRexxar/Cobra-W/all.svg)](https://github.com/LoRexxar/Cobra-W/releases)
[![license](https://img.shields.io/github/license/mashape/apistatus.svg?maxAge=2592000)](https://github.com/wufeifei/cobra/blob/master/LICENSE)
[![Build Status](https://travis-ci.org/LoRexxar/Cobra-W.svg?branch=master)](https://travis-ci.org/LoRexxar/Cobra-W)
[![Coverage Status](https://coveralls.io/repos/github/LoRexxar/Cobra-W/badge.svg?branch=master)](https://coveralls.io/github/LoRexxar/Cobra-W?branch=master)
![](https://img.shields.io/badge/language-python3.7-orange.svg)

## Introduction
Cobra是一款**源代码安全审计**工具,支持检测多种开发语言源代码中的**大部分显著**的安全问题和漏洞。
Expand Down
2 changes: 1 addition & 1 deletion cobra/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ def start(target, formatter, output, special_rules, a_sid=None, language=None, s

# Pretreatment ast object
ast_object.init_pre(target_directory, files)
ast_object.pre_ast(main_language)
ast_object.pre_ast_all(main_language)

# scan
scan(target_directory=target_directory, a_sid=a_sid, s_sid=s_sid, special_rules=pa.special_rules,
Expand Down
19 changes: 17 additions & 2 deletions cobra/engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,9 @@
import json
import os
import re
import asyncio
import traceback
import functools

import portalocker
from phply import phpast as php
Expand Down Expand Up @@ -167,11 +169,20 @@ def store(result):
else:
logger.debug('[SCAN] [STORE] Not found vulnerabilities on this rule!')

async def start_scan(target_directory, rule, files, language, secret_name):

result = scan_single(target_directory, rule, files, language, secret_name)
store(result)

if len(rules) == 0:
logger.critical('no rules!')
return False
logger.info('[PUSH] {rc} Rules'.format(rc=len(rules)))
push_rules = []
scan_list = []

# async start task
loop = asyncio.get_event_loop()

for idx, single_rule in enumerate(sorted(rules.keys())):

Expand All @@ -189,8 +200,12 @@ def store(result):
vulnerability=rule.vulnerability,
language=rule.language
))
result = scan_single(target_directory, rule, files, language, secret_name)
store(result)
# result = scan_single(target_directory, rule, files, language, secret_name)
scan_list.append(start_scan(target_directory, rule, files, language, secret_name))
# store(result)

loop.run_until_complete(asyncio.gather(*scan_list))
loop.close()

# print
data = []
Expand Down
27 changes: 22 additions & 5 deletions cobra/pretreatment.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@
import codecs
import traceback
import zipfile
import queue
import asyncio


could_ast_pase_lans = ["php", "chromeext", "javascript"]
Expand Down Expand Up @@ -76,12 +78,13 @@ class Pretreatment:

def __init__(self):
self.file_list = []
self.target_queue = queue.Queue()
self.target_directory = ""

self.pre_result = {}
self.define_dict = {}

self.pre_ast()
self.pre_ast_all()

def init_pre(self, target_directory, files):
self.file_list = files
Expand All @@ -99,7 +102,7 @@ def get_path(self, filepath):
else:
return os.path.join(self.target_directory, filepath)

def pre_ast(self, lan=None):
def pre_ast_all(self, lan=None):

if lan is not None:
# 检查是否在可ast pasre列表中
Expand All @@ -109,6 +112,17 @@ def pre_ast(self, lan=None):
return True

for fileext in self.file_list:
self.target_queue.put(fileext)

loop = asyncio.get_event_loop()
scan_list = (self.pre_ast() for i in range(20))
loop.run_until_complete(asyncio.gather(*scan_list))

async def pre_ast(self):

while not self.target_queue.empty():

fileext = self.target_queue.get()

if fileext[0] in ext_dict['php']:
# 下面是对于php文件的处理逻辑
Expand Down Expand Up @@ -145,7 +159,8 @@ def pre_ast(self, lan=None):
for node in all_nodes:
if isinstance(node, php.FunctionCall) and node.name == "define":
define_params = node.params
logger.debug("[AST][Pretreatment] new define {}={}".format(define_params[0].node, define_params[1].node))
logger.debug(
"[AST][Pretreatment] new define {}={}".format(define_params[0].node, define_params[1].node))

self.define_dict[define_params[0].node] = define_params[1].node

Expand All @@ -168,7 +183,7 @@ def pre_ast(self, lan=None):

# target可能是单个文件,这里需要专门处理
if not self.target_directory.endswith("/") and not self.target_directory.endswith("\\"):
relative_path = os.path.join(re.split(r'[\\|/]', self.target_directory)[-1]+"_files")
relative_path = os.path.join(re.split(r'[\\|/]', self.target_directory)[-1] + "_files")
else:
relative_path = target_files_path.split(self.target_directory)[-1]

Expand Down Expand Up @@ -212,7 +227,6 @@ def pre_ast(self, lan=None):
new_filepath = filepath + ".pretty"

if not os.path.isfile(new_filepath):

fi2 = codecs.open(new_filepath, "w", encoding='utf-8', errors='ignore')
code_content = jsbeautifier.beautify(code_content)
fi2.write(code_content)
Expand All @@ -235,6 +249,9 @@ def pre_ast(self, lan=None):
except:
logger.warning('[AST] something error, {}'.format(traceback.format_exc()))

return True


def get_nodes(self, filepath, vul_lineno=None, lan=None):
filepath = os.path.normpath(filepath)

Expand Down
2 changes: 1 addition & 1 deletion tests/test_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@

files = [('.php', {'list': ["v_parser.php", "v.php"]})]
ast_object.init_pre(project_directory + '/tests/vulnerabilities/', files)
ast_object.pre_ast()
ast_object.pre_ast_all()


target_projects = project_directory + '/tests/vulnerabilities/v_parser.php'
Expand Down

0 comments on commit ed4fafa

Please sign in to comment.