import json
import subprocess

from django.core.cache import cache


class SafetyCheck:
    def __init__(self):
        pass

    def get_check_result(self):
        commands = [
            # 安装或升级 safety 包
            "pip install --upgrade safety -i https://pypi.tuna.tsinghua.edu.cn/simple >> a.txt",
            # 运行 safety check 并输出为 JSON
            "safety check --json",
        ]
        result = {}
        # 逐个执行命令
        for cmd in commands:
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        return json.loads(result.stdout)

    def get_filter_result(self):
        stdout = self.get_check_result()
        # 过滤输出
        filter_result = {
            "timestamp": stdout["report_meta"]["timestamp"],  # 时间
            "pkg_count": stdout["report_meta"]["packages_found"],  # 扫描包数
            "vulnerabilities_count": stdout["report_meta"]["vulnerabilities_found"],  # 存在漏洞的包数
            "vulnerabilities": [
                # 问题包详情
                {
                    "name": vul["package_name"],
                    "vulnerable_spec": vul["vulnerable_spec"],
                    "analyzed_version": vul["analyzed_version"],
                    "advisory": vul["advisory"],  # 修复建议
                    "CVE": vul["CVE"],
                    "fixed_versions": vul["fixed_versions"],
                    "recommended_version": stdout["remediations"]
                    .get(vul["package_name"], {})
                    .get("recommended_version"),
                }
                for vul in stdout["vulnerabilities"]
            ],
        }

        return filter_result

    def check(self):
        result = self.get_filter_result()
        cache.set("checkpkg", result, 86400)
        cache.set("vulnerabilities_count", result.get("vulnerabilities_count", 0), 86400)
        return result
