Spaces:
Sleeping
Sleeping
File size: 6,024 Bytes
f513a95 0700577 20e24b1 0700577 f513a95 20e24b1 f513a95 20e24b1 0700577 f513a95 20e24b1 f513a95 20e24b1 f513a95 20e24b1 f513a95 20e24b1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
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)
st.snow()
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]
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)
st.snow() |