Skip to content
This repository was archived by the owner on Nov 28, 2023. It is now read-only.

Commit e23f4b9

Browse files
authored
Merge pull request #923 from BlBana/master
解决了Python3下代码拉取bug
2 parents aff9982 + 306fbd4 commit e23f4b9

File tree

4 files changed

+68
-33
lines changed

4 files changed

+68
-33
lines changed

cobra/export.py

+21-17
Original file line numberDiff line numberDiff line change
@@ -144,24 +144,28 @@ def write_to_file(target, sid, output_format='', filename=None):
144144
logger.info('Vulnerabilities\n' + str(dict_to_pretty_table(scan_data.get('vulnerabilities'))))
145145

146146
elif output_format == 'json' or output_format == 'JSON':
147-
if not os.path.exists(filename):
148-
with open(filename, 'w', encoding='utf-8') as f:
149-
json_data = {
150-
sid: scan_data,
151-
}
152-
f.write(dict_to_json(json_data))
153-
else:
154-
with open(filename, 'r+', encoding='utf-8') as f:
155-
try:
156-
json_data = json.load(f)
157-
json_data.update({sid: scan_data})
158-
# 使用 r+ 模式不会覆盖,调整文件指针到开头
159-
f.seek(0)
160-
f.truncate()
147+
try:
148+
if not os.path.exists(filename):
149+
with open(filename, 'w', encoding='utf-8') as f:
150+
json_data = {
151+
sid: scan_data,
152+
}
161153
f.write(dict_to_json(json_data))
162-
except ValueError:
163-
logger.warning('[EXPORT] The json file have invaild token or None: {}'.format(os.path.join(export_path, filename)))
164-
return False
154+
else:
155+
with open(filename, 'r+', encoding='utf-8') as f:
156+
try:
157+
json_data = json.load(f)
158+
json_data.update({sid: scan_data})
159+
# 使用 r+ 模式不会覆盖,调整文件指针到开头
160+
f.seek(0)
161+
f.truncate()
162+
f.write(dict_to_json(json_data))
163+
except ValueError:
164+
logger.warning('[EXPORT] The json file have invaild token or None: {}'.format(os.path.join(export_path, filename)))
165+
return False
166+
except IOError:
167+
logger.warning('[EXPORT] Please input a file path after the -o parameter')
168+
return False
165169

