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

Commit a961418

Browse files
committed
解决了测试用例不通过的问题
1 parent c4b71a1 commit a961418

File tree

2 files changed

+34
-7
lines changed

2 files changed

+34
-7
lines changed

cobra/cve.py

+30-5
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import gzip
1919
import xml.etree.cElementTree as eT
2020
import multiprocessing
21+
import subprocess
2122
from .config import project_directory, Config, config_path
2223
from .log import logger
2324
from .dependencies import Dependencies
@@ -283,6 +284,24 @@ def download_rule_gz():
283284
for t in threads:
284285
t.join()
285286
end_time = datetime.datetime.now()
287+
for afile in files:
288+
param = ['file', afile]
289+
p = subprocess.Popen(param, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
290+
res_out, res_err = p.communicate()
291+
292+
res_out = res_out.decode('utf-8')
293+
res_err = res_err.decode('utf-8')
294+
295+
if 'HTML' in res_out:
296+
os.remove(afile)
297+
afile_name = os.path.split(afile)[1]
298+
year = afile_name.split('.')[0]
299+
url = "https://static.nvd.nist.gov/feeds/xml/cve/2.0/nvdcve-2.0-" + str(year) + ".xml.gz"
300+
try:
301+
urlretrieve(url, afile)
302+
except IOError:
303+
logger.warning('[CVE] The {} download fail'.format(afile))
304+
286305
logger.info("All CVE xml file already download success, use time:%ds" % (end_time - start_time).seconds)
287306
return files
288307

@@ -292,11 +311,17 @@ def un_gz(gz_files):
292311
start_time = datetime.datetime.now()
293312
logger.info("Start decompress rule files, Please wait a moment....")
294313
for gz_file in gz_files:
295-
f_name = gz_file.replace(".gz", "")
296-
g_file = gzip.GzipFile(gz_file)
297-
open(f_name, "wb+").write(g_file.read())
298-
g_file.close()
299-
os.remove(gz_file)
314+
if os.path.exists(gz_file):
315+
f_name = gz_file.replace(".gz", "")
316+
317+
try:
318+
g_file = gzip.GzipFile(gz_file, "rb")
319+
open(f_name, "wb+").write(g_file.read())
320+
g_file.close()
321+
except IOError:
322+
logger.warning('[CVE] The {} download fail'.format(gz_file))
323+
324+
os.remove(gz_file)
300325
end_time = datetime.datetime.now()
301326
logger.info("Decompress success, use time:%ds" % (end_time - start_time).seconds)
302327
return True

tests/test_cve_parse.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -112,15 +112,17 @@ def test_download_rule_gz():
112112
files = download_rule_gz()
113113
assert isinstance(files, list)
114114
for file_ in files:
115-
os.remove(file_)
115+
if os.path.exists(file_):
116+
os.remove(file_)
116117

117118

118119
def test_un_gz():
119120
files = download_rule_gz()
120121
res = un_gz(files)
121122
assert res is True
122123
for year in range(2002, datetime.datetime.now().year+1):
123-
os.remove(project_directory+"/rules/%d.xml" % year)
124+
if os.path.exists(project_directory+"/rules/%d.xml" % year):
125+
os.remove(project_directory+"/rules/%d.xml" % year)
124126

125127

126128
def test_rule_single():

0 commit comments

Comments
 (0)