增加镜像打包和从环境读取变量覆盖配置文件中的配置项

.dockerignore

@@ -0,0 +1,5 @@
+*.md
+Dockerfile
+tests/
+test.py
+README/

Dockerfile

@@ -0,0 +1,21 @@
+FROM python:3.8.10-slim-buster
+ENV TZ="Asia/Shanghai"
+# 修改更新源, 设置时区
+RUN echo "deb https://mirrors.aliyun.com/debian/ buster main non-free contrib" > /etc/apt/sources.list \
+    && echo "deb-src https://mirrors.aliyun.com/debian/ buster main non-free contrib" >> /etc/apt/sources.list \
+    && echo "deb https://mirrors.aliyun.com/debian-security buster/updates main" >> /etc/apt/sources.list \
+    && echo "deb-src https://mirrors.aliyun.com/debian-security buster/updates main" >> /etc/apt/sources.list \
+    && echo "deb https://mirrors.aliyun.com/debian/ buster-updates main non-free contrib" >> /etc/apt/sources.list \
+    && echo "deb-src https://mirrors.aliyun.com/debian/ buster-updates main non-free contrib" >> /etc/apt/sources.list \
+    && apt-get update \
+    && apt-get install -y --no-install-recommends tzdata gcc libmariadb-dev \
+    && ln -snf /usr/share/zoneinfo/$TZ /etc/localtime  \
+    && echo $TZ > /etc/timezone
+# 设置工作目录, 安装依赖
+WORKDIR /app
+COPY requirements.txt .
+RUN pip install -r requirements.txt -i http://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com
+# 拷贝代码
+COPY . .
+# 启动
+ENTRYPOINT ["python", "main.py"]

README.md

