push speed

app/safety_check.py

@@ -8,117 +8,176 @@ from PIL import Image
 import torch
 import re
 from urllib.parse import urlparse, unquote
+from functools import lru_cache
+import threading
 
-# ==== Load models 1 lần duy nhất ====
-# Text độc hại
+# Load models
 detox_model = Detoxify('original')
 
-# NSFW
 nsfw_model_id = "Falconsai/nsfw_image_detection"
 nsfw_processor = AutoProcessor.from_pretrained(nsfw_model_id)
 nsfw_model = AutoModelForImageClassification.from_pretrained(nsfw_model_id)
 
-# Bạo lực
 violence_model_id = "jaranohaal/vit-base-violence-detection"
-violence_processor = ViTFeatureExtractor.from_pretrained(violence_model_id)
 violence_model = ViTForImageClassification.from_pretrained(violence_model_id)
+violence_processor = ViTFeatureExtractor.from_pretrained(violence_model_id)
 
-# Caption ảnh
 blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
 blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 
-# AI phân loại URL
-url_classifier = pipeline("zero-shot-classification")
+# Load mô hình zero-shot nhẹ
+classifier = pipeline("zero-shot-classification", model="MoritzLaurer/distilbert-zero-shot-v1")
 
-# ==== HÀM CHÍNH ====
 
+# ========================== TEXT (Prompt) ==========================
+@lru_cache(maxsize=128)
 def is_prompt_safe(prompt: str):
     results = detox_model.predict(prompt)
     threshold = 0.5
     flagged = {label: score for label, score in results.items() if score > threshold}
-    return False if flagged else True, list(flagged.keys())
+    if flagged:
+        return False, list(flagged.keys())
+    return True, []
+
 
+# ========================== IMAGE (NSFW + Violence) ==========================
 def generate_caption(image: Image.Image):
     inputs = blip_processor(images=image, return_tensors="pt")
     with torch.no_grad():
