From 2fca05a18ffbe1f679408cce242466956c3e3c45 Mon Sep 17 00:00:00 2001
From: kim <kim@if010.com>
Date: Wed, 2 Jul 2025 11:54:20 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E4=BA=86=E4=B8=80=E4=BA=9B?=
 =?UTF-8?q?=E9=87=87=E9=9B=86=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

修复：关于POC库被删除，采集是报404问题，将采集adysec
修复：采集nvd时出现403问题，增加头部信息
修复：采集时报错会跳出循环，增加了`try...except...`
优化：增加日志记录，会在执行的当前路径输出athena.log文件记录每次运行采集的结果
---
 collectors/collector_poc.py | 4 ++--
 collectors/utils.py         | 6 +++++-
 main.py                     | 8 ++++++--
 processing/filter.py        | 7 ++++++-
 4 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/collectors/collector_poc.py b/collectors/collector_poc.py
index 4d8825f..72781d6 100644
--- a/collectors/collector_poc.py
+++ b/collectors/collector_poc.py
@@ -8,7 +8,7 @@
 
 class POCCollector(VulnerabilityCollector):
     def __init__(self):
-        poc_url = "https://api.github.com/repos/wy876/POC/commits"
+        poc_url = "https://api.github.com/repos/adysec/POC/commits"
         poc_headers = {
             "Authorization": f"token {cfg['github']['token']}",
         }
@@ -42,7 +42,7 @@ def parse_data(self, raw_data):
                             'description': description,
                             'source': self.source_name,
                             'date': datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-                            'link': 'https://github.com/wy876/POC/blob/main/' + file['filename']
+                            'link': 'https://github.com/adysec/POC/blob/main/' + file['filename']
                         }
                         vulnerabilities_dict[vulnerability['name']] = vulnerability
 
diff --git a/collectors/utils.py b/collectors/utils.py
index c930e78..f130712 100644
--- a/collectors/utils.py
+++ b/collectors/utils.py
@@ -16,8 +16,12 @@ def extract_score(soup, selector):
 
 def get_cve_info(cve):
     url = "https://nvd.nist.gov/vuln/detail/" + cve
+    headers = {
+        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
+        'Referer': 'https://nvd.nist.gov/'
+    }
     try:
-        res = requests.get(url)
+        res = requests.get(url, headers=headers)
         res.raise_for_status()
     except requests.RequestException as e:
         print(f"Error establishing connection: {e}")
diff --git a/main.py b/main.py
index e7ba129..81e4019 100644
--- a/main.py
+++ b/main.py
@@ -38,8 +38,12 @@ def main():
             daily_task()
             last_sent_date = current_date
 
-        vulnerabilities = gather_data()
-        filter_high_risk_vuls(vulnerabilities)
+        try:
+            vulnerabilities = gather_data()
+            filter_high_risk_vuls(vulnerabilities)
+        except Exception as e:
+            with open('athena.log', 'a') as f:
+                print(f"{current_time}\tError: {e}", file=f)
         time.sleep(600)
 
 
diff --git a/processing/filter.py b/processing/filter.py
index 08a39bc..41e6b44 100644
--- a/processing/filter.py
+++ b/processing/filter.py
@@ -1,4 +1,5 @@
 import html
+import datetime
 from tqdm import tqdm
 from config import cfg
 from database.db_class import MySQLDatabase
@@ -47,5 +48,9 @@ def filter_high_risk_vuls(vulnerabilities):
                     content += f"{key}: {value}\n"
             content = content.rstrip("\n")
             send_realtime_notifications(content)
-    print(f"Selected severity vulnerabilities found: {selected_severity_num}\n\n")
+
+    current_time = datetime.datetime.now()
+    with open('athena.log', 'a') as f:
+        print(f"{current_time}\tSelected severity vulnerabilities found: {selected_severity_num}", file=f)
+    print(f"{current_time}\tSelected severity vulnerabilities found: {selected_severity_num}")