diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..c8e733c --- /dev/null +++ b/.dockerignore @@ -0,0 +1,5 @@ +*.md +Dockerfile +tests/ +test.py +README/ \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..c65faf4 --- /dev/null +++ b/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"] \ No newline at end of file diff --git a/README.md b/README.md index fde39d3..149d411 100644 --- a/README.md +++ b/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: [] diff --git a/collectors/collector_afrog.py b/collectors/collector_afrog.py index 549b369..c4795b1 100644 --- a/collectors/collector_afrog.py +++ b/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) diff --git a/collectors/collector_msf.py b/collectors/collector_msf.py index 3234851..cf74cb0 100644 --- a/collectors/collector_msf.py +++ b/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 diff --git a/collectors/collector_vulhub.py b/collectors/collector_vulhub.py index 1cce50d..90ec587 100644 --- a/collectors/collector_vulhub.py +++ b/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() diff --git a/collectors/decorators.py b/collectors/decorators.py index 5a4d399..ad5b896 100644 --- a/collectors/decorators.py +++ b/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 diff --git a/config.py b/config.py index 398a038..e631a4d 100644 --- a/config.py +++ b/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() diff --git a/config.yaml b/config.yaml index 902ae26..1d2093c 100644 --- a/config.yaml +++ b/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: [] diff --git a/main.py b/main.py index e7ba129..f8273eb 100644 --- a/main.py +++ b/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__": diff --git a/notifications/feishu.py b/notifications/feishu.py index 81fd866..0d92b01 100644 --- a/notifications/feishu.py +++ b/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) diff --git a/requirements.txt b/requirements.txt index ccca9ec..e7c2280 100644 --- a/requirements.txt +++ b/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