Merge branch 'master' of https://github.com/wufeifei/cobra

cobra/cve_parse.py

@@ -13,14 +13,19 @@
 import datetime
 import os
 import requests
-import urllib
 import threading
 import gzip
 import xml.etree.cElementTree as eT
 import multiprocessing
 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:
@@ -52,7 +57,7 @@ def cve_parse(self):
         if not isinstance(cve_file, list):
             tree = self.parse_xml(cve_file)
             root = tree.getroot()
-            childs = root.iter('%sentry' % self.NS)
+            childs = root.findall('.//%sentry' % self.NS)
             for child in childs:  # child is entry Element
                 cve_id = child.attrib['id']
                 cve_info = self.cve_info(child)
@@ -62,7 +67,7 @@ def cve_parse(self):
             for filename in cve_file:
                 tree = self.parse_xml(filename)
                 root = tree.getroot()
-                childs = root.iter('%sentry' % self.NS)
+                childs = root.findall('.//%sentry' % self.NS)
                 for child in childs:  # child is entry Element
                     cve_id = child.attrib['id']
                     cve_info = self.cve_info(child)
@@ -97,7 +102,7 @@ def cve_info(self, entry):
                        'wu-ftpd', 'cluster_server', 'catos', 'mantis', 'quicktime', 'security_linux', 'firefox',
                        'jetty_http_server', 'php:', 'enterprise_linux', 'oracle10g', 'oracle9g', 'oracle8g', 'firehol',
                        'fetchmail', 'postgresql', 'freebsd', 'chrome']
-        products = entry.iter('%sproduct' % self.VULN)
+        products = entry.findall('.//%sproduct' % self.VULN)
         access_complexity = entry.find('.//%saccess-complexity' % self.CVSS)
         for product in products:
             module_version = product.text.split(':')
@@ -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):
         """
@@ -179,7 +185,7 @@ def rule_parse(self, file_):
         """
         tree = self.parse_xml(file_)
         root = tree.getroot()
-        cves = root.iter('cve')
+        cves = root.findall('.//cve')
         for cve_child in cves:
             cve_id = cve_child.attrib['id']
             cve_level = cve_child.attrib['level']
@@ -191,7 +197,7 @@ def rule_parse(self, file_):
     def rule_info(cve_child):
         rule_info = {}
         cpe_list = []
-        products = cve_child.iter('product')
+        products = cve_child.findall('.//product')
         for product in products:
             cpe_list.append(product.text.lower())
         rule_info['cpe'] = cpe_list
@@ -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)')
 
@@ -246,14 +252,14 @@ def rule_parse():
     if is_update():
         gz_files = download_rule_gz()
         un_gz(gz_files)
-        pool = multiprocessing.Pool(processes=50)
-        for year in range(2002, datetime.datetime.now().year+1):
+        pool = multiprocessing.Pool()
+        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

cobra/dependencies.py

@@ -91,7 +91,7 @@ def find_java_mvn(self, file_path):
         for pom in file_path:
             tree = self.parse_xml(pom)
             root = tree.getroot()
-            childs = root.iter('%sdependency' % pom_ns)
+            childs = root.findall('.//%sdependency' % pom_ns)
             for child in childs:
                 group_id = child.getchildren()[0].text
                 artifact_id = child.getchildren()[1].text

cobra/git_projects.py

@@ -13,64 +13,115 @@
 """
 import json
 import requests
+import re
+import threading
+from .log import logger
+from .config import code_path, Config
+try:
+    # Python 3
+    import queue
+except ImportError:
+    # Python 2
+    import Queue as queue
 
 
-def start(url, private_token, cobra_ip, key):
+def start():
+    url = Config('git', 'gitlab_url').value
+    private_token = Config('git', 'private_token').value
+    cobra_ip = Config('git', 'cobra_ip').value
+    key = Config('cobra', 'secret_key').value
+    threads = []
+    pages = get_pages(url, private_token)
+    q_pages = queue.Queue(pages)
+    result_path = code_path + '/result_sid'
+    fi = open(result_path, 'w+')
+    for i in range(int(pages)):
+        q_pages.put(i+1)
+
+    for i in range(10):
+        thread = threading.Thread(target=get_git_urls, args=(url, private_token, cobra_ip, key, q_pages, fi))
+        thread.start()
+        threads.append(thread)
+
+    for thread in threads:
+        thread.join()
+
+    fi.close()
+    logger.info("All projects have been pushed")
+
+
+def get_git_urls(url, private_token, cobra_ip, key, q_pages, fi):
     """
     :param url: The gitlab's projects api ,example:http://xxx.gitlab.com/api/v3/projects
     :param private_token: The user's private_token
     :param cobra_ip: The Cobra server's ip
     :param key: The Cobra api key
