main.py

from copy import deepcopy
from collections import Counter

import numpy as np
import pandas as pd
import streamlit as st


@st.cache_data
def load_data() -> pd.DataFrame:
    data = pd.read_csv("Data/cve_data.csv")
    data["created_at"] = pd.to_datetime(data["created_at"])
    data["modified_at"] = pd.to_datetime(data["modified_at"])
    return data


def populate_second_column(col, data):
    with col:
        st.header("Vulnerabilities by Year")
        st.line_chart(data.groupby(data["created_at"].dt.year).count()["cve_id"])
        st.header("Vulnerabilities by Status")
        st.bar_chart(data=data.groupby("status").count()["cve_id"])


def populate_first_column(col, data):
    with col:
        st.header("Most Identified CWEs")
        cwe_counter = dict(sorted(Counter(data["cwe_id"]).items(), key=lambda x: x[1], reverse=True))
        cwe_counter = pd.DataFrame({"CWE ID": list(cwe_counter.keys())[:16], "Count": list(cwe_counter.values())[:16]})
        st.bar_chart(data=cwe_counter, x="CWE ID", y="Count")
        st.header("Most Vulnerable Products")
        data["cpe_version_id"] = ["_".join(x.split(":")[3:5]) if pd.isnull(x) is False else None for x in data["cpe_id"]]
        cpe_counter = dict(sorted(Counter(data["cpe_version_id"]).items(), key=lambda x: x[1], reverse=True))
        cpe_counter = pd.DataFrame({"CPE ID": list(cpe_counter.keys())[1:32], "Count": list(cpe_counter.values())[1:32]})
        st.bar_chart(data=cpe_counter, x="CPE ID", y="Count")



def main():
    st.set_page_config(page_title="Vulnerabilities in the Wild", page_icon="🐍", layout="wide")
    st.title("Vulnerabilities in the Wild")
    st.write("A look at the vulnerabilities that have been reported to the NVD over the years.")
    data = load_data()
    col1, col2 = st.columns(2)
    populate_first_column(col1, deepcopy(data))
    populate_second_column(col2, deepcopy(data))
    st.header("Dataset")
    st.dataframe(data)

if __name__ == "__main__":
    main()