step5_integrate_predict.py

import re
import joblib
import numpy as np
import pandas as pd
from urllib.parse import urlparse
import sys
import os

# --- Configuration for Model Paths ---
MODELS_DIR = "models" # Assuming all models/scalers/feature lists are in a 'models' directory

# ----------------------------
# Load Models and Artifacts
# ----------------------------

def load_all_artifacts():
    """
    Loads all necessary models, scalers, and feature lists from the 'models' directory.
    Handles missing files gracefully.
    """
    artifacts = {}
    
    # Define expected files and their keys in the artifacts dictionary
    # Updated paths and keys to match the `step4_train_main_with_attacktypes.py` output
    expected_files = {
        "primary_model": os.path.join(MODELS_DIR, "main_phishing_model.pkl"),
        "primary_scaler": os.path.join(MODELS_DIR, "scaler.pkl"),
        "primary_features": os.path.join(MODELS_DIR, "training_features_list.pkl"),
        "attack_model_info": os.path.join(MODELS_DIR, "attack_type_model.pkl"), # This will be the dummy or actual model
    }

    missing_files = []
    for key, path in expected_files.items():
        if not os.path.exists(path):
            missing_files.append(path)
    
    if missing_files:
        print("ERROR: The following required model/artifact files are missing. "
              "Please ensure 'step4_train_main_with_attacktypes.py' was run successfully "
              "and these exist in the 'models/' directory:")
        for f in missing_files:
            print(f" - {f}")
        sys.exit(1)

    try:
        artifacts["primary_model"] = joblib.load(expected_files["primary_model"])
        artifacts["primary_scaler"] = joblib.load(expected_files["primary_scaler"])
        artifacts["primary_features"] = joblib.load(expected_files["primary_features"])
        # For rule-based attack type, `attack_model_info` is expected to be the dummy dictionary
        artifacts["attack_model_info"] = joblib.load(expected_files["attack_model_info"]) 
        
        print("✅ All models and artifacts loaded successfully.")
        return artifacts
    except Exception as e:
        print(f"CRITICAL ERROR: Failed to load one or more artifacts. Details: {e}")
        sys.exit(1)

# Load artifacts once at the start
artifacts = load_all_artifacts()

# Assign loaded artifacts to more readable variables
primary_model = artifacts["primary_model"]
primary_scaler = artifacts["primary_scaler"]
primary_features = artifacts["primary_features"]
attack_model_info = artifacts["attack_model_info"] 

# -------------------------
# RULE-BASED ATTACK TYPES (Duplicate from STEP 4 for self-containment)
# -------------------------
def categorize_attack_type(url: str) -> str:
    """
    Categorizes a URL into an attack type based on keyword matching.
    This function should be identical to the one used in `step4_train_main_with_attacktypes.py`.
    """
    u = str(url).lower()

    if "login" in u or "signin" in u or "sign-in" in u or "log-in" in u:
        return "credential_phishing"

    if "bank" in u or "pay" in u or "wallet" in u or "financial" in u or "creditcard" in u or "upi" in u:
        return "financial_theft"

    if "verify" in u or "update" in u or "secure" in u or "confirm" in u or "reset" in u:
        return "account_verification"

    if "free" in u or "win" in u or "bonus" in u or "gift" in u or "prize" in u:
        return "prize_scam"

    if "download" in u or "exe" in u or "apk" in u or "software" in u or "install" in u:
        return "malware_attack"

    return "unknown"


# ----------------------------
# Feature Extraction (must match STEP 4 exactly)
# ----------------------------