+    :param q_pages: The Queue of pages
+    :param fi: The result in this file
     :return:
     """
-    page = 1
-    while True:
+    while not q_pages.empty():
         git_urls = []
+        page = q_pages.get()
         params = {'private_token': private_token, 'page': page}
         url = url
-        r = requests.get(url, params=params)
+        r = request_target(url, params, method="get")
 
         if r.status_code == 200:
             data = r.json()  # 一个页面中的Json数据，默认20条
-            if len(data) == 0:
-                print("url收集完毕")
-                break
             for j in range(len(data)):
                 git_url = data[j]['http_url_to_repo']
                 git_branch = data[j]['default_branch']
 
                 if git_branch is not None:
                     request_url = git_url+':'+git_branch
+
                 else:
                     request_url = git_url
+
                 git_urls.append(request_url)
-            # res = push_to_api(git_urls, cobra_ip, key)
-            # if res:
-            #     print ("page %d git push success" % page)
-            # else:
-            #     print ("page %d git push fail" % page)
+            res = push_to_api(git_urls, cobra_ip, key, fi)
+            if res:
+                logger.info("page %d git push success" % page)
+            else:
+                logger.info("page %d git push fail" % page)
 
         elif r.status_code == 404:
-            print("404")
+            logger.warning("page %d 404" % page)
 
         else:
-            print(r.status_code)
+            logger.warning("page %d is %d" % page, r.status_code)
+        q_pages.task_done()
+
 
-        page += 1
+def request_target(target_url, params=None, header=None, method="get"):
+    if method == "get":
+        response = requests.get(url=target_url, params=params, headers=header)
+        return response
+    if method == "post":
+        response = requests.post(url=target_url, data=json.dumps(params), headers=header)
+        return response
 
 
-def push_to_api(urls, cobra_ip, key):
+def push_to_api(urls, cobra_ip, key, fi):
     headers = {"Content-Type": "application/json"}
     url = cobra_ip + "/api/add"
     payload = {"key": key, "target": urls}
-    r = requests.post(url=url, data=json.dumps(payload), headers=headers)
+    r = request_target(url, payload, headers, method="post")
     if r.status_code == 200:
-        print(r.json())
+        fi.write(str(r.json()) + '\n')
+        logger.info(r.json())
         return True
     else:
-        print(r.json())
+        logger.info(r.json())
         return False
 
 
-def get_pages():
-    pass
+def get_pages(url, private_token):
+    params = {"private_token": private_token}
+    response = request_target(url, params)
+    res = response.headers['link'].split(",")
+    res = res[2]
+    res = re.search(r"all\?page=(\d*)&per_page=0", res)
+    pages = res.group(1)
+    return pages

cobra/utils.py

@@ -18,6 +18,7 @@
 import string
 import random
 import hashlib
+import locale
 from .log import logger
 from .config import Config
 from .exceptions import PickupException, NotExistException, AuthFailedException
@@ -178,20 +179,32 @@ def convert_time(seconds):
     if minute == 0:
         return str(seconds % one_minute) + "\""
     else:
-        return str(minute) + "'" + str(seconds % one_minute) + "\""
+        return str(int(minute)) + "'" + str(seconds % one_minute) + "\""
 
 
-def convert_number(number):
+def convert_number(n):
     """
     Convert number to , split
     Ex: 123456 -> 123,456
-    :param number:
+    :param n:
     :return:
     """
-    if number is None or number == 0:
-        return 0
-    number = int(number)
-    return '{:20,}'.format(number).strip()
+    if n is None: return '0'
+    n = str(n)
+    if '.' in n:
+        dollars, cents = n.split('.')
+    else:
+        dollars, cents = n, None
+
+    r = []
+    for i, c in enumerate(str(dollars)[::-1]):
+        if i and (not (i % 3)):
+            r.insert(0, ',')
+        r.insert(0, c)
+    out = ''.join(r)
+    if cents:
+        out += '.' + cents
+    return out
 
 
 def md5(content):
@@ -200,6 +213,7 @@ def md5(content):
     :param content:
     :return:
     """
+    content = content.encode('utf8')
     return hashlib.md5(content).hexdigest()
 
 
@@ -230,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])))
@@ -258,8 +273,9 @@ def path_to_file(path):
     :return:
     """
     paths = path.split('/')
-    paths = filter(None, paths)
-    return '.../{0}'.format(paths[len(paths) - 1])
+    paths = list(filter(None, paths))
+    length = len(paths)
+    return '.../{0}'.format(paths[length - 1])
 
 
 def percent(part, whole, need_per=True):

config.template

@@ -53,6 +53,11 @@ password:
 [git]
 username:
 password:
+private_token:
+# http://xxx.gitlab.com/api/v3/projects
+gitlab_url:
+cobra_ip:
+
 
 [cve]
 modified = 00B6A630D5515C9D6F61B2FC97ED7EE689A787AC52E1654D1E991BC26F4647CD