Upload app.py

app.py

@@ -12,6 +12,7 @@ def show_cwe_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"
@@ -63,23 +64,23 @@ if __name__ == "__main__":
                     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:")                                
+                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:
-                    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:")
+                    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
@@ -100,17 +101,19 @@ if __name__ == "__main__":
                 # inference complete
                 st.snow()
                 print_code = input_code.split("\n")[:MAX_NUM_STATEMENTS]
-                st.markdown("### Scanning Results:")
                 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:
-                    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)
+                    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()