Merge pull request #923 from BlBana/master

 解决了Python3下代码拉取bug

Author: FeeiCN

cobra/export.py

@@ -144,24 +144,28 @@ def write_to_file(target, sid, output_format='', filename=None):
         logger.info('Vulnerabilities\n' + str(dict_to_pretty_table(scan_data.get('vulnerabilities'))))
 
     elif output_format == 'json' or output_format == 'JSON':
-        if not os.path.exists(filename):
-            with open(filename, 'w', encoding='utf-8') as f:
-                json_data = {
-                    sid: scan_data,
-                }
-                f.write(dict_to_json(json_data))
-        else:
-            with open(filename, 'r+', encoding='utf-8') as f:
-                try:
-                    json_data = json.load(f)
-                    json_data.update({sid: scan_data})
-                    # 使用 r+ 模式不会覆盖，调整文件指针到开头
-                    f.seek(0)
-                    f.truncate()
+        try:
+            if not os.path.exists(filename):
+                with open(filename, 'w', encoding='utf-8') as f:
+                    json_data = {
+                        sid: scan_data,
+                    }
                     f.write(dict_to_json(json_data))
-                except ValueError:
-                    logger.warning('[EXPORT] The json file have invaild token or None: {}'.format(os.path.join(export_path, filename)))
-                    return False
+            else:
+                with open(filename, 'r+', encoding='utf-8') as f:
+                    try:
+                        json_data = json.load(f)
+                        json_data.update({sid: scan_data})
+                        # 使用 r+ 模式不会覆盖，调整文件指针到开头
+                        f.seek(0)
+                        f.truncate()
+                        f.write(dict_to_json(json_data))
+                    except ValueError:
+                        logger.warning('[EXPORT] The json file have invaild token or None: {}'.format(os.path.join(export_path, filename)))
+                        return False
+        except IOError:
+            logger.warning('[EXPORT] Please input a file path after the -o parameter')
+            return False
 
     elif output_format == 'xml' or output_format == 'XML':
         xml_data = {

cobra/pickup.py

@@ -387,6 +387,10 @@ def clone(self):
 
         p = subprocess.Popen(cmd, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
         (clone_out, clone_err) = p.communicate()
+
+        clone_out = clone_out.decode('utf-8')
+        clone_err = clone_err.decode('utf-8')
+
         clone_err = clone_err.replace('{0}:{1}'.format(self.repo_username, self.repo_password), '')
 
         logger.debug('[PICKUP] [CLONE] ' + clone_out.strip())
@@ -421,6 +425,10 @@ def diff(self, new_version, old_version, raw_output=False):
         cmd = 'git diff ' + old_version + ' ' + new_version
         p = subprocess.Popen(cmd, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
         (diff_out, diff_err) = p.communicate()
+
+        diff_out = diff_out.decode('utf-8')
+        diff_err = diff_err.decode('utf-8')
+
         logger.info(diff_out)
 
         # change the work directory back.
@@ -448,6 +456,10 @@ def checkout(self, branch):
         cmd = "git fetch origin && git checkout " + branch
         p = subprocess.Popen(cmd, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
         (checkout_out, checkout_err) = p.communicate()
+
+        checkout_out = checkout_out.decode('utf-8')
+        checkout_err = checkout_err.decode('utf-8')
+
         logger.info('[PICKUP] [CHECKOUT] ' + checkout_err.strip())
 
         # Already on
@@ -572,6 +584,10 @@ def __init__(self, filename, current_version=None, online_version=None):
         )
         p = subprocess.Popen(cmd, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
         (diff_out, diff_err) = p.communicate()
+
+        diff_out = diff_out.decode('utf-8')
+        diff_err = diff_err.decode('utf-8')
+
         if len(diff_err) == 0:
             logger.debug("[PICKUP] svn diff success")
         elif 'authorization failed' in diff_err:
@@ -589,6 +605,8 @@ def log(self):
         )
         p = subprocess.Popen(cmd, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
         log_out = p.communicate()[0]
+        log_out = log_out.decode('utf-8')
+
         return log_out
 
     def diff(self):
@@ -601,6 +619,7 @@ def diff(self):
         )
         p = subprocess.Popen(cmd, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
         diff_out = p.communicate()[0]
+        diff_out = diff_out.decode('utf-8')
 
         added, removed, changed = [], [], []
         diff = {}

cobra/send_mail.py

@@ -31,10 +31,14 @@ def send_mail(target, filename, receiver):
 
     msg.attach(MIMEText('扫描项目：{t}\n报告见附件'.format(t=target), 'plain', 'utf-8'))
 
-    with open(filename, 'rb') as f:
-        attachment = MIMEApplication(f.read())
-        attachment.add_header('Content-Disposition', 'attachment', filename=os.path.split(filename)[1])
-        msg.attach(attachment)
+    try:
+        with open(filename, 'rb') as f:
+            attachment = MIMEApplication(f.read())
+            attachment.add_header('Content-Disposition', 'attachment', filename=os.path.split(filename)[1])
+            msg.attach(attachment)
+    except IOError:
+        logger.warning('[EMAIL] No such file {}, please check input parameter'.format(filename))
+        return False
 
     try:
         server.login(user=username, password=password)

git_projects.py

@@ -26,6 +26,9 @@
     import Queue as queue
 
 
+git_urls = []
+
+
 def start():
     url = Config('git', 'gitlab_url').value
     private_token = Config('git', 'private_token').value
@@ -40,34 +43,37 @@ def start():
         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 = threading.Thread(target=get_git_urls, args=(url, private_token, q_pages, fi))
         thread.start()
         threads.append(thread)
 
     for thread in threads:
         thread.join()
 
+    res = push_to_api(git_urls, cobra_ip, key, fi)
+
+    if res:
+        logger.info("Git push success: {}".format(len(git_urls)))
+    else:
+        logger.info("Git push fail")
+
     fi.close()
     logger.info("All projects have been pushed")
 
 
-def get_git_urls(url, private_token, cobra_ip, key, q_pages, fi):
+def get_git_urls(url, private_token, 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:
     """
     while not q_pages.empty():
-        git_urls = []
         page = q_pages.get()
         params = {'private_token': private_token, 'page': page}
         url = url
         r = request_target(url, params, method="get")
-
         if r.status_code == 200:
             data = r.json()  # 一个页面中的Json数据，默认20条
             for j in range(len(data)):
@@ -80,12 +86,8 @@ def get_git_urls(url, private_token, cobra_ip, key, q_pages, fi):
                 else:
                     request_url = git_url
 
+                fi.write(request_url + '\n')
                 git_urls.append(request_url)
-            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:
             logger.warning("page %d 404" % page)
@@ -107,12 +109,14 @@ def request_target(target_url, params=None, header=None, method="get"):
 def push_to_api(urls, cobra_ip, key, fi):
     headers = {"Content-Type": "application/json"}
     url = cobra_ip + "/api/add"
-    payload = {"key": key, "target": urls}
+    payload = {"key": key, "target": urls, "dels": True, "rule": "cvi-190009"}
     r = request_target(url, payload, headers, method="post")
     if r.status_code == 200:
         fi.write(str(r.json()) + '\n')
         logger.info(r.json())
         return True
+    elif r.status_code == 404:
+        logger.info("The page is 404")
     else:
         logger.info(r.json())
         return False
@@ -126,3 +130,7 @@ def get_pages(url, private_token):
     res = re.search(r"all\?page=(\d*)&per_page=0", res)
     pages = res.group(1)
     return pages
+
+
+if __name__ == '__main__':
+    start()