cvealert.py

#!/usr/bin/python3
# ------------------ conf start -----------------------------
# 群机器人的webhook写这里,支持钉钉/企业微信,一行一个
# 钉钉需要设置自定义关键词为`CVE`或IP白名单
webhooks = """\


""".split("\n")

# 关键词写这里
keywords = [
    # "windows",

]

# 消息通知模板(markdown)
# 模板可用的变量: $cveId, $description, $refs, $published, $sourceIdentifier
template = """\
**[$cveId](https://nvd.nist.gov/vuln/detail/$cveId)**
> $description

References:
> $refs

"""

# ------------------ conf end -------------------------------
from calendar import timegm
from string import Template
import requests
import time
import sys
import os
import re

headers = {
    "Referer": "https://github.com/p7e4/CVE-Alert",
    "User-Agent": "CVE-Alert"
}
keywordRegx = [re.compile(f"({i})", re.I) for i in keywords]


def getCve(lastRunTimestamp):
    now = time.strftime("%Y-%m-%dT%X.000", time.gmtime(time.time()))
    last = time.strftime("%Y-%m-%dT%X.000", time.gmtime(lastRunTimestamp))
    url = f"https://services.nvd.nist.gov/rest/json/cves/2.0/?resultsPerPage=2000&pubStartDate={last}&pubEndDate={now}"

    data = []
    templateVars = ("cveId", "description", "refs", "published", "sourceIdentifier")

    try:
        r = requests.get(url, timeout=30, headers=headers)
        vulnerabilities = r.json()["vulnerabilities"]
    except:
        print(f"[{r.status_code}] {r.url} {r.text}")
        raise

    for item in vulnerabilities:
        refs = []
        for ref in item["cve"]["references"]:
            refs.append(f"[{ref['url']}]({ref['url']})")

        refs = "\n".join(refs)
        cveId = item["cve"]["id"]
        description = item["cve"]["descriptions"][0]["value"].replace("\n\n", "\n").rstrip("\r\n")
        published = time.strftime("%Y-%m-%d %H:%M", time.gmtime(timegm(time.strptime(item["cve"]["published"][:-4], "%Y-%m-%dT%X")) + 60*60*8))
        sourceIdentifier = item["cve"]["sourceIdentifier"]

        if keywordRegx:
            khit = False
            for i in keywordRegx:
                if i.search(description):
                    khit = True
                    description = i.sub("**\\1**", description)
            if not khit: continue

        data.append(Template(template).safe_substitute({key: value for key, value in vars().items() if key in templateVars}))

    return data


def main():
    runtime = int(time.time())

    recordFile = f"{sys.path[0]}/.cvealert"

    if os.path.exists(recordFile):
        with open(recordFile) as f:
            sendMessages(getCve(int(f.read()) - 1))
    else:
        sendMessages(getCve(time.time() - 60 * 60 * 24)[:5])

    with open(recordFile, "w") as f:
        f.write(str(runtime))


def sendMessages(messages):
    # 钉钉 https://open.dingtalk.com/document/group/custom-robot-access
    # 企业微信 https://developer.work.weixin.qq.com/document/path/91770
    fq = 0
    s = requests.Session()
    for message in messages[::-1]:
        if len(message) > 4096:
            message = message[:4092] + "..."
        for webhook in webhooks:
            if not webhook: continue
            if webhook.startswith("https://oapi.dingtalk.com/robot/send?access_token="):
                send = {
                   "msgtype": "markdown",
                   "markdown": {
                        "title":"有新的推送消息",
                        "text": message
                    }
                }
            elif webhook.startswith("https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key="):
                send = {
                   "msgtype": "markdown",
                   "markdown": {
                        "content": message
                    }
                }
            else:
                exit(f"未知的通知地址: `{webhook}`")

            r = s.post(webhook, json=send, timeout=20)
            if r.json()["errcode"] != 0:
                print(f"发送消息失败: `{webhook}` {r.text}")
                print(f"message: {message}")

        fq += 1
        if fq % 20 == 0:
            time.sleep(60)


if __name__ == '__main__':
    main()