diff --git a/cobra/export.py b/cobra/export.py index 9330d707..7da1ec94 100644 --- a/cobra/export.py +++ b/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 = { diff --git a/cobra/pickup.py b/cobra/pickup.py index d5422ba3..a602827c 100644 --- a/cobra/pickup.py +++ b/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 = {} diff --git a/cobra/send_mail.py b/cobra/send_mail.py index bd02f744..0a630c41 100644 --- a/cobra/send_mail.py +++ b/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) diff --git a/git_projects.py b/git_projects.py index 88ff6954..4bcf9da4 100644 --- a/git_projects.py +++ b/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()