Look for commercial STOP to lower score.

app.py

@@ -8,6 +8,7 @@ from phishing_datasets import submit_entry
 from url_tools import extract_urls, resolve_short_url, extract_domain_from_url
 from urlscan_client import UrlscanClient
 import requests
+import re
 
 app = FastAPI()
 urlscan = UrlscanClient()
@@ -68,14 +69,39 @@ def get_robots_txt():
 
 @app.post("/predict")
 def predict(model: InputModel) -> OutputModel:
+    sender = model.query.sender
     text = model.query.message.text
 
-    print(f"Predict: {text}")
+    print(f"[{sender}] {text}")
+
+    result = pipe(text)
+    label = result[0]['label']
+    score = result[0]['score']
+
+    print(f"classification {label} score {score}")
+
+    if label == 'LABEL_0':
+        score = 1 - score
+
+    commercial_sender_pattern = r'\b[2-8]\d{4}\b'
+    commercial_stop_pattern = r'\bSTOP(?:\s+SMS)?(?:\s+au)?\s+([2-8]\d{4})\b'
+    commercial_stop = False
+
+    if re.search(commercial_sender_pattern, sender):
+        print("commercial sender")
+        score = score - 0.1
+        if re.search(commercial_stop_pattern, text):
+            print("STOP founded")
+            score = score - 0.2
+            commercial_stop = True
+        else:
+            print("STOP missing")
 
     urls = extract_urls(text)
 
     if urls:
-        print("Searching for past scans")
+        print(f"found URLs: {urls}")
+        print("searching for past scans")
         search_results = [urlscan.search(f"domain:{extract_domain_from_url(url)}") for url in urls]
 
         scan_results = []
@@ -87,32 +113,37 @@ def predict(model: InputModel) -> OutputModel:
                 scan_results.append(scan_result)
 
         if not scan_results:
-            print("Scanning...")
+            print("scanning...")
             scan_results = [urlscan.scan(url) for url in urls]
 
         for result in scan_results:
             overall = result.get('verdicts', {}).get('overall', {})
-            print(f"Checking overall verdict: {overall}")
+            print(f"overall verdict: {overall}")
             if overall.get('hasVerdicts'):
                 score = overall.get('score')
-                print(f"Has verdicts score {score}")
+                print(f"verdict score {score}")
 
                 if 0 < overall.get('score'):
-                    print("Submitting entry and returning JUNK.")
-                    submit_entry(model.query.sender, model.query.message.text)
-                    return OutputModel(action=ActionModel.JUNK, sub_action=SubActionModel.NONE)
-                # elif overall.get('score') < 0:
-                #     print("Returning ALLOW.")
-                #     return OutputModel(action=ActionModel.ALLOW, sub_action=SubActionModel.NONE)
-
-    label = pipe(text)
-    if label[0]['label'] == 'LABEL_1':
-        print("Classify LABEL_1. Submitting entry and returning JUNK.")
-        submit_entry(model.query.sender, model.query.message.text)
-        return OutputModel(action=ActionModel.JUNK, sub_action=SubActionModel.NONE)
+                    score = 1.0
+                    break
+                elif overall.get('score') < 0:
+                    score = score - 0.1
     else:
-        print("Classify LABEL_0. Submitting entry and returning NONE.")
-        return OutputModel(action=ActionModel.NONE, sub_action=SubActionModel.NONE)
+        print(f"no URL found")
+        score = score - 0.1
+
+    print(f"final score {score}")
+    action = ActionModel.NONE
+    if score > 0.7:
+        action=ActionModel.JUNK
+    elif score > 0.5:
+        if commercial_stop:
+            action=ActionModel.PROMOTION
+        else:
+            action=ActionModel.JUNK
+
+    print(f"final action {action}")
+    return OutputModel(action=action, sub_action=SubActionModel.NONE)
 
 class ReportModel(BaseModel):
     sender: str