main.py

import rtoml
from pathlib import Path
from modules.info_push import post_message_to_WXWork
from modules.cve_MS import getMSDATA
from modules.cmdb import (
    Base,
    engine,
)
from modules.atk_tools import ATKTools
from modules.followed_users import FollowedUsers
from modules.cve_github import CVEGithub
from rocketry import Rocketry
from rocketry.conds import daily, hourly
import asyncio
from rocketry.conds import cron
from modules.atom_tool import InitAtomFile, add_atom_entry
import html


# 读取配置项信息
CONFIG = rtoml.load(Path(__file__).parent / "config.toml")
GITHUB_TOKEN = CONFIG["github_token"]
ATOM_FILE_PATH = Path(__file__).parent / "releases.atom"

app = Rocketry()

print("程序已载入完成, 启动中...")
print("正在初始化数据库...")
Base.metadata.create_all(engine)
print("数据库初始化完成, 开始企微Bot连接测试...")
# post_message_to_WXWork("Github-CVE监控-企微Bot连接测试")
print("企微Bot连接测试完成")
# 初始化 Atom 文件
InitAtomFile(ATOM_FILE_PATH)
add_atom_entry(
    atom_file_path=ATOM_FILE_PATH,
    entry_title="CVE monitor feed",
    entry_link="https://github.com/Ayusummer/cve-monitor-wxworkbot",
    entry_id="tag:github.com,2024:https://github.com/Ayusummer/cve-monitor-wxworkbot/20240305",
    entry_summary="testsummary",
    entry_author_name="Ayusummer",
)

# 初始化 Github 关注用户仓库更新监控类
followedUsers = FollowedUsers()
# 初始化 Github CVE 新仓库监控类
cveGithub = CVEGithub()
# 初始化 Github 关注仓库监控类
atkTools = ATKTools()
# getMSDATA()


@app.task(cron("0 9-19/2 * * 1-5"))
def normal_task_2_hour():
    """每个工作日的9-19点之间的每两个小时执行一次 CVE_Github 仓库监控 + 关注用户仓库监控 + 关注仓库监控 + MS-CVE监控"""
    followedUsers.get_users_repo()
    cveGithub.get_2day_cve_info()
    atkTools.get_tools_info()
    getMSDATA()


@app.task(cron("0 9 * * 1-5"))
def github_cve_task_1_day_push_and_clear_high_freq_org_repo():
    """每个工作日的 9:00 对 Github CVE 监控中的高频 CVE 信息进行一次推送"""
    cveGithub.push_and_clear_high_freq_user_org_repo_push_info_list()


if __name__ == "__main__":
    app.run()