-        output = blip_model.generate(**inputs)
-    return blip_processor.decode(output[0], skip_special_tokens=True)
-
-def check_nsfw_image(image: Image.Image):
-    inputs = nsfw_processor(images=image, return_tensors="pt")
-    with torch.no_grad():
-        logits = nsfw_model(**inputs).logits
-        probs = torch.nn.functional.softmax(logits, dim=1)[0]
-    labels = list(nsfw_model.config.id2label.values())
-    pred_idx = probs.argmax().item()
-    return labels[pred_idx], probs[pred_idx].item() * 100
-
-def check_violence_image(image: Image.Image):
-    inputs = violence_processor(images=image, return_tensors="pt")
-    with torch.no_grad():
-        logits = violence_model(**inputs).logits
-        probs = torch.nn.functional.softmax(logits, dim=1)[0]
-    labels = ["Non-Violent", "Violent"]
-    pred_idx = probs.argmax().item()
-    return labels[pred_idx], probs[pred_idx].item() * 100
-
-def analyze_image(image: Image.Image) -> str:
-    nsfw_label, nsfw_score = check_nsfw_image(image)
-    violence_label, violence_score = check_violence_image(image)
-    caption = generate_caption(image)
-
-    result = f"🖼️ Mô tả ảnh: {caption}\n\n"
-
-    # NSFW
-    if nsfw_label.lower() in ["porn", "hentai", "sex", "nsfw"] and nsfw_score > 60:
-        result += f"🚨 NSFW: {nsfw_label} ({nsfw_score:.2f}%)\n"
-    else:
-        result += f"✅ An toàn NSFW: {nsfw_label} ({nsfw_score:.2f}%)\n"
+        out = blip_model.generate(**inputs)
+    caption = blip_processor.decode(out[0], skip_special_tokens=True)
+    return caption
+
+def analyze_image(image: Image.Image):
+    result = {
+        "nsfw": "",
+        "violence": ""
+    }
+
+    def nsfw_task():
+        nsfw_inputs = nsfw_processor(images=image, return_tensors="pt")
+        with torch.no_grad():
+            nsfw_outputs = nsfw_model(**nsfw_inputs)
+            nsfw_probs = torch.nn.functional.softmax(nsfw_outputs.logits, dim=1)[0]
+        nsfw_labels = list(nsfw_model.config.id2label.values())
+        nsfw_pred = nsfw_probs.argmax().item()
+        nsfw_label = nsfw_labels[nsfw_pred]
+        nsfw_score = nsfw_probs[nsfw_pred].item() * 100
+        caption = generate_caption(image)
+
+        if nsfw_label.lower() in ["porn", "hentai", "sex", "nsfw"]:
+            result["nsfw"] = f"""\ud83d\udea8 Ảnh KHÔNG an toàn (NSFW):\n- Loại: {nsfw_label}\n- Độ chính xác: {nsfw_score:.2f}%\n- Mô tả: {caption}"""
+        else:
+            result["nsfw"] = f"""✅ Ảnh an toàn (NSFW):\n- Loại: {nsfw_label}\n- Độ chính xác: {nsfw_score:.2f}%\n- Mô tả: {caption}"""
+
+    def violence_task():
+        violence_inputs = violence_processor(images=image, return_tensors="pt")
+        with torch.no_grad():
+            violence_outputs = violence_model(**violence_inputs)
+            violence_probs = torch.nn.functional.softmax(violence_outputs.logits, dim=1)[0]
+        violence_labels = ["Non-Violent", "Violent"]
+        violence_pred = violence_probs.argmax().item()
+        violence_label = violence_labels[violence_pred]
+        violence_score = violence_probs[violence_pred].item() * 100
+        caption = generate_caption(image)
+
+        is_violent = False
+        if violence_label.lower() == "non-violent" and violence_score > 50:
+            is_violent = True
+        elif violence_label.lower() == "violent" and violence_score > 80:
+            is_violent = True
+
+        if is_violent:
+            result["violence"] = f"""\ud83d\udea8 Ảnh KHÔNG an toàn (Bạo lực):\n- Loại: {violence_label}\n- Độ chính xác: {violence_score:.2f}%\n- Mô tả: {caption}"""
+        else:
+            result["violence"] = f"""✅ Ảnh an toàn (Bạo lực):\n- Loại: {violence_label}\n- Độ chính xác: {violence_score:.2f}%\n- Mô tả: {caption}"""
 
-    # Bạo lực
-    if violence_label == "Violent" and violence_score > 80:
-        result += f"🚨 Bạo lực: {violence_label} ({violence_score:.2f}%)"
-    else:
-        result += f"✅ An toàn bạo lực: {violence_label} ({violence_score:.2f}%)"
+    # Chạy song song
+    t1 = threading.Thread(target=nsfw_task)
+    t2 = threading.Thread(target=violence_task)
+    t1.start(); t2.start()
+    t1.join(); t2.join()
+    return result["nsfw"], result["violence"]
 
-    return result
 
-def check_url(url: str) -> str:
+# ========================== URL ==========================
+@lru_cache(maxsize=128)
+def check_url(url: str):
     try:
-        decoded = unquote(url)
-        parsed = urlparse(decoded)
+        decoded_url = unquote(url)
+        parsed = urlparse(decoded_url)
         warnings = []
 
-        # Kiểm tra dấu hiệu đáng ngờ
-        if re.match(r'^https?://\d{1,3}(\.\d{1,3}){3}', decoded):
-            warnings.append("🚨 Truy cập IP trực tiếp")
-        if re.search(r'\.(exe|msi|bat|js|jar|apk|dmg)(\?|$)', parsed.path.lower()):
-            warnings.append("🚨 URL chứa file thực thi")
+        if re.match(r'^https?://\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}', decoded_url):
+            warnings.append("\ud83d\udea8 Nguy hiểm: Truy cập trực tiếp bằng IP (thường dùng cho tấn công)")
+
+        if re.search(r'\\.(exe|msi|bat|js|jar|apk|dmg)(\\?|$)', parsed.path.lower()):
+            warnings.append("\ud83d\udea8 Nguy hiểm: URL chứa file thực thi có thể độc hại")
+
         if 'redirect' in parsed.path.lower() or 'url=' in parsed.query.lower():
