app.py

import streamlit as st
import os
import pandas as pd
from utils import *

PATH = os.getcwd()

def show_cwe_sev(input_code):
    # more predictions
    cwe_pred = predict_cweid([input_code])["cwe_id"][0]
    sev_pred = predict_sev([input_code])
    sev_score = int(sev_pred["batch_sev_score"][0])
    sev_class = sev_pred["batch_sev_class"][0]
    
    st.markdown("### Scanning Results:")
    if "CWE-" in cwe_pred:
        id_ = cwe_pred.split("CWE-")[-1]
        link = f"https://cwe.mitre.org/data/definitions/{id_}.html"
        st.markdown(f'**CWE ID: <a href="{link}">{cwe_pred}</a>**', unsafe_allow_html=True)
    if sev_score < 4:
        st.markdown(f"**Severity: <span style='color:blue'>{sev_class}</span>**", unsafe_allow_html=True)
        st.markdown(f"**Severity Score: <span style='color:blue'>{sev_score}</span>**", unsafe_allow_html=True)
    elif sev_score < 7:
        st.markdown(f"**Severity: <span style='color:orange'>{sev_class}</span>**", unsafe_allow_html=True)
        st.markdown(f"**Severity Score: <span style='color:orange'>{sev_score}</span>**", unsafe_allow_html=True)
    elif sev_score < 9:
        st.markdown(f"**Severity: <span style='color:red'>{sev_class}</span>**", unsafe_allow_html=True)
        st.markdown(f"**Severity Score: <span style='color:red'>{sev_score}</span>**", unsafe_allow_html=True)
    else:
        st.markdown(f"**Severity: <span style='color:red'>{sev_class}</span>**", unsafe_allow_html=True)
        st.markdown(f"**Severity Score: <span style='color:red'>{sev_score}</span>**", unsafe_allow_html=True)
        

if __name__ == "__main__":
    MAX_NUM_STATEMENTS = 155
    
    st.set_page_config(page_title="AIBugHunter")
    # sidebar
    st.sidebar.title("AIBugHunter Web App")
    behavior = st.sidebar.selectbox(label="NAVIGATOR IS HERE:",
                                    options=["DEMO", "Analyze my own"])
    if behavior == "DEMO":
        # function title
        st.title("C/C++ Vulnerability Dataset Viewer")
        dataset_path = PATH + "/data/test.csv"
        st.dataframe(pd.read_csv(dataset_path))
        
        with st.form("input_form_a"):
            idx = st.selectbox('Select an index', (str(i) for i in range(100)))
            sub = st.form_submit_button("Select")
            if sub:
                idx = int(idx)
                df = pd.read_csv(dataset_path)
                input_code = df["function"][idx]
                
                input_code = input_code.split("\n")[:MAX_NUM_STATEMENTS]
                input_code = "\n".join(input_code)
                # load model
                with st.spinner("Scanning security issues..."):
                    # do inference
                    out = predict_vul_lines([input_code])
                    func_pred = out["batch_func_pred"][0]
                    func_confidence = out["batch_func_pred_prob"][0]
                    line_pred = out["batch_statement_pred"][0]
                    line_confidence = out["batch_statement_pred_prob"][0]      
                output = None
                print_code = input_code.split("\n")[:MAX_NUM_STATEMENTS]               
                if func_pred == 0:
                    st.markdown("### Scanning Results:")
                    st.markdown("<span style='color:green'>" + "**No vulnerabilities detected**"+ "</span>", unsafe_allow_html=True)
                    st.markdown("### Non-Vulnerable Function:")
                else:
                    with st.spinner("Identifying vulnerability types and severity..."):
                        show_cwe_sev(input_code)
                        for i in range(len(print_code)):
                            c = print_code[i]
                            vul = line_pred[i]
                            if vul == 1:
                                st.markdown(f"<span style='color:red'> Vulnerable Line **{i+1}** </span>", unsafe_allow_html=True)
                                st.code(c)
                        st.markdown("### Vulnerable Function:")
                st.code(input_code, language="cpp", line_numbers=True)
                    
    elif behavior == "Analyze my own":
        # user input of project title
        ## todo- limit the input to 150 lines
        with st.form("input_form_b"):
            input_code = st.text_area("Input a C/C++ function:", height=275)
            submitted = st.form_submit_button("Analyze")
            if submitted:
                # load model
                with st.spinner("Scanning security issues..."):
                    # do inference
                    out = predict_vul_lines([input_code])
                    func_pred = out["batch_func_pred"][0]
                    func_confidence = out["batch_func_pred_prob"][0]
                    line_pred = out["batch_statement_pred"][0]
                    line_confidence = out["batch_statement_pred_prob"][0]      
                output = None
                print_code = input_code.split("\n")[:MAX_NUM_STATEMENTS]
                if func_pred == 0:
                    st.markdown("### Scanning Results:")
                    st.write("<span style='color:green'>" + "No vulnerabilities detected"+ "</span>", unsafe_allow_html=True)
                    st.markdown("### Non-Vulnerable Function:")
                else:
                    with st.spinner("Identifying vulnerability types and severity..."):
                        show_cwe_sev(input_code)
                        for i in range(len(print_code)):
                            c = print_code[i]
                            vul = line_pred[i]
                            if vul == 1:
                                st.write(f"<span style='color:red'> Vulnerable Line {i+1} </span>", unsafe_allow_html=True)
                                st.code(c)
                        st.markdown("### Vulnerable Function:")   
                st.code(input_code, language="cpp", line_numbers=True)