""" DecisionTree_inference.py ========================= Loads the trained Decision Tree model (phishing_model_dt.pkl) and features.pkl to predict whether a given URL is legitimate or phishing. """ import joblib import pandas as pd import re import time import math import psutil import os from collections import Counter from urllib.parse import urlparse # ============================================================================= # CONSTANTS (Synchronized with Training) # ============================================================================= PHISHING_KEYWORDS = [ "login", "verify", "update", "secure", "account", "bank", "confirm", "signin", "password", "reset", "admin" ] SUSPICIOUS_TLDS = ["xyz", "top", "club", "tk", "ml", "ga", "cf", "buzz", "cn", "ru"] BRANDS = ["paypal", "google", "facebook", "bank", "apple", "microsoft", "amazon"] SHORTENERS = ["bit.ly", "tinyurl.com", "t.co", "goo.gl", "is.gd", "ow.ly", "buff.ly"] # ----------------------------------------------------------------------------- # WHITELIST — Only include domains that: # (a) Never host user-generated / third-party content, AND # (b) Are not abused as phishing infrastructure # ----------------------------------------------------------------------------- TRUSTED_DOMAINS = [ # "google.com", # DISABLED: docs/drive/sites subdomains host real phishing # "facebook.com", # DISABLED: used as phishing redirect infrastructure "github.com", "youtube.com", "microsoft.com", "apple.com", ] # Brand domain mapping — used to distinguish impersonation from genuine brand URLs BRAND_REAL_DOMAINS = { "paypal": ["paypal.com"], "google": ["google.com", "google.co"], "facebook": ["facebook.com", "fb.com"], "bank": [], # Generic word — flag all uses in domain "apple": ["apple.com"], "microsoft": ["microsoft.com", "live.com", "outlook.com"], "amazon": ["amazon.com", "aws.amazon.com"], } # ============================================================================= # HELPER FUNCTIONS # ============================================================================= def is_ip(domain: str) -> bool: return bool(re.match(r"^\d{1,3}(\.\d{1,3}){3}$", domain)) def entropy(s: str) -> float: if not s: return 0.0 prob = [n / len(s) for _, n in Counter(s).items()] return -sum(p * math.log2(p) for p in prob) def has_brand_impersonation(url_lower: str, domain: str) -> int: for brand, real_domains in BRAND_REAL_DOMAINS.items(): if brand in url_lower: # If the domain IS one of the real brand domains, it is not impersonation if any(domain == rd or domain.endswith("." + rd) for rd in real_domains): continue return 1 return 0 def count_domain_phishing_keywords(domain: str) -> int: return sum(k in domain for k in PHISHING_KEYWORDS) def count_path_phishing_keywords(path: str) -> int: return sum(k in path for k in PHISHING_KEYWORDS) def is_trusted_domain(domain: str, path: str) -> bool: for trusted in TRUSTED_DOMAINS: if domain == trusted or domain.endswith("." + trusted): # Extra guard: if the path itself looks phishy, don't whitelist if count_path_phishing_keywords(path) == 0: return True return False # ============================================================================= # FEATURE EXTRACTION # ============================================================================= def extract_features(url: str) -> dict: url = url.strip() # Normalise scheme so urlparse works correctly if not url.startswith(("http://", "https://")): url = "https://" + url parsed = urlparse(url) domain = parsed.netloc.lower().split(":")[0] path = parsed.path.lower() query = parsed.query.lower() length = len(url) num_digits = sum(c.isdigit() for c in url) num_special = sum(not c.isalnum() for c in url) tld = domain.split(".")[-1] if "." in domain else "" return { "url_length": length, "domain_length": len(domain), "path_length": len(path), "num_subdomains": domain.count("."), "num_hyphens_url": url.count("-"), "num_hyphens_domain": domain.count("-"), "num_digits": num_digits, "digit_ratio": num_digits / length if length > 0 else 0, "num_special_chars": num_special, "special_char_ratio": num_special / length if length > 0 else 0, "num_query_params": query.count("&") + (1 if query else 0), "num_phishing_keywords_domain": count_domain_phishing_keywords(domain), "num_phishing_keywords_path": count_path_phishing_keywords(path), "is_http_only": int(url.startswith("http://") and not url.startswith("https://")), "https_in_domain": int("https" in domain), "has_www": int("www." in domain), "tld_length": len(tld), "is_ip_address": int(is_ip(domain)), "is_suspicious_tld": int(tld in SUSPICIOUS_TLDS), "is_shortener": int(any(s in domain for s in SHORTENERS)), "has_brand_name": has_brand_impersonation(url.lower(), domain), "has_at_symbol": int("@" in url), "has_double_slash_path": int("//" in path), "domain_entropy": entropy(domain), } # ============================================================================= # LOAD MODEL # ============================================================================= load_start = time.perf_counter() # Ensure these files exist in the same directory and were produced by DecisionTree_train.py model = joblib.load("phishing_model_dt.pkl") FEATURES = joblib.load("features_dt.pkl") load_end = time.perf_counter() print(f"Model loading time: {load_end - load_start:.6f} seconds") # ============================================================================= # PREDICTION FUNCTION # ============================================================================= def predict_url(url: str, threshold: float = 0.5): timings = {} url = url.strip() # Normalise for domain extraction url_to_parse = url if url.startswith(("http://", "https://")) else "https://" + url parsed = urlparse(url_to_parse) domain = parsed.netloc.lower().split(":")[0] path = parsed.path.lower() # ── Whitelist check (with path guard) ───────────────────────────────── if is_trusted_domain(domain, path): return "LEGITIMATE (Whitelisted)", 0.00, {} # ── Feature extraction ───────────────────────────────────────────────── feat_start = time.perf_counter() features = extract_features(url) # Debug: find missing keys missing = [f for f in FEATURES if f not in features] if missing: print(f"Missing features: {missing}") # Fill missing with 0 as fallback for key in missing: features[key] = 0 X = pd.DataFrame([features])[FEATURES] feat_end = time.perf_counter() timings["feature_extraction"] = feat_end - feat_start # ── Model inference ──────────────────────────────────────────────────── pred_start = time.perf_counter() prob = model.predict_proba(X)[0][1] pred_end = time.perf_counter() timings["prediction"] = pred_end - pred_start label = "PHISHING" if prob >= threshold else "LEGITIMATE" return label, prob, timings # ============================================================================= # TESTING & PROFILING # ============================================================================= test_urls = [ # Should be LEGITIMATE "https://www.tradingview.com", "www.tradingview.com", "https://www.reynoldstransfer.com/", # Should be PHISHING "https://docs.google.com/forms/d/e/1FAIpQLSel0WzlUSRy5Yq7D5GuYeKTGE6eCIY12-8pYpvhFYihnkOQ8g/viewform", "https://sites.google.com/view/paypal-customer-services/", "paypal-secure-update.com/login", "bit.ly/3xYz8q", "https://httpsservices.runescape.com-ov.ru/", ] for url in test_urls: total_start = time.perf_counter() process = psutil.Process(os.getpid()) mem_before = process.memory_info().rss / (1024 * 1024) label, prob, timings = predict_url(url) mem_after = process.memory_info().rss / (1024 * 1024) total_end = time.perf_counter() print("\n" + "=" * 55) print(f"URL : {url}") print(f"Prediction: {label}") print(f"Phishing probability: {prob:.3f}") if timings: print(f"Feature extraction : {timings.get('feature_extraction', 0):.6f}s") print(f"Model inference : {timings.get('prediction', 0):.6f}s") else: print("Skipped ML inference (Whitelisted)") print(f"Total inference time : {total_end - total_start:.6f}s") print(f"Memory fluctuation : {mem_after - mem_before:.2f} MB")