diff --git a/cobra/cve_parse.py b/cobra/cve_parse.py index 8f9ed393..cc6e6187 100644 --- a/cobra/cve_parse.py +++ b/cobra/cve_parse.py @@ -13,7 +13,6 @@ import datetime import os import requests -import urllib import threading import gzip import xml.etree.cElementTree as eT @@ -21,6 +20,12 @@ from .config import project_directory, Config, config_path from .log import logger from .dependencies import Dependencies + +try: + from urllib import urlretrieve # Python2 +except ImportError: + from urllib.request import urlretrieve # Python3 + try: from configparser import ConfigParser except ImportError: @@ -158,7 +163,8 @@ def rule_xml(self): rule_path = project_directory + '/rules/CVI-999' tree.write(rule_path + str(self.year)[1:] + '.xml') endtime = datetime.datetime.now() - logger.info('CVE-999' + str(self.year)[1:] + '.xml Rule update succeeds, times:%ds' % (endtime - starttime).seconds) + logger.info( + 'CVE-999' + str(self.year)[1:] + '.xml Rule update succeeds, times:%ds' % (endtime - starttime).seconds) def pretty(self, e, level=0): """ @@ -234,7 +240,7 @@ def log_result(self): for cve_child in self._scan_result[module_]: cve_id = cve_child level = self._scan_result[module_][cve_id] - logger.warning('Find the module ' + module_ + ' have ' + cve_id +',level: ' +level) + logger.warning('Find the module ' + module_ + ' have ' + cve_id + ',level: ' + level) count = len(self._scan_result[module_]) logger.warning('The ' + module_ + ' module have ' + str(count) + ' CVE Vul(s)') @@ -247,13 +253,13 @@ def rule_parse(): gz_files = download_rule_gz() un_gz(gz_files) pool = multiprocessing.Pool() - for year in range(2002, datetime.datetime.now().year+1): + for year in range(2002, datetime.datetime.now().year + 1): cve_xml = "../rules/%d.xml" % year pool.apply_async(rule_single, args=(cve_xml, year)) pool.close() pool.join() - for year in range(2002, datetime.datetime.now().year+1): - os.remove(project_directory+"/rules/%d.xml" % year) + for year in range(2002, datetime.datetime.now().year + 1): + os.remove(project_directory + "/rules/%d.xml" % year) logger.info("The rule update success, start scan cve vuls") return True else: @@ -264,18 +270,19 @@ def download_rule_gz(): threads = [] files = [] start_time = datetime.datetime.now() - for year in range(2002, datetime.datetime.now().year+1): + for year in range(2002, datetime.datetime.now().year + 1): url = "https://static.nvd.nist.gov/feeds/xml/cve/2.0/nvdcve-2.0-" + str(year) + ".xml.gz" logger.info("start download " + str(year) + ".xml.gz") - thread = threading.Thread(target=urllib.urlretrieve, args=(url, project_directory+"/rules/"+str(year)+".xml.gz")) + thread = threading.Thread(target=urlretrieve, + args=(url, project_directory + "/rules/" + str(year) + ".xml.gz")) thread.start() threads.append(thread) logger.info('CVE-' + str(year) + " is download success") - files.append(project_directory+"/rules/" + str(year) + ".xml.gz") + files.append(project_directory + "/rules/" + str(year) + ".xml.gz") for t in threads: t.join() end_time = datetime.datetime.now() - logger.info("All CVE xml file already download success, use time:%ds" % (end_time-start_time).seconds) + logger.info("All CVE xml file already download success, use time:%ds" % (end_time - start_time).seconds) return files @@ -286,11 +293,11 @@ def un_gz(gz_files): for gz_file in gz_files: f_name = gz_file.replace(".gz", "") g_file = gzip.GzipFile(gz_file) - open(f_name, "w+").write(g_file.read()) + open(f_name, "wb+").write(g_file.read()) g_file.close() os.remove(gz_file) end_time = datetime.datetime.now() - logger.info("Decompress success, use time:%ds" % (end_time-start_time).seconds) + logger.info("Decompress success, use time:%ds" % (end_time - start_time).seconds) return True @@ -300,9 +307,9 @@ def rule_single(target_directory, year): def is_update(): url = "https://static.nvd.nist.gov/feeds/xml/cve/2.0/nvdcve-2.0-modified.meta" - r = requests.get(url) + r = requests.get(url, verify=False) index = r.text.find('sha256:') - sha256_now = r.text[index+7:].strip() + sha256_now = r.text[index + 7:].strip() sha256_local = Config(level1='cve', level2='modified').value if sha256_local != sha256_now: logger.info("The CVE Rule already update, start update local rule") @@ -314,7 +321,7 @@ def is_update(): config.write(fi) fi.close() except IOError as e: - logger.warning(e.message) + logger.warning(e) logger.info("The sha256 been update") return True return False diff --git a/cobra/utils.py b/cobra/utils.py index 5e1d18c8..2c370e43 100644 --- a/cobra/utils.py +++ b/cobra/utils.py @@ -213,6 +213,7 @@ def md5(content): :param content: :return: """ + content = content.encode('utf8') return hashlib.md5(content).hexdigest() @@ -243,6 +244,7 @@ def path_to_short(path, max_length=36): return path paths = path.split('/') paths = filter(None, paths) + paths = list(paths) tmp_path = '' for i in range(0, len(paths)): logger.debug((i, str(paths[i]), str(paths[len(paths) - i - 1]))) diff --git a/tests/test_cve_parse.py b/tests/test_cve_parse.py index f729b32f..35cbb568 100644 --- a/tests/test_cve_parse.py +++ b/tests/test_cve_parse.py @@ -18,7 +18,7 @@ from cobra.cve_parse import * from cobra.cve_parse import CveParse, project_directory try: - from configparser import ConfigParser + from configparser import ConfigParser, NoSectionError except ImportError: from ConfigParser import ConfigParser, NoSectionError diff --git a/tests/test_detection.py b/tests/test_detection.py index bb74dd4c..a312d7dd 100644 --- a/tests/test_detection.py +++ b/tests/test_detection.py @@ -50,7 +50,8 @@ def test_get_dict(): extension = ['php', 'js', 'java'] type_num = {} type_num = detection.get_dict(extension, type_num) - assert isinstance(extension, type(type_num.keys())) + print(type(type_num)) + assert type_num['php']['blank'] == 0 def test_project_information(): @@ -61,37 +62,32 @@ def test_project_information(): def test_count_py_line(): count = Detection.count_py_line(examples_path+'/cloc.py') - type_ = count.keys() type_count = ['count_blank', 'count_code', 'count_pound'] - assert type_ == type_count + assert count['count_code'] == 5 def test_count_php_line(): count = Detection.count_php_line(examples_path+'/cloc.php') - type_ = count.keys() type_count = ['count_blank', 'count_code', 'count_pound'] - assert type_ == type_count + assert count['count_code'] == 2 def test_count_java_line(): count = Detection.count_java_line(examples_path+'/cloc.java') - type_ = count.keys() type_count = ['count_blank', 'count_code', 'count_pound'] - assert type_ == type_count + assert count['count_code'] == 1 def test_count_html_line(): count = Detection.count_html_line(examples_path+'/cloc.html') - type_ = count.keys() type_count = ['count_blank', 'count_code', 'count_pound'] - assert type_ == type_count + assert count['count_code'] == 9 def test_count_data_line(): count = Detection.count_data_line(examples_path+'/param_xml.xml') - type_ = count.keys() type_count = ['count_blank', 'count_code', 'count_pound'] - assert type_ == type_count + assert count['count_code'] == 81 def test_countnum():