updated to improve google safebrowsing scoring

app.py

@@ -7,10 +7,16 @@ from langdetect import detect
 from deep_translator import GoogleTranslator
 import openai
 import os
+import requests
+import json
 
 # Set your OpenAI API key
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
+# Retrieve Google Safe Browsing API key from environment
+SAFE_BROWSING_API_KEY = os.getenv("GOOGLE_SAFE_BROWSING_API_KEY")
+SAFE_BROWSING_URL = "https://safebrowsing.googleapis.com/v4/threatMatches:find"
+
 # Translator instance
 translator = GoogleTranslator(source="auto", target="es")
 
@@ -26,6 +32,68 @@ model_name = "joeddav/xlm-roberta-large-xnli"
 classifier = pipeline("zero-shot-classification", model=model_name)
 CANDIDATE_LABELS = ["SMiShing", "Other Scam", "Legitimate"]
 
+def check_urls_with_google_safebrowsing(urls):
+    """
+    Queries the Google Safe Browsing API to see if any URLs are malicious.
+    Returns a dict {url: bool is_malicious}.
+    If the API key is missing or an error occurs, returns {url: False} for all.
+    """
+    result = {}
+    if not SAFE_BROWSING_API_KEY:
+        # No key: we can't check. Return all as safe.
+        for u in urls:
+            result[u] = False
+        return result
+
+    # Build threatEntries for each URL
+    threat_entries = [{"url": u} for u in urls]
+
+    payload = {
+        "client": {
+            "clientId": "my-smishing-detector",
+            "clientVersion": "1.0"
+        },
+        "threatInfo": {
+            "threatTypes": [
+                "MALWARE",
+                "SOCIAL_ENGINEERING",
+                "UNWANTED_SOFTWARE",
+                "POTENTIALLY_HARMFUL_APPLICATION"
+            ],
+            "platformTypes": ["ANY_PLATFORM"],
+            "threatEntryTypes": ["URL"],
+            "threatEntries": threat_entries
+        }
+    }
+
+    try:
+        resp = requests.post(
+            SAFE_BROWSING_URL,
+            params={"key": SAFE_BROWSING_API_KEY},
+            json=payload,
+            timeout=10
+        )
+        data = resp.json()
+        # If "matches" is present, some URL is flagged
+        # Each match has "threat": {"url": "..."}
+        malicious_urls = set()
+        if "matches" in data:
+            for match in data["matches"]:
+                threat_url = match.get("threat", {}).get("url")
+                if threat_url:
+                    malicious_urls.add(threat_url)
+
+        for u in urls:
+            result[u] = (u in malicious_urls)
+
+    except Exception as e:
+        print(f"Error contacting Safe Browsing API: {e}")
+        # default: everything safe if error
+        for u in urls:
+            result[u] = False
+
+    return result
+
 def get_keywords_by_language(text: str):
     """
     Detect language using langdetect and translate keywords if needed.
@@ -49,7 +117,8 @@ def get_keywords_by_language(text: str):
 
 def boost_probabilities(probabilities: dict, text: str):
     """