def extract_features(url: str) -> dict:
    """
    Extracts a consistent set of features from a given URL string.
    This function's output must exactly match the features generated by STEP 2/3 and used by STEP 4.
    """
    parsed = urlparse(url)
    domain = parsed.netloc or ""
    path = parsed.path or ""
    query = parsed.query or ""

    return {
        "url_length": len(url),
        "domain_length": len(domain),
        "num_digits": len(re.findall(r"[0-9]", url)),
        "num_dots": url.count("."),
        "num_hyphens": url.count("-"),
        "has_at": 1 if "@" in url else 0,
        "has_ip": 1 if re.match(r"^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$", domain) else 0, # Robust IPv4 check
        "https": 1 if url.startswith("https") else 0,
        "num_subdirs": path.count("/") if path else 0,
        "num_query_params": query.count("=") if query else 0,
        "contains_login": 1 if "login" in url.lower() else 0,
        "contains_secure": 1 if "secure" in url.lower() else 0,
        "contains_account": 1 if "account" in url.lower() else 0,
        "contains_update": 1 if "update" in url.lower() else 0,
        "contains_verify": 1 if "verify" in url.lower() else 0
    }

# ----------------------------
# Additional RULE-BASED PHISHING CHECKS
# ----------------------------
def rule_based_check(url: str) -> tuple[int, list[str]]:
    """
    Applies additional rule-based checks for phishing indicators.
    Returns a score and a list of reasons.
    """
    phishing_signals = 0
    reasons = []

    parsed = urlparse(url)
    domain = parsed.netloc

    # 1. Too many dots (often indicates subdomains used for trickery)
    # A common heuristic is 3 or more dots in the full URL (not just domain), but domain is also valid.
    # Let's consider 3+ dots in the *domain* itself as suspicious.
    if domain.count(".") >= 3: 
        phishing_signals += 1
        reasons.append("Multiple dots/subdomains in domain")

    # 2. Suspicious keywords in the full URL
    keywords = ["verify", "update", "secure", "account", "login", "alert", "paypal", "amazon", "google"] # Added more keywords
    for k in keywords:
        if k in url.lower() and k not in domain.lower(): # Keyword in URL but not in the primary domain
            phishing_signals += 1
            reasons.append(f"Suspicious keyword '{k}' found outside primary domain")
        elif k in url.lower() and domain.lower().count(k) > 1: # Keyword repeated in domain or path
             phishing_signals += 1
             reasons.append(f"Repeated suspicious keyword '{k}'")

    # 3. Hyphen attack (e.g., faceb00k-login.com) or long hyphens (e.g., brand-security-update-id123.com)
    if "-" in domain and domain.count('-') > 1: # More than one hyphen in domain can be suspicious
        phishing_signals += 1
        reasons.append("Multiple hyphens in domain name")
    elif len(domain) > 20 and domain.count('-') > 0 and (domain.lower().count('login') or domain.lower().count('secure')):
        phishing_signals += 1
        reasons.append("Long domain with hyphens and suspicious keywords")


    # 4. Domain impersonation (subdomain trickery: e.g., facebook.com.malicious.com)
    # This checks if a legitimate brand name is followed by a dot and then another domain
    impersonation_brands = ["facebook.com", "amazon.com", "google.com", "apple.com", "microsoft.com", "paypal.com"]
    for brand in impersonation_brands:
        if brand in url.lower() and url.lower().find(brand) < url.lower().find(domain.lower()): # Brand appears before the actual domain
             phishing_signals += 1
             reasons.append(f"Potential domain impersonation (e.g., '{brand}.malicious.com')")
        elif domain.lower().startswith(brand.split('.')[0] + "-"): # e.g., facebook-login.com
            phishing_signals += 1
            reasons.append(f"Domain starts with brand name '{brand.split('.')[0]}' and a hyphen")
        

    # 5. Non-HTTPS for sensitive keywords (less reliable but can add to score)
    if not url.startswith("https") and any(k in url.lower() for k in ["login", "account", "secure", "bank"]):
        phishing_signals += 1
        reasons.append("Non-HTTPS connection for sensitive content")

    # 6. URL Redirection (often indicated by multiple slashes or special characters like @ in path)
    if "://" in parsed.path or "@" in urlparse(url).netloc: # @ in netloc is a known trick
        phishing_signals += 1
        reasons.append("URL uses redirection tricks (e.g., '@' symbol)")

    return phishing_signals, reasons

