Update app.py

app.py

@@ -28,40 +28,44 @@ def extract_urls(text):
     return re.findall(url_pattern, text)
 
 # ---------------- UI START ----------------
-st.set_page_config(page_title="📩 Email Malicious Detector", layout="wide")
-st.markdown("<h1 style='text-align: center;'>📩 Malicious Email Detection App</h1>", unsafe_allow_html=True)
+def main():
+    st.set_page_config(page_title="📩 Email Malicious Detector", layout="wide")
+    st.markdown("<h1 style='text-align: center;'>📩 Malicious Email Detection App</h1>", unsafe_allow_html=True)
 
-st.markdown("### ✉️ Enter your email content:")
-email_text = st.text_area("Paste your email content here:", height=200)
+    st.markdown("### ✉️ Enter your email content:")
+    email_text = st.text_area("Paste your email content here:", height=200)
 
-if st.button("🚨 Scan Email & Analyze URL"):
-    if not email_text.strip():
-        st.warning("⚠️ Please input some email content.")
-    else:
-        result1 = pipe1(email_text)[0]
-        label1, score1 = result1['label'], result1['score']
-        pred1 = normalize_label(label1)
-
-        if pred1 == "benign":
-            st.markdown("## 🛡️ **Prediction Result:**")
-            st.success(f"✅ BENIGN EMAIL CONTENT (Confidence Score: {score1:.2%})")
+    if st.button("🚨 Scan Email & Analyze URL"):
+        if not email_text.strip():
+            st.warning("⚠️ Please input some email content.")
         else:
-            urls = extract_urls(email_text)
-            if not urls:
-                st.warning("⚠️ Email content is malicious, but no URL found for further analysis.")
-            else:
-                url = urls[0]
-                result2 = pipe2(url)[0]
-                result3 = pipe3(url)[0]
-                label2, score2 = result2['label'], result2['score']
-                label3, score3 = result3['label'], result3['score']
-
-                final_label, final_score = calculate_weighted_prediction(label2, score2, label3, score3)
+            result1 = pipe1(email_text)[0]
+            label1, score1 = result1['label'], result1['score']
+            pred1 = normalize_label(label1)
 
+            if pred1 == "benign":
                 st.markdown("## 🛡️ **Prediction Result:**")
-                if final_score < 0.6:
-                    st.warning(f"🤔 URLs in email content are UNCERTAIN - Confidence too low ({final_score:.2%}). Please review manually.")
-                elif final_label == "benign":
-                    st.success(f"✅ URLs in email content are BENIGN (Confidence Score: {final_score:.2%})")
+                st.success(f"✅ BENIGN EMAIL CONTENT (Confidence Score: {score1:.2%})")
+            else:
+                urls = extract_urls(email_text)
+                if not urls:
+                    st.warning("⚠️ Email content is malicious, but no URL found for further analysis.")
                 else:
-                    st.error(f"⚠️ URLs in email content are MALICIOUS (Confidence Score: {final_score:.2%})")
+                    url = urls[0]
+                    result2 = pipe2(url)[0]
+                    result3 = pipe3(url)[0]
+                    label2, score2 = result2['label'], result2['score']
+                    label3, score3 = result3['label'], result3['score']
+
+                    final_label, final_score = calculate_weighted_prediction(label2, score2, label3, score3)
+
+                    st.markdown("## 🛡️ **Prediction Result:**")
+                    if final_score < 0.6:
+                        st.warning(f"🤔 URLs in email content are UNCERTAIN - Confidence too low ({final_score:.2%}). Please review manually.")
+                    elif final_label == "benign":
+                        st.success(f"✅ URLs in email content are BENIGN (Confidence Score: {final_score:.2%})")
+                    else:
+                        st.error(f"⚠️ URLs in email content are MALICIOUS (Confidence Score: {final_score:.2%})")
+
+if __name__ == "__main__":
+    main()