@@ -90,9 +90,15 @@ pip install -r requirements.txt
 config.yaml
 
 ```yaml
+global:
+  max_retries: 3
+  delay: 3
+  timeout: 30
+
 #https://github.com/settings/tokens
 github:
   token: ""
+  proxy: ""
 
 #collectors默认为空，表示爬取所有漏洞源信息，如需指定特定源，可修改此项.可选项为['POC','Afrog','PacketStorm','Github','Seebug','OSCS','Ali','QAX','ThreatBook','Vulhub','MSF','ExploitDB']
 collectors: []

collectors/collector_afrog.py

@@ -24,7 +24,10 @@ def fetch_data(self, timeout):
     @staticmethod
     @retry()
     def extract_info(file_path, timeout):
-        url = f"https://raw.githubusercontent.com/zan8in/afrog/master/{file_path}"
+        if cfg['github']['proxy'] == '':
+            url = f"https://raw.githubusercontent.com/zan8in/afrog/master/{file_path}"
+        else:
+            url = f"{cfg['github']['proxy']}https://raw.githubusercontent.com/zan8in/afrog/master/{file_path}"
         response = requests.get(url, timeout=timeout)
         body = response.text
         yaml_content = yaml.safe_load(body)

collectors/collector_msf.py

@@ -30,7 +30,10 @@ def fetch_data(self, timeout):
     @staticmethod
     @retry()
     def extract_info(file_path, timeout):
-        url = f"https://raw.githubusercontent.com/rapid7/metasploit-framework/master/{file_path}"
+        if cfg['github']['proxy'] == '':
+            url = f"https://raw.githubusercontent.com/rapid7/metasploit-framework/master/{file_path}"
+        else:
+            url = f"{cfg['github']['proxy']}https://raw.githubusercontent.com/rapid7/metasploit-framework/master/{file_path}"
         response = requests.get(url, timeout=timeout)
         body = response.text
 

collectors/collector_vulhub.py

@@ -24,7 +24,10 @@ def fetch_data(self, timeout):
 
     @retry()
     def extract_name(self, file_path, timeout):
-        url = f"https://raw.githubusercontent.com/vulhub/vulhub/master/{file_path}/README.zh-cn.md"
+        if cfg['github']['proxy'] == '':
+            url = f"https://raw.githubusercontent.com/vulhub/vulhub/master/{file_path}/README.zh-cn.md"
+        else:
+            url = f"{cfg['github']['proxy']}https://raw.githubusercontent.com/vulhub/vulhub/master/{file_path}/README.zh-cn.md"
         response = requests.get(url, headers=self.headers, timeout=timeout)
         body = response.text
         title = body.split('\n', 1)[0].split('#')[1].strip()

collectors/decorators.py

@@ -3,10 +3,18 @@
 from functools import wraps
 import time
 
+from config import cfg
 from notifications.notifier import send_realtime_notifications
 
 
 def retry(max_retries=3, delay=2, timeout=10):
+    if cfg['global']['max_retries']:
+        max_retries = cfg['global']['max_retries']
+    if cfg['global']['delay']:
+        delay = cfg['global']['delay']
+    if cfg['global']['timeout']:
+        timeout = cfg['global']['timeout']
+
     def decorator(func):
         @wraps(func)
         def wrapper(self, *args, **kwargs):
@@ -18,7 +26,10 @@ def wrapper(self, *args, **kwargs):
                     print(f"Request failed: {e}. Retrying {retries + 1}/{max_retries}...")
                     retries += 1
                     time.sleep(delay)
-            msg = f"fail to fetch {self.source_name} data due to network error"
+            if isinstance(self, str):
+                msg = f"fail to fetch {self} data due to network error"
+            else:
+                msg = f"fail to fetch {self.source_name} data due to network error"
             send_realtime_notifications(msg)
             return None
 

config.py

@@ -1,3 +1,5 @@
+import os
+
 import yaml
 import os
 
@@ -8,9 +10,40 @@ def load_config(config_path=None):
     if config_path is None:
         config_path = os.path.join(config_dir, 'config.yaml')
 
-    with open(config_path, 'r') as f:
+    with open(config_path, 'r', encoding='UTF-8') as f:
         config = yaml.safe_load(f)
+        update_by_env(config)
     return config
 
 
+def update_by_env(config):
+    traverse_dict(config)
+
+
+def traverse_dict(config, obj=None, parent_key=None):
+    if obj is None:
+        obj = config
+    if isinstance(obj, dict):
+        for key, value in obj.items():
+            if parent_key is not None:
+                key_list = parent_key.copy()
+            else:
+                key_list = []
+            key_list.append(key)
+            if not isinstance(value, dict):
+                update_config_key_from_env(config, key=key_list)
+            else:
+                traverse_dict(config=config, obj=value, parent_key=key_list)
+
+
+def update_config_key_from_env(config, key):
+    env_key = '_'.join(key).upper()
+    value = os.getenv(env_key)
+    update_config = config
+    for k in key[:-1]:
+        update_config = update_config[k]
+    if value is not None:
+        update_config[key[-1]] = value
+
+
 cfg = load_config()

config.yaml

@@ -1,6 +1,12 @@
+global:
+  max_retries: 3
+  delay: 3
+  timeout: 30
+
 #https://github.com/settings/tokens
 github:
   token: ""
+  proxy: ""
 
 #collectors默认为空，表示爬取所有漏洞源信息，如需指定特定源，可修改此项.可选项为['POC','Afrog','PacketStorm','Github','Seebug','OSCS','Ali','QAX','ThreatBook','Vulhub','MSF','ExploitDB']
 collectors: []

main.py

@@ -31,16 +31,20 @@ def main():
 
     last_sent_date = None
     while True:
-        current_time = datetime.datetime.now()
-        current_date = current_time.date()
-
-        if current_time.hour == 6 and last_sent_date != current_date:
-            daily_task()
-            last_sent_date = current_date
-
-        vulnerabilities = gather_data()
-        filter_high_risk_vuls(vulnerabilities)
-        time.sleep(600)
+        try:
+            current_time = datetime.datetime.now()
+            current_date = current_time.date()
+
+            if current_time.hour == 6 and last_sent_date != current_date:
+                daily_task()
+                last_sent_date = current_date
+
+            vulnerabilities = gather_data()
+            filter_high_risk_vuls(vulnerabilities)
+            time.sleep(600)
+        except Exception as e:
+            print(f"Error: {e}")
+            time.sleep(60)
 
 
 if __name__ == "__main__":

notifications/feishu.py

@@ -1,9 +1,10 @@
-import hashlib
 import base64
+import hashlib
 import hmac
-import requests
 import time
 
+import requests
+
 
 def gen_sign(timestamp, secret):
     string_to_sign = '{}\n{}'.format(timestamp, secret)
@@ -15,16 +16,16 @@ def gen_sign(timestamp, secret):
 def feishu_notification(webhook, secret, content):
     webhook_url = 'https://open.feishu.cn/open-apis/bot/v2/hook/' + webhook
     timestamp = str(int(time.time()))
-    sign = gen_sign(timestamp, secret)
     headers = {
         "Content-Type": "application/json"
     }
     data = {
         "timestamp": timestamp,
-        "sign": sign,
         "msg_type": "text",
         "content": {
             "text": content
         }
     }
+    if secret:
+        data['sign'] = gen_sign(timestamp, secret)
     requests.post(webhook_url, json=data, headers=headers)

requirements.txt

@@ -1,6 +1,7 @@
 Flask==2.2.5
 art==6.2
 mysql==0.0.3
+mysqlclient==2.1.1
 mysql-connector-python==9.0.0
 python-dateutil==2.8.2
 PyYAML==6.0.1