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

Commit

Permalink
解决了测试用例不通过的问题
Browse files Browse the repository at this point in the history
  • Loading branch information
BlBana committed Mar 31, 2018
1 parent c4b71a1 commit a961418
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 7 deletions.
35 changes: 30 additions & 5 deletions cobra/cve.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import gzip
import xml.etree.cElementTree as eT
import multiprocessing
import subprocess
from .config import project_directory, Config, config_path
from .log import logger
from .dependencies import Dependencies
Expand Down Expand Up @@ -283,6 +284,24 @@ def download_rule_gz():
for t in threads:
t.join()
end_time = datetime.datetime.now()
for afile in files:
param = ['file', afile]
p = subprocess.Popen(param, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
res_out, res_err = p.communicate()

res_out = res_out.decode('utf-8')
res_err = res_err.decode('utf-8')

if 'HTML' in res_out:
os.remove(afile)
afile_name = os.path.split(afile)[1]
year = afile_name.split('.')[0]
url = "https://static.nvd.nist.gov/feeds/xml/cve/2.0/nvdcve-2.0-" + str(year) + ".xml.gz"
try:
urlretrieve(url, afile)
except IOError:
logger.warning('[CVE] The {} download fail'.format(afile))

logger.info("All CVE xml file already download success, use time:%ds" % (end_time - start_time).seconds)
return files

Expand All @@ -292,11 +311,17 @@ def un_gz(gz_files):
start_time = datetime.datetime.now()
logger.info("Start decompress rule files, Please wait a moment....")
for gz_file in gz_files:
f_name = gz_file.replace(".gz", "")
g_file = gzip.GzipFile(gz_file)
open(f_name, "wb+").write(g_file.read())
g_file.close()
os.remove(gz_file)
if os.path.exists(gz_file):
f_name = gz_file.replace(".gz", "")

try:
g_file = gzip.GzipFile(gz_file, "rb")
open(f_name, "wb+").write(g_file.read())
g_file.close()
except IOError:
logger.warning('[CVE] The {} download fail'.format(gz_file))

os.remove(gz_file)
end_time = datetime.datetime.now()
logger.info("Decompress success, use time:%ds" % (end_time - start_time).seconds)
return True
Expand Down
6 changes: 4 additions & 2 deletions tests/test_cve_parse.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,15 +112,17 @@ def test_download_rule_gz():
files = download_rule_gz()
assert isinstance(files, list)
for file_ in files:
os.remove(file_)
if os.path.exists(file_):
os.remove(file_)


def test_un_gz():
files = download_rule_gz()
res = un_gz(files)
assert res is True
for year in range(2002, datetime.datetime.now().year+1):
os.remove(project_directory+"/rules/%d.xml" % year)
if os.path.exists(project_directory+"/rules/%d.xml" % year):
os.remove(project_directory+"/rules/%d.xml" % year)


def test_rule_single():
Expand Down

0 comments on commit a961418

Please sign in to comment.