-            warnings.append("⚠️ Chứa chức năng chuyển hướng")
+            warnings.append("\u26a0\ufe0f Cảnh báo: URL chứa chức năng redirect (có thể lừa đảo)")
+
         if re.search(r'%[0-9a-f]{2}|[\x00-\x1f\x7f]', url):
-            warnings.append("🚨 Ký tự mã hóa bất thường")
+            warnings.append("\ud83d\udea8 Nguy hiểm: URL chứa ký tự mã hóa đáng ngờ (có thể tấn công)")
+
         if '@' in parsed.netloc:
-            warnings.append("🚨 URL giả mạo domain (dùng @)")
-        if any(k in parsed.netloc.lower() for k in ['login', 'secure', 'account']):
-            warnings.append("⚠️ Domain trông giống dịch vụ đăng nhập")
+            warnings.append("\ud83d\udea8 Lừa đảo: URL chứa kỹ thuật giả mạo domain (user@fake-domain)")
+
+        deceptive_domains = ['login', 'secure', 'account', 'verify', 'update']
+        if any(keyword in parsed.netloc.lower() for keyword in deceptive_domains):
+            warnings.append("\u26a0\ufe0f Cảnh báo: Domain có dấu hiệu giả mạo dịch vụ đăng nhập")
+
         if parsed.scheme == 'http':
-            warnings.append("⚠️ Kết nối không mã hóa (HTTP)")
+            warnings.append("\u26a0\ufe0f Cảnh báo: Kết nối không mã hóa (HTTP)")
 
-        # Phân tích bằng AI
-        ai_result = url_classifier(url, candidate_labels=["malicious", "safe"])
+        ai_result = classifier(url, candidate_labels=["malicious", "safe"])
         ai_label = ai_result["labels"][0]
         ai_score = ai_result["scores"][0] * 100
 
-        # Kết luận
-        result = f"🔗 URL: {url}\n"
+        report = {
+            "url": url,
+            "decoded_url": decoded_url,
+            "domain": parsed.netloc,
+            "path": parsed.path,
+            "warnings": warnings,
+            "ai_analysis": {
+                "label": ai_label,
+                "confidence": ai_score
+            }
+        }
+
         if warnings or ai_label == "malicious":
-            result += f"🚨 KHÔNG AN TOÀN\n\n📢 Cảnh báo:\n" + "\n".join(f"- {w}" for w in warnings)
+            return format_report(report, is_safe=False)
         else:
-            result += "✅ An toàn\n"
-
-        result += f"\n🤖 AI đánh giá: {ai_label} ({ai_score:.2f}%)"
-        return result
+            return format_report(report, is_safe=True)
 
     except Exception as e:
-        return f"⚠️ Lỗi phân tích URL: {str(e)}"
+        return f"\u26a0\ufe0f Lỗi khi phân tích URL: {str(e)}"
+
+def format_report(report: dict, is_safe: bool):
+    warning_text = "\n".join(f"- {w}" for w in report["warnings"]) if report["warnings"] else "- Không phát hiện cảnh báo"
+
+    if not is_safe:
+        return f"""\ud83d\udea8 URL KHÔNG AN TOÀN
+\ud83d\udd0d Phân tích chi tiết:
+• URL gốc: {report['url']}
+• Domain: {report['domain']}
+• Đường dẫn: {report['path']}
+
+\ud83d\udce3 CẢNH BÁO:
+{warning_text}
+
+🤖 Phân tích AI:
+- Kết quả: {report['ai_analysis']['label']}
+- Độ tin cậy: {report['ai_analysis']['confidence']:.2f}%
+
+🛡️ Khuyến nghị: KHÔNG TRUY CẬP!"""
+    else:
+        return f"""✅ URL AN TOÀN
+\ud83d\udd0d Phân tích chi tiết:
+• URL gốc: {report['url']}
+• Domain: {report['domain']}
+
+🤖 Phân tích AI:
+- Kết quả: {report['ai_analysis']['label']}
+- Độ tin cậy: {report['ai_analysis']['confidence']:.2f}%"""