Update app.py

app.py

@@ -252,43 +252,73 @@ def calculate_malware_risk(features):
 # -------------------------------
 def get_final_prediction(phishing_pred, malware_pred, phishing_risk, malware_risk):
     """
-    Simple risk-based decision system with whitelist protection:
-    1. Check for trusted domains first
-    2. Compare risk scores directly
-    3. Use higher risk score for final prediction
-    4. Add safety thresholds for benign classification
+    Enhanced decision system:
+    1. Prioritize model predictions first
+    2. Use risk scores for confidence and tie-breaking
+    3. Whitelist protection for trusted domains
     """
     
-    # Safety thresholds
-    THREAT_THRESHOLD = 30  # Increased minimum score to consider as threat
-    HIGH_CONFIDENCE_THRESHOLD = 15  # Risk difference for high confidence
+    # Trusted domains whitelist (exact match)
+    trusted_domains = [
+        'google.com', 'www.google.com', 'facebook.com', 'www.facebook.com',
+        'microsoft.com', 'www.microsoft.com', 'apple.com', 'www.apple.com',
+        'amazon.com', 'www.amazon.com', 'youtube.com', 'www.youtube.com',
+        'twitter.com', 'www.twitter.com', 'linkedin.com', 'www.linkedin.com',
+        'github.com', 'www.github.com', 'stackoverflow.com', 'www.stackoverflow.com'
+    ]
+    
+    # Extract domain from URL for whitelist check
+    from urllib.parse import urlparse
+    try:
+        parsed_url = urlparse(url if 'url' in locals() else "")
+        domain = parsed_url.netloc.lower()
+        if domain in trusted_domains:
+            return "Benign", f"Whitelisted trusted domain: {domain}"
+    except:
+        pass
     
-    # Case 1: Both risks are very low - definitely benign
-    if phishing_risk < THREAT_THRESHOLD and malware_risk < THREAT_THRESHOLD:
-        return "Benign", f"Low risk scores (Phishing: {phishing_risk}, Malware: {malware_risk}) - Safe to browse"
+    # Model prediction priorities
+    RISK_BOOST_THRESHOLD = 15  # Minimum risk to boost model prediction
     
-    # Case 2: One or both risks are above threshold
-    risk_difference = abs(phishing_risk - malware_risk)
+    # Case 1: Both models detect threats
+    if phishing_pred == "Phishing" and malware_pred == "malicious":
+        if phishing_risk > malware_risk:
+            return "Phishing", f"Both models detected threat - phishing characteristics stronger (Risk: {phishing_risk} vs {malware_risk})"
+        else:
+            return "Malicious", f"Both models detected threat - malware characteristics stronger (Risk: {malware_risk} vs {phishing_risk})"
     
-    if phishing_risk > malware_risk:
-        if phishing_risk >= THREAT_THRESHOLD:
-            confidence = "High" if risk_difference >= HIGH_CONFIDENCE_THRESHOLD else "Medium"
-            return "Phishing", f"Phishing risk higher ({phishing_risk} vs {malware_risk}) - {confidence} confidence"
+    # Case 2: Only phishing model detects threat
+    elif phishing_pred == "Phishing" and malware_pred != "malicious":
+        if phishing_risk >= RISK_BOOST_THRESHOLD or phishing_risk > malware_risk:
+            return "Phishing", f"Phishing model detected threat with supporting risk indicators (Risk: {phishing_risk})"
         else:
-            return "Benign", f"Phishing risk slightly higher but below threat threshold ({phishing_risk})"
+            return "Phishing", f"Phishing model detected threat (Risk score: {phishing_risk})"
     
-    elif malware_risk > phishing_risk:
-        if malware_risk >= THREAT_THRESHOLD:
-            confidence = "High" if risk_difference >= HIGH_CONFIDENCE_THRESHOLD else "Medium"
-            return "Malicious", f"Malware risk higher ({malware_risk} vs {phishing_risk}) - {confidence} confidence"
+    # Case 3: Only malware model detects threat  
+    elif malware_pred == "malicious" and phishing_pred != "Phishing":
+        if malware_risk >= RISK_BOOST_THRESHOLD or malware_risk > phishing_risk:
+            return "Malicious", f"Malware model detected threat with supporting risk indicators (Risk: {malware_risk})"
         else:
-            return "Benign", f"Malware risk slightly higher but below threat threshold ({malware_risk})"
+            return "Malicious", f"Malware model detected threat (Risk score: {malware_risk})"
     
-    else:  # Equal risks
-        if phishing_risk >= THREAT_THRESHOLD:
-            return "Suspicious", f"Equal risk scores ({phishing_risk}) - requires manual review"
+    # Case 4: Both models report benign - check high risk scores
+    else:
+        HIGH_RISK_THRESHOLD = 40  # High risk threshold for override
+        MEDIUM_RISK_THRESHOLD = 25  # Medium risk threshold
+        
+        if phishing_risk >= HIGH_RISK_THRESHOLD and malware_risk >= HIGH_RISK_THRESHOLD:
+            if phishing_risk > malware_risk:
+                return "Phishing", f"Models missed but high phishing risk detected ({phishing_risk})"
+            else:
+                return "Malicious", f"Models missed but high malware risk detected ({malware_risk})"
+        elif phishing_risk >= HIGH_RISK_THRESHOLD:
+            return "Phishing", f"Models reported benign but high phishing risk indicators ({phishing_risk})"
+        elif malware_risk >= HIGH_RISK_THRESHOLD:
+            return "Malicious", f"Models reported benign but high malware risk indicators ({malware_risk})"
+        elif phishing_risk >= MEDIUM_RISK_THRESHOLD or malware_risk >= MEDIUM_RISK_THRESHOLD:
+            return "Suspicious", f"Models reported benign but moderate risk present (P:{phishing_risk}, M:{malware_risk})"
         else:
-            return "Benign", f"Equal low risk scores ({phishing_risk}) - Safe to browse"
+            return "Benign", f"Models and risk analysis confirm safe (P:{phishing_risk}, M:{malware_risk})"
 
 def analyze_url(url):
     try:
@@ -308,7 +338,7 @@ def analyze_url(url):
         phishing_risk = calculate_phishing_risk(phishing_features)
         malware_risk = calculate_malware_risk(malware_features)
         
-        # Get final prediction using simple risk-based system
+        # Get final prediction using enhanced model-priority system
         final_result, decision_reason = get_final_prediction(
             phishing_pred, malware_pred, phishing_risk, malware_risk
         )