main.py

import databases
import sqlalchemy

from fastapi import FastAPI, Query
from pydantic import BaseModel
from typing import List

DATABASE_URL = "sqlite:///./cves.db"

database = databases.Database(DATABASE_URL)

metadata = sqlalchemy.MetaData()

cves = sqlalchemy.Table(
    "cves",
    metadata,
    sqlalchemy.Column("cve_id", sqlalchemy.String),
    sqlalchemy.Column("description_eng", sqlalchemy.String),
    sqlalchemy.Column("affects", sqlalchemy.JSON),
    sqlalchemy.Column("cve_data_meta", sqlalchemy.JSON),
    sqlalchemy.Column("description", sqlalchemy.JSON),
    sqlalchemy.Column("problemtype", sqlalchemy.JSON),
    sqlalchemy.Column("cve_references", sqlalchemy.JSON),
    sqlalchemy.Column("data_format", sqlalchemy.String),
    sqlalchemy.Column("data_type", sqlalchemy.String),
    sqlalchemy.Column("data_version", sqlalchemy.String),
)

engine = sqlalchemy.create_engine(
    DATABASE_URL, connect_args={"check_same_thread": False}
)
metadata.create_all(engine)


# This was probably overkill, but I chose to split out the DB like this incase I want to do some graphql
class CVE(BaseModel):
    cve_id: str
    description_eng: str
    affects: dict
    cve_data_meta: dict
    description: dict
    problemtype: dict
    cve_references: dict
    data_format: str
    data_type: str
    data_version: str


# The output model removes cve_id and description_eng because I just use them to make life easier in db queries
class CVEOut(BaseModel):
    affects: dict
    cve_data_meta: dict
    description: dict
    problemtype: dict
    cve_references: dict
    data_format: str
    data_type: str
    data_version: str


app = FastAPI(
    title="Simple CVE API",
    description="Provides a simple API to query cvelist data",
    version="1.0.0",
)


@app.on_event("startup")
async def startup():
    await database.connect()


@app.on_event("shutdown")
async def shutdown():
    await database.disconnect()


@app.get("/v1/cve/{cve_id}", response_model=List[CVEOut])
async def cve_id_search(cve_id: str):
    query = cves.select().where(cves.c.cve_id == cve_id)
    return await database.fetch_all(query)


@app.get("/v1/search/", response_model=List[CVEOut])
async def cve_description_search(
    q: str = Query(
        ...,
        title="Query string",
        description="Wraps parameter in wildcards and searches the eng CVE descriptions. 100 result limit",
        min_length=1,
    )
):
    wildcard_q = "%{0}%".format(q)
    query = cves.select().where(cves.c.description.like(wildcard_q)).limit(100)
    return await database.fetch_all(query)