from PDAPI import *
import math

__metadata__ = {
    "methods": ["GET"],
    "blueprint": "SVE",
    "permission": 0.0,
    "attributes": [],
    "limit": "1 per second",
    "args": [
        {"Name": "title", "Type": "str", "MaxLen": 256},
        {"Name": "cvssSeverity", "Type": "str", "AllowList": ["CRITICAL", "HIGH", "MEDIUM", "LOW", ""], "Required": False},
        {"Name": "publishYear", "Type": "int", "AllowList": [2025, 2024, 2023, 2022, 2021], "Required": False},
        {"Name": "cvssAv", "Type": "str", "AllowList": ["N", "A", "L", "P", ""], "Required": False},
        {"Name": "cvssAc", "Type": "str", "AllowList": ["L", "H", ""], "Required": False},
        {"Name": "cvssPr", "Type": "str", "AllowList": ["N", "L", "H", ""], "Required": False},
        {"Name": "page", "Type": "int", "Required": False},
        {"Name": "pageSize", "Type": "int", "MaxLen": 100, "Required": False},
    ]
}

def PreSVE(SVEID: str):
    SVEID = str(SVEID).replace("SVE-", "").replace("-", "")
    if not len(SVEID) == 9: return False
    for i in SVEID:
        if i not in "0123456789": return False
    SVEID = "0" + SVEID
    return SVEID

def PreSQLSVE(SVEID: str):
    return SVEID[1:]

def HandleRequest():
    title = flask.request.args.get("title", "")
    cvssSeverity = flask.request.args.get("cvssSeverity", "")
    publishYear = flask.request.args.get("publishYear", type=int)
    cvssAv = flask.request.args.get("cvssAv", "")
    cvssAc = flask.request.args.get("cvssAc", "")
    cvssPr = flask.request.args.get("cvssPr", "")
    page = flask.request.args.get("page", 1, type=int)
    pageSize = flask.request.args.get("pageSize", 10, type=int)
    pageSize = min(pageSize, 100)
    query = SVEBase.query
    if title:
        title_like = f"%{title}%"
        query = query.filter(
            db.or_(
                SVEBase.SVEName.like(title_like),
                SVEBase.SVEDescription.like(title_like)
            )
        )
    if publishYear:
        query = query.filter(db.func.year(SVEBase.SVECreateTime) == publishYear)
    if cvssSeverity and cvssSeverity != "":
        severity_map = {"CRITICAL": (9.0, 10.0), "HIGH": (7.0, 8.9), "MEDIUM": (4.0, 6.9), "LOW": (0.1, 3.9)}
        if cvssSeverity in severity_map:
            min_score, max_score = severity_map[cvssSeverity]
            query = query.filter(SVEBase.SVECvss["score"].cast(db.Float).between(min_score, max_score))
    if cvssAv and cvssAv != "":
        query = query.filter(SVEBase.SVECvss["av"].cast(db.String) == cvssAv)
    if cvssAc and cvssAc != "":
        query = query.filter(SVEBase.SVECvss["ac"].cast(db.String) == cvssAc)
    if cvssPr and cvssPr != "":
        query = query.filter(SVEBase.SVECvss["pr"].cast(db.String) == cvssPr)
    total_count = query.count()
    total_pages = math.ceil(total_count / pageSize) if pageSize else 0
    page = max(1, min(page, total_pages)) if total_pages else 1
    offset = (page - 1) * pageSize
    results = query.offset(offset).limit(pageSize).all()
    data = []
    for item in results:
        sve_id_hex = item.SVEID.hex()
        sve_id_num = str(int(sve_id_hex))
        sve_id = f"SVE-{sve_id_num[:4]}-{sve_id_num[4:]}"
        cvss_data = item.SVECvss or {}
        publish_date = item.SVECreateTime.strftime("%Y-%m-%d") if item.SVECreateTime else ""
        data.append({
            "sveId": sve_id,
            "vulnTitle": item.SVEName or "",
            "cvss": {
                "score": str(cvss_data.get("score", "")),
                "severity": cvss_data.get("severity", "")
            },
            "publishDate": publish_date
        })
    return flask.jsonify({
        "code": 0,
        "data": data,
        "page": page,
        "pageSize": pageSize,
        "totalPages": total_pages,
        "totalCount": total_count
    })