Spaces:
Sleeping
Sleeping
| 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) | |