Spaces:
Running
Running
import streamlit as st | |
import os | |
import pandas as pd | |
from utils import * | |
PATH = os.getcwd() | |
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 | |
# inference complete | |
st.snow() | |
print_code = input_code.split("\n")[:MAX_NUM_STATEMENTS] | |
st.markdown("### Scanning Results:") | |
if func_pred == 0: | |
st.write("<span style='color:green'>" + "No vulnerabilities detected"+ "</span>", unsafe_allow_html=True) | |
st.markdown("### Non-Vulnerable Function:") | |
else: | |
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) | |
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 | |
# inference complete | |
st.snow() | |
print_code = input_code.split("\n")[:MAX_NUM_STATEMENTS] | |
st.markdown("### Scanning Results:") | |
if func_pred == 0: | |
st.write("<span style='color:green'>" + "No vulnerabilities detected"+ "</span>", unsafe_allow_html=True) | |
st.markdown("### Non-Vulnerable Function:") | |
else: | |
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) |