import onnxruntime as ort import numpy as np from transformers import AutoTokenizer from fastapi import FastAPI from pydantic import BaseModel import uvicorn from concurrent.futures import ThreadPoolExecutor import tldextract from fastapi.middleware.cors import CORSMiddleware import asyncio import time import re from urllib.parse import urlparse import string from collections import Counter class ONNXPhishingDetector: def __init__(self, model_path="phishing_detector.onnx"): # Initialize with optimized settings and cache self.tokenizer = AutoTokenizer.from_pretrained("microsoft/deberta-v3-base", local_files_only=False) self.session = ort.InferenceSession( model_path, providers=['CPUExecutionProvider'], # Removed CoreMLExecutionProvider due to errors sess_options=self._get_optimized_options() ) self.model_expected_length = 128 self.extract = tldextract.TLDExtract(include_psl_private_domains=True) self.url_cache = {} # Cache for URL analysis results self.thread_pool = ThreadPoolExecutor(max_workers=32) # Increased thread pool size for better parallelism # Comprehensive lists for pattern matching self.suspicious_keywords = { 'login': 'credential-related', 'signin': 'credential-related', 'account': 'credential-related', 'password': 'credential-related', 'verify': 'verification-related', 'secure': 'security-related', 'banking': 'financial-related', 'paypal': 'financial-related', 'wallet': 'financial-related', 'bitcoin': 'cryptocurrency-related', 'crypto': 'cryptocurrency-related', 'authenticate': 'authentication-related', 'authorize': 'authentication-related', 'validation': 'verification-related', 'confirm': 'verification-related' } self.legitimate_tlds = {'.com', '.org', '.net', '.edu', '.gov', '.mil', '.int'} self.suspicious_tlds = {'.xyz', '.top', '.buzz', '.country', '.stream', '.gq', '.tk', '.ml'} # Brand protection patterns self.common_brands = { 'google', 'facebook', 'apple', 'microsoft', 'amazon', 'paypal', 'netflix', 'linkedin', 'twitter', 'instagram' } def _get_optimized_options(self): # Optimize ONNX session options options = ort.SessionOptions() options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL # Note: Changing these thread values doesn't significantly impact performance options.intra_op_num_threads = 4 options.inter_op_num_threads = 4 options.enable_mem_pattern = True options.enable_cpu_mem_arena = True return options def _calculate_entropy(self, text): """Calculate Shannon entropy of domain to detect random-looking strings""" prob = [float(text.count(c)) / len(text) for c in set(text)] entropy = -sum(p * np.log2(p) for p in prob) return entropy def _check_character_distribution(self, domain): """Analyze character distribution patterns""" if not domain: # Handle empty domain case return 0.0, 0.0 char_counts = Counter(domain) total_chars = len(domain) # Check for unusual character distributions digit_ratio = sum(c.isdigit() for c in domain) / total_chars consonant_ratio = sum(c in 'bcdfghjklmnpqrstvwxyz' for c in domain.lower()) / total_chars return digit_ratio, consonant_ratio def _analyze_url_structure(self, url, ext): reasons = [] parsed = urlparse(url) domain = ext.domain # 1. Domain Analysis domain_length = len(domain) entropy = self._calculate_entropy(domain) digit_ratio, consonant_ratio = self._check_character_distribution(domain) # Check domain composition if domain_length > 20: reasons.append(f"Suspicious: Domain length ({domain_length} chars) exceeds normal range") if entropy > 4.5: reasons.append(f"Suspicious: High domain entropy ({entropy:.2f}) suggests randomly generated name") if digit_ratio > 0.4: reasons.append(f"Suspicious: Unusual number of digits ({digit_ratio:.1%} of domain)") if consonant_ratio > 0.7: reasons.append(f"Suspicious: Unusual consonant pattern ({consonant_ratio:.1%} of domain)") # 2. Brand Impersonation Detection for brand in self.common_brands: if brand in domain and brand != domain: if re.search(f"{brand}[^a-zA-Z]", domain) or re.search(f"[^a-zA-Z]{brand}", domain): reasons.append(f"High Risk: Potential brand impersonation of {brand}") # 3. URL Component Analysis if parsed.username or parsed.password: reasons.append("High Risk: URL contains embedded credentials") if parsed.port and parsed.port not in (80, 443): reasons.append(f"Suspicious: Non-standard port number ({parsed.port})") # 4. Path Analysis if parsed.path: path_segments = parsed.path.split('/') if len(path_segments) > 4: reasons.append(f"Suspicious: Deep URL structure ({len(path_segments)} levels)") # Check for suspicious file extensions if any(segment.endswith(('.exe', '.dll', '.bat', '.sh')) for segment in path_segments): reasons.append("High Risk: Contains executable file extension") # 5. Query Parameter Analysis if parsed.query: query_params = parsed.query.split('&') suspicious_params = [p for p in query_params if any(k in p.lower() for k in ['pass', 'pwd', 'token', 'key'])] if suspicious_params: reasons.append("Suspicious: Query contains sensitive parameter names") # 6. Special Pattern Detection if len(re.findall(r'[.-]', domain)) > 4: reasons.append("Suspicious: Excessive use of dots/hyphens in domain") if re.search(r'([a-zA-Z0-9])\1{3,}', domain): reasons.append("Suspicious: Repeated character pattern detected") # 7. TLD Analysis if ext.suffix in self.suspicious_tlds: reasons.append(f"Suspicious: Known high-risk TLD (.{ext.suffix})") elif ext.suffix not in [tld.strip('.') for tld in self.legitimate_tlds]: reasons.append(f"Suspicious: Uncommon TLD (.{ext.suffix})") # 8. Keyword Analysis found_keywords = [] for keyword, category in self.suspicious_keywords.items(): if keyword in f"{domain}{parsed.path}".lower(): found_keywords.append(f"{keyword} ({category})") if found_keywords: reasons.append(f"Suspicious: Contains sensitive keywords: {', '.join(found_keywords)}") return reasons def _batch_preprocess(self, urls): processed = [] for url in urls: url = url.strip().lower() if not url.startswith(('http://', 'https://')): url = f'http://{url}' processed.append(url) return processed def _batch_tokenize(self, urls): return self.tokenizer( urls, max_length=self.model_expected_length, truncation=True, padding="max_length", return_tensors="np" ) def _predict_thread(self, urls): """Process a batch of URLs in a separate thread""" processed_urls = self._batch_preprocess(urls) inputs = self._batch_tokenize(processed_urls) ort_inputs = { "input_ids": inputs["input_ids"].astype(np.int64), "attention_mask": inputs["attention_mask"].astype(np.int64) } try: logits = self.session.run(None, ort_inputs)[0] probabilities = np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True) results = [] for url, prob in zip(urls, probabilities[:, 1]): ext = self.extract(url) reasons = self._analyze_url_structure(url, ext) if prob > 0.99: reasons.append(f"Critical: ML model detected strong phishing patterns (confidence: {prob:.2%})") verdict = "phishing" else: if not reasons: reasons = ["No suspicious patterns detected"] verdict = "legitimate" result = { "url": url, "verdict": verdict, "confidence": float(prob), "reasons": set(reasons) } self.url_cache[url] = result results.append(result) return results except Exception as e: # Fallback to rule-based analysis if model inference fails results = [] for url in urls: ext = self.extract(url) reasons = self._analyze_url_structure(url, ext) # Determine verdict based on rule analysis only if any("High Risk" in reason for reason in reasons): verdict = "phishing" confidence = 0.95 elif len(reasons) > 2: verdict = "phishing" confidence = 0.85 else: verdict = "legitimate" confidence = 0.70 if not reasons: reasons = ["No suspicious patterns detected"] result = { "url": url, "verdict": verdict, "confidence": float(confidence), "reasons": set(reasons + ["Note: Using rule-based analysis due to model inference error"]) } self.url_cache[url] = result results.append(result) return results async def _batch_predict(self, inputs): ort_inputs = { "input_ids": inputs["input_ids"].astype(np.int64), "attention_mask": inputs["attention_mask"].astype(np.int64) } return self.session.run(None, ort_inputs)[0] async def _batch_analyze(self, urls): processed_urls = self._batch_preprocess(urls) inputs = self._batch_tokenize(processed_urls) logits = await self._batch_predict(inputs) probabilities = np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True) return probabilities[:, 1] async def analyze_batch(self, urls): results = [] uncached_urls = [] # Check cache first for url in urls: if url in self.url_cache: results.append(self.url_cache[url]) else: uncached_urls.append(url) if uncached_urls: # Split URLs into smaller batches for multithreaded processing batch_size = 10 # Process 10 URLs per thread url_batches = [uncached_urls[i:i+batch_size] for i in range(0, len(uncached_urls), batch_size)] # Submit each batch to thread pool futures = [] for batch in url_batches: futures.append(self.thread_pool.submit(self._predict_thread, batch)) # Collect results from all threads for future in futures: try: batch_results = future.result() results.extend(batch_results) except Exception as e: # Handle any unexpected errors in thread execution print(f"Error processing batch: {str(e)}") # Create fallback results for this batch for url in batch: results.append({ "url": url, "verdict": "error", "confidence": 0.0, "reasons": {f"Error analyzing URL: {str(e)}"} }) return results app = FastAPI() # Add CORS middleware app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) detector = ONNXPhishingDetector() class UrlList(BaseModel): urls: list[str] @app.post("/scan") async def scan_urls(url_list: UrlList): start_time = time.time() # Process all URLs in a single batch with internal multithreading results = await detector.analyze_batch(url_list.urls) phishing_count = sum(1 for r in results if r["verdict"] == "phishing") avg_confidence = sum(r["confidence"] for r in results) / len(results) if results else 0 if avg_confidence >= 0.99: overall_verdict = "malicious" else: overall_verdict = "safe" return { "time_taken": f"{time.time() - start_time:.2f}s", "total_urls": len(url_list.urls), "legitimate": len(url_list.urls) - phishing_count, "phishing": phishing_count, "overall_verdict": overall_verdict, "average_confidence": avg_confidence, "results": results } # To run this file in terminal: # 1. Make sure you have all dependencies installed: # pip install fastapi uvicorn onnxruntime numpy transformers tldextract # 2. Navigate to the directory containing this file # 3. Run the command: # python -m backend.main # or if you're already in the backend directory: # python main.py # 4. The API will be available at http://localhost:8000 # 5. You can test it with curl: # curl -X POST "http://localhost:8000/scan" -H "Content-Type: application/json" -d '{"urls":["google.com", "suspicious-phishing-site.xyz"]}' # 6. Or use tools like Postman to send POST requests to the /scan endpoint if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=8000)