-    Boost probabilities based on keyword matches and presence of URLs.
+    Boost probabilities based on keyword matches, presence of URLs,
+    and Google Safe Browsing checks.
     """
     lower_text = text.lower()
     smishing_keywords, other_scam_keywords, detected_lang = get_keywords_by_language(text)
@@ -60,7 +129,12 @@ def boost_probabilities(probabilities: dict, text: str):
     smishing_boost = 0.30 * smishing_count
     other_scam_boost = 0.30 * other_scam_count
 
-    found_urls = re.findall(r"(https?://[^\s]+|\b(?:[a-zA-Z0-9.-]+\.(?:com|net|org|edu|gov|mil|io|ai|co|info|biz|us|uk|de|fr|es|ru|jp|cn|in|au|ca|br|mx|it|nl|se|no|fi|ch|pl|kr|vn|id|tw|sg|hk))\b)", lower_text)
+    # More robust URL pattern
+    found_urls = re.findall(
+        r"(https?://[^\s]+|\b[a-zA-Z0-9.-]+\.(?:com|net|org|edu|gov|mil|io|ai|co|info|biz|us|uk|de|fr|es|ru|jp|cn|in|au|ca|br|mx|it|nl|se|no|fi|ch|pl|kr|vn|id|tw|sg|hk)\b)",
+        lower_text
+    )
+    # If any URL is found, add 0.35 to smishing
     if found_urls:
         smishing_boost += 0.35
 
@@ -72,12 +146,11 @@ def boost_probabilities(probabilities: dict, text: str):
     p_other_scam += other_scam_boost
     p_legit -= (smishing_boost + other_scam_boost)
 
-    # Clamp
+    # Preliminary clamp & normalization
     p_smishing = max(p_smishing, 0.0)
     p_other_scam = max(p_other_scam, 0.0)
     p_legit = max(p_legit, 0.0)
 
-    # Re-normalize
     total = p_smishing + p_other_scam + p_legit
     if total > 0:
         p_smishing /= total
@@ -86,6 +159,16 @@ def boost_probabilities(probabilities: dict, text: str):
     else:
         p_smishing, p_other_scam, p_legit = 0.0, 0.0, 1.0
 
+    # **Now** check Safe Browsing. If any URL is malicious => p_smishing=1.0
+    if found_urls:
+        malicious_results = check_urls_with_google_safebrowsing(found_urls)
+        # If any malicious => set p_smishing=1.0
+        if any(malicious_results[u] for u in malicious_results):
+            # Bump SMiShing to max
+            p_smishing = 1.0
+            p_other_scam = 0.0
+            p_legit = 0.0
+
     return {
         "SMiShing": p_smishing,
         "Other Scam": p_other_scam,
@@ -119,7 +202,6 @@ def query_llm_for_classification(raw_message: str) -> dict:
         )
         raw_reply = response["choices"][0]["message"]["content"].strip()
 
-        import json
         llm_result = json.loads(raw_reply)
         if "label" not in llm_result or "reason" not in llm_result:
             return {"label": "Unknown", "reason": f"Unexpected format: {raw_reply}"}
@@ -176,17 +258,14 @@ def query_llm_for_explanation(
     Second LLM call: provides a holistic explanation of the final classification
     in the same language as detected_lang (English or Spanish).
     """
-    # Decide the language for final explanation
     if detected_lang == "es":
-        # Spanish
         system_prompt = (
             "Eres un experto en ciberseguridad. Proporciona una explicación final al usuario en español. "
             "Combina la clasificación local, la clasificación LLM y la etiqueta final en una sola explicación breve. "
             "No reveles el código interno ni el JSON bruto; simplemente da una breve explicación fácil de entender. "
-            "Termina con la etiqueta final. "
+            "Termina con la etiqueta final."
         )
     else:
-        # Default to English
         system_prompt = (
             "You are a cybersecurity expert providing a final explanation to the user in English. "
             "Combine the local classification, the LLM classification, and the final label "
@@ -259,7 +338,6 @@ def smishing_detector(input_type, text, image):
     boosted = boost_probabilities(original_probs, combined_text)
     detected_lang = boosted.pop("detected_lang", "en")
 
-    # Convert to float only
     for k in boosted:
         boosted[k] = float(boosted[k])
 
@@ -274,7 +352,6 @@ def smishing_detector(input_type, text, image):
     # 4. Incorporate LLM’s label into final probabilities
     boosted = incorporate_llm_label(boosted, llm_label)
 
-    # Now we have updated probabilities
     final_label = max(boosted, key=boosted.get)
     final_confidence = round(boosted[final_label], 3)
 
@@ -282,11 +359,14 @@ def smishing_detector(input_type, text, image):
     lower_text = combined_text.lower()
     smishing_keys, scam_keys, _ = get_keywords_by_language(combined_text)
 
-    found_urls = re.findall(r"(https?://[^\s]+|\b(?:[a-zA-Z0-9.-]+\.(?:com|net|org|edu|gov|mil|io|ai|co|info|biz|us|uk|de|fr|es|ru|jp|cn|in|au|ca|br|mx|it|nl|se|no|fi|ch|pl|kr|vn|id|tw|sg|hk))\b)", lower_text)
+    found_urls = re.findall(
+        r"(https?://[^\s]+|\b[a-zA-Z0-9.-]+\.(?:com|net|org|edu|gov|mil|io|ai|co|info|biz|us|uk|de|fr|es|ru|jp|cn|in|au|ca|br|mx|it|nl|se|no|fi|ch|pl|kr|vn|id|tw|sg|hk)\b)",
+        lower_text
+    )
     found_smishing = [kw for kw in smishing_keys if kw in lower_text]
     found_other_scam = [kw for kw in scam_keys if kw in lower_text]
 
-    # 6. Final LLM explanation (in detected_lang)
+    # 6. Final explanation in user's language
     final_explanation = query_llm_for_explanation(
         text=combined_text,
         final_label=final_label,
@@ -333,7 +413,7 @@ def toggle_inputs(choice):
         return gr.update(visible=False), gr.update(visible=True)
 
 with gr.Blocks() as demo:
-    gr.Markdown("## SMiShing & Scam Detector with LLM-Enhanced Logic (Multilingual Explanation)")
+    gr.Markdown("## SMiShing & Scam Detector with Google Safe Browsing + LLM")
     
     with gr.Row():
         input_type = gr.Radio(
@@ -377,4 +457,6 @@ with gr.Blocks() as demo:
 if __name__ == "__main__":
     if not openai.api_key:
         print("WARNING: OPENAI_API_KEY not set. LLM calls will fail or be skipped.")
+    if not SAFE_BROWSING_API_KEY:
+        print("WARNING: GOOGLE_SAFE_BROWSING_API_KEY not set. URL checks will be skipped.")
     demo.launch()

app.py.bestoftues

@@ -8,7 +8,7 @@ from deep_translator import GoogleTranslator
 import openai
 import os
 
-# Set up your OpenAI API key
+# Set your OpenAI API key
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
 # Translator instance
@@ -60,7 +60,7 @@ def boost_probabilities(probabilities: dict, text: str):
     smishing_boost = 0.30 * smishing_count
     other_scam_boost = 0.30 * other_scam_count
 
-    found_urls = re.findall(r"(https?://[^\s]+)", lower_text)
+    found_urls = re.findall(r"(https?://[^\s]+|\b(?:[a-zA-Z0-9.-]+\.(?:com|net|org|edu|gov|mil|io|ai|co|info|biz|us|uk|de|fr|es|ru|jp|cn|in|au|ca|br|mx|it|nl|se|no|fi|ch|pl|kr|vn|id|tw|sg|hk))\b)", lower_text)
     if found_urls:
         smishing_boost += 0.35
 
@@ -120,7 +120,6 @@ def query_llm_for_classification(raw_message: str) -> dict:
         raw_reply = response["choices"][0]["message"]["content"].strip()
 
         import json
-        # Expect something like {"label": "...", "reason": "..."}
         llm_result = json.loads(raw_reply)
         if "label" not in llm_result or "reason" not in llm_result:
             return {"label": "Unknown", "reason": f"Unexpected format: {raw_reply}"}
@@ -170,19 +169,31 @@ def query_llm_for_explanation(
     llm_reason: str,
     found_smishing: list,
     found_other_scam: list,
-    found_urls: list
+    found_urls: list,
+    detected_lang: str
 ) -> str:
     """
-    Second LLM call: provides a holistic explanation of the final classification.
-    We include the local classification info, the LLM’s own classification, and
-    relevant details (keywords, URLs).
+    Second LLM call: provides a holistic explanation of the final classification
+    in the same language as detected_lang (English or Spanish).
     """
-    system_prompt = (
-        "You are a cybersecurity expert providing a final explanation to the user. "
-        "Combine the local classification, the LLM classification, and the final label "
-        "into one concise explanation. Do not reveal internal code or raw JSON; just give "
-        "a short, user-friendly explanation. End with a final statement of the final label."
-    )
+    # Decide the language for final explanation
+    if detected_lang == "es":
+        # Spanish
+        system_prompt = (
+            "Eres un experto en ciberseguridad. Proporciona una explicación final al usuario en español. "
+            "Combina la clasificación local, la clasificación LLM y la etiqueta final en una sola explicación breve. "
+            "No reveles el código interno ni el JSON bruto; simplemente da una breve explicación fácil de entender. "
+            "Termina con la etiqueta final. "
+        )
+    else:
+        # Default to English
+        system_prompt = (
+            "You are a cybersecurity expert providing a final explanation to the user in English. "
+            "Combine the local classification, the LLM classification, and the final label "
+            "into one concise explanation. Do not reveal internal code or raw JSON. "
+            "End with a final statement of the final label."
+        )
+
     user_context = f"""
 User Message:
 {text}
@@ -195,7 +206,7 @@ Suspicious SMiShing Keywords => {found_smishing}
 Suspicious Other Scam Keywords => {found_other_scam}
 URLs => {found_urls}
 """
-    # The LLM can combine these facts into a short paragraph.
+
     try:
         response = openai.ChatCompletion.create(
             model="gpt-3.5-turbo",
@@ -215,7 +226,7 @@ def smishing_detector(input_type, text, image):
     Main detection function combining text (if 'Text') & OCR (if 'Screenshot'),
     plus two LLM calls: 
       1) classification to adjust final probabilities,
-      2) a final explanation summarizing the outcome.
+      2) a final explanation summarizing the outcome in the detected language.
     """
     if input_type == "Text":
         combined_text = text.strip() if text else ""
@@ -271,11 +282,11 @@ def smishing_detector(input_type, text, image):
     lower_text = combined_text.lower()
     smishing_keys, scam_keys, _ = get_keywords_by_language(combined_text)
 
-    found_urls = re.findall(r"(https?://[^\s]+)", lower_text)
+    found_urls = re.findall(r"(https?://[^\s]+|\b(?:[a-zA-Z0-9.-]+\.(?:com|net|org|edu|gov|mil|io|ai|co|info|biz|us|uk|de|fr|es|ru|jp|cn|in|au|ca|br|mx|it|nl|se|no|fi|ch|pl|kr|vn|id|tw|sg|hk))\b)", lower_text)
     found_smishing = [kw for kw in smishing_keys if kw in lower_text]
     found_other_scam = [kw for kw in scam_keys if kw in lower_text]
 
-    # 6. Final LLM explanation
+    # 6. Final LLM explanation (in detected_lang)
     final_explanation = query_llm_for_explanation(
         text=combined_text,
         final_label=final_label,
@@ -286,7 +297,8 @@ def smishing_detector(input_type, text, image):
         llm_reason=llm_reason,
         found_smishing=found_smishing,
         found_other_scam=found_other_scam,
-        found_urls=found_urls
+        found_urls=found_urls,
+        detected_lang=detected_lang
     )
 
     return {
@@ -321,7 +333,7 @@ def toggle_inputs(choice):
         return gr.update(visible=False), gr.update(visible=True)
 
 with gr.Blocks() as demo:
-    gr.Markdown("## SMiShing & Scam Detector with LLM-Enhanced Logic")
+    gr.Markdown("## SMiShing & Scam Detector with LLM-Enhanced Logic (Multilingual Explanation)")
     
     with gr.Row():
         input_type = gr.Radio(
@@ -363,7 +375,6 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    # Warn if openai.api_key not set
     if not openai.api_key:
         print("WARNING: OPENAI_API_KEY not set. LLM calls will fail or be skipped.")
     demo.launch()