# ----------------------------
# Main Prediction Function
# ----------------------------
def predict_url(url: str):
    """
    Analyzes a given URL for phishing using an ML model and rule-based checks.
    Also determines the attack type for detected phishing URLs.
    """
    print("\n==============================")
    print("🔍 URL ANALYSIS REPORT")
    print("==============================")
    print(f"🔗 URL: {url}")

    # 1. Feature Extraction
    extracted_feats = extract_features(url)

    # 2. Prepare Features for ML Model
    features_df = pd.DataFrame([extracted_feats])
    
    # Reindex to match the exact features and their order from training
    X_processed = features_df.reindex(columns=primary_features, fill_value=0)
    
    # Ensure all columns are numeric
    for col in X_processed.columns:
        if not pd.api.types.is_numeric_dtype(X_processed[col]):
            X_processed[col] = pd.to_numeric(X_processed[col], errors='coerce').fillna(0)

    # 3. Scale Features
    X_scaled = primary_scaler.transform(X_processed)

    # 4. ML Model Prediction
    main_pred = primary_model.predict(X_scaled)[0]
    phishing_proba = primary_model.predict_proba(X_scaled)[0][1] # Probability of being phishing (class 1)

    # 5. Rule-Based Check
    rule_score, rule_reasons = rule_based_check(url)

    print(f"\n📈 ML Phishing Probability: {phishing_proba:.4f}")
    print(f"🕵️ Rule-based Suspicion Score: {rule_score}")

    # 6. FINAL DECISION Logic
    # Combine ML prediction with rule-based score
    # A URL is considered phishing if ML predicts it or if the rule-based score is high enough.
    # The threshold of 2 for rule_score is a heuristic, adjust based on desired sensitivity.
    is_phishing = False
    if main_pred == 1 and phishing_proba > 0.5: # ML model explicitly predicts phishing with reasonable confidence
        is_phishing = True
    elif phishing_proba > 0.7: # ML model predicts high probability even if binary output is 0 (due to default 0.5 threshold)
        is_phishing = True
    elif rule_score >= 3: # Higher rule score indicates stronger rule-based suspicion
        is_phishing = True
    elif main_pred == 1 and rule_score >= 1: # ML and at least one rule-based signal
        is_phishing = True
    
    if is_phishing:
        print("\n🔴 FINAL RESULT: **Phishing Website Detected!**")

        # Attack Type Classification (Rule-Based from `categorize_attack_type` function)
        attack_type_derived = categorize_attack_type(url)
        print(f"🎯 Attack Type: **{attack_type_derived}**")

        if rule_reasons:
            print("\n📌 Rule-Based Reasons for Suspicion:")
            for r in rule_reasons:
                print(f" - {r}")
        
        # Optional: Add ML model confidence to reasons if it contributed to decision
        if main_pred == 1 and phishing_proba > 0.5:
            print(f" - ML Model classified as phishing with {phishing_proba:.2f} probability.")
        
    else:
        print("\n🟢 FINAL RESULT: **Legitimate Website**")
        if rule_reasons: # Even if legitimate, show minor warnings
            print("\n⚠️ Minor Rule-Based Observations:")
            for r in rule_reasons:
                print(f" - {r}")


    print("==============================")

# ----------------------------
# MAIN LOOP
# ----------------------------
if __name__ == "__main__":
    print("\n--- Initializing Phishing Detector ---")
    # `load_all_artifacts()` is called once at the beginning
    print("Real-Time URL Phishing & Attack Type Detection Tool.")
    print("Type 'exit' or 'quit' to stop.\n")

    while True:
        url = input("Enter URL to analyze: ")
        if url.lower() in ["exit", "quit"]:
            print("Exiting detector. Goodbye!")
            break
        
        if not url:
            print("Please enter a URL.")
            continue

        try:
            predict_url(url)
        except Exception as e:
            print(f"An error occurred during prediction for '{url}': {e}")
            import traceback
            traceback.print_exc() # Print full traceback for debugging
            print("Please ensure the URL format is valid and models are correctly loaded.")