166170
elif output_format == 'xml' or output_format == 'XML':
167171
xml_data = {

cobra/pickup.py

+19
Original file line numberDiff line numberDiff line change
@@ -387,6 +387,10 @@ def clone(self):
387387

388388
p = subprocess.Popen(cmd, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
389389
(clone_out, clone_err) = p.communicate()
390+
391+
clone_out = clone_out.decode('utf-8')
392+
clone_err = clone_err.decode('utf-8')
393+
390394
clone_err = clone_err.replace('{0}:{1}'.format(self.repo_username, self.repo_password), '')
391395

392396
logger.debug('[PICKUP] [CLONE] ' + clone_out.strip())
@@ -421,6 +425,10 @@ def diff(self, new_version, old_version, raw_output=False):
421425
cmd = 'git diff ' + old_version + ' ' + new_version
422426
p = subprocess.Popen(cmd, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
423427
(diff_out, diff_err) = p.communicate()
428+
429+
diff_out = diff_out.decode('utf-8')
430+
diff_err = diff_err.decode('utf-8')
431+
424432
logger.info(diff_out)
425433

426434
# change the work directory back.
@@ -448,6 +456,10 @@ def checkout(self, branch):
448456
cmd = "git fetch origin && git checkout " + branch
449457
p = subprocess.Popen(cmd, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
450458
(checkout_out, checkout_err) = p.communicate()
459+
460+
checkout_out = checkout_out.decode('utf-8')
461+
checkout_err = checkout_err.decode('utf-8')
462+
451463
logger.info('[PICKUP] [CHECKOUT] ' + checkout_err.strip())
452464

453465
# Already on
@@ -572,6 +584,10 @@ def __init__(self, filename, current_version=None, online_version=None):
572584
)
573585
p = subprocess.Popen(cmd, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
574586
(diff_out, diff_err) = p.communicate()
587+
588+
diff_out = diff_out.decode('utf-8')
589+
diff_err = diff_err.decode('utf-8')
590+
575591
if len(diff_err) == 0:
576592
logger.debug("[PICKUP] svn diff success")
577593
elif 'authorization failed' in diff_err:
@@ -589,6 +605,8 @@ def log(self):
589605
)
590606
p = subprocess.Popen(cmd, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
591607
log_out = p.communicate()[0]
608+
log_out = log_out.decode('utf-8')
609+
592610
return log_out
593611

594612
def diff(self):
@@ -601,6 +619,7 @@ def diff(self):
601619
)
602620
p = subprocess.Popen(cmd, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
603621
diff_out = p.communicate()[0]
622+
diff_out = diff_out.decode('utf-8')
604623

605624
added, removed, changed = [], [], []
606625
diff = {}

cobra/send_mail.py

+8-4
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,14 @@ def send_mail(target, filename, receiver):
3131

3232
msg.attach(MIMEText('扫描项目:{t}\n报告见附件'.format(t=target), 'plain', 'utf-8'))
3333

34-
with open(filename, 'rb') as f:
35-
attachment = MIMEApplication(f.read())
36-
attachment.add_header('Content-Disposition', 'attachment', filename=os.path.split(filename)[1])
37-
msg.attach(attachment)
34+
try:
35+
with open(filename, 'rb') as f:
36+
attachment = MIMEApplication(f.read())
37+
attachment.add_header('Content-Disposition', 'attachment', filename=os.path.split(filename)[1])
38+
msg.attach(attachment)
39+
except IOError:
40+
logger.warning('[EMAIL] No such file {}, please check input parameter'.format(filename))
41+
return False
3842

3943
try:
4044
server.login(user=username, password=password)

git_projects.py

+20-12
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@
2626
import Queue as queue
2727

2828

29+
git_urls = []
30+
31+
2932
def start():
3033
url = Config('git', 'gitlab_url').value
3134
private_token = Config('git', 'private_token').value
@@ -40,34 +43,37 @@ def start():
4043
q_pages.put(i + 1)
4144

4245
for i in range(10):
43-
thread = threading.Thread(target=get_git_urls, args=(url, private_token, cobra_ip, key, q_pages, fi))
46+
thread = threading.Thread(target=get_git_urls, args=(url, private_token, q_pages, fi))
4447
thread.start()
4548
threads.append(thread)
4649

4750
for thread in threads:
4851
thread.join()
4952

53+
res = push_to_api(git_urls, cobra_ip, key, fi)
54+
55+
if res:
56+
logger.info("Git push success: {}".format(len(git_urls)))
57+
else:
58+
logger.info("Git push fail")
59+
5060
fi.close()
5161
logger.info("All projects have been pushed")
5262

5363

54-
def get_git_urls(url, private_token, cobra_ip, key, q_pages, fi):
64+
def get_git_urls(url, private_token, q_pages, fi):
5565
"""
5666
:param url: The gitlab's projects api ,example:http://xxx.gitlab.com/api/v3/projects
5767
:param private_token: The user's private_token
58-
:param cobra_ip: The Cobra server's ip
59-
:param key: The Cobra api key
6068
:param q_pages: The Queue of pages
6169
:param fi: The result in this file
6270
:return:
6371
"""
6472
while not q_pages.empty():
65-
git_urls = []
6673
page = q_pages.get()
6774
params = {'private_token': private_token, 'page': page}
6875
url = url
6976
r = request_target(url, params, method="get")
70-
7177
if r.status_code == 200:
7278
data = r.json() # 一个页面中的Json数据,默认20条
7379
for j in range(len(data)):
@@ -80,12 +86,8 @@ def get_git_urls(url, private_token, cobra_ip, key, q_pages, fi):
8086
else:
8187
request_url = git_url
8288

89+
fi.write(request_url + '\n')
8390
git_urls.append(request_url)
84-
res = push_to_api(git_urls, cobra_ip, key, fi)
85-
if res:
86-
logger.info("page %d git push success" % page)
87-
else:
88-
logger.info("page %d git push fail" % page)
8991

9092
elif r.status_code == 404:
9193
logger.warning("page %d 404" % page)
@@ -107,12 +109,14 @@ def request_target(target_url, params=None, header=None, method="get"):
107109
def push_to_api(urls, cobra_ip, key, fi):
108110
headers = {"Content-Type": "application/json"}
109111
url = cobra_ip + "/api/add"
110-
payload = {"key": key, "target": urls}
112+
payload = {"key": key, "target": urls, "dels": True, "rule": "cvi-190009"}
111113
r = request_target(url, payload, headers, method="post")
112114
if r.status_code == 200:
113115
fi.write(str(r.json()) + '\n')
114116
logger.info(r.json())
115117
return True
118+
elif r.status_code == 404:
119+
logger.info("The page is 404")
116120
else:
117121
logger.info(r.json())
118122
return False
@@ -126,3 +130,7 @@ def get_pages(url, private_token):
126130
res = re.search(r"all\?page=(\d*)&per_page=0", res)
127131
pages = res.group(1)
128132
return pages
133+
134+
135+
if __name__ == '__main__':
136+
start()

0 commit comments

Comments
 (0)