app.py

"""
Gradio app for phishing detection using URLScan.io and Hopsworks.

This app:
1. Loads a trained model from Hopsworks Model Registry
2. Takes a URL as input
3. Scans it using URLScan.io API
4. Extracts features from the scan results
5. Runs inference using the loaded model
6. Returns whether the URL is likely phishing or legitimate
"""

import os
import sys
import re
import logging
import gradio as gr
from typing import Tuple

# Add src directory to path for imports
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from phising_detection.inference import PhishingDetectionPipeline

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Global inference pipeline
pipeline = None
model_info_cache = None


def initialize_app():
    """Initialize the app by loading the inference pipeline."""
    global pipeline

    try:
        # Get URLScan API key
        urlscan_api_key = os.getenv("URLSCAN_API_KEY")

        # Initialize pipeline
        pipeline = PhishingDetectionPipeline(
            model_name="phishing_detector",
            model_version=5,  # Use latest version
            urlscan_api_key=urlscan_api_key
        )

        # Load model from Hopsworks
        pipeline.load_model_from_hopsworks()
        logger.info("Inference pipeline initialized successfully!")

        # Load model metadata (metrics and images)
        logger.info("Loading model metadata and images...")
        load_model_metadata()

        return True
    except Exception as e:
        logger.error(f"Failed to initialize app: {e}")
        return False


def load_model_metadata():
    """Load model metadata and images from Hopsworks."""
    global model_info_cache

    try:
        from phising_detection.utils.hopsworks_utils import connect_to_hopsworks

        project = connect_to_hopsworks()
        mr = project.get_model_registry()

        # Get the same model that pipeline loaded
        model_version = getattr(pipeline, 'model_version', None) if pipeline else None
        model_name = pipeline.model_name if pipeline else "phishing_detector"

        if model_version:
            model_registry = mr.get_model(model_name, version=model_version)
        else:
            model_registry = mr.get_model(model_name)

        # Download model artifacts to get images and metrics
        model_dir = model_registry.download()

        # Parse hyperparameters.txt for metrics
        metrics = {}
        hyperparams_path = os.path.join(model_dir, "hyperparameters.txt")
        if os.path.exists(hyperparams_path):
            with open(hyperparams_path, 'r') as f:
                content = f.read()

                # Extract model name (first line)
                lines = content.split('\n')
                if lines:
                    first_line = lines[0].strip()
                    if 'Phishing Detection Model' in first_line:
                        model_type = first_line.split('-')[-1].strip() if '-' in first_line else 'Unknown'
                        metrics['model_type'] = model_type

                # Extract training info
                if 'CV Folds:' in content:
                    cv_match = re.search(r'CV Folds:\s*(\d+)', content)
                    if cv_match:
                        metrics['cv_folds'] = int(cv_match.group(1))

                    iter_match = re.search(r'RandomizedSearchCV iterations:\s*(\d+)', content)
                    if iter_match:
                        metrics['search_iterations'] = int(iter_match.group(1))

                    cv_score_match = re.search(r'Best CV Score:\s*([\d.]+)', content)
                    if cv_score_match:
                        metrics['best_cv_score'] = float(cv_score_match.group(1))

                # Extract Test Performance metrics
                if 'Test Performance:' in content:
                    test_section = content.split('Test Performance:')[1].split('Features:')[0]

                    acc_match = re.search(r'Accuracy:\s*([\d.]+)', test_section)
                    if acc_match:
                        metrics['test_accuracy'] = float(acc_match.group(1))

                    prec_match = re.search(r'Precision:\s*([\d.]+)', test_section)
                    if prec_match:
                        metrics['test_precision'] = float(prec_match.group(1))

                    rec_match = re.search(r'Recall:\s*([\d.]+)', test_section)
                    if rec_match:
                        metrics['test_recall'] = float(rec_match.group(1))

                    f1_match = re.search(r'F1 Score:\s*([\d.]+)', test_section)
                    if f1_match:
                        metrics['test_f1_score'] = float(f1_match.group(1))

                    roc_match = re.search(r'ROC-AUC:\s*([\d.]+)', test_section)
                    if roc_match:
                        metrics['test_roc_auc'] = float(roc_match.group(1))

                # Extract number of features
                if 'Features:' in content:
                    features_section = content.split('Features:')[1].strip()
                    # Count features in the list
                    feature_list = re.findall(r"'([^']+)'", features_section)
                    metrics['n_features'] = len(feature_list)
                    metrics['feature_names'] = ', '.join(feature_list)

        # Get image paths
        confusion_path = os.path.join(model_dir, "evaluation_image_1.png")
        importance_path = os.path.join(model_dir, "evaluation_image_2.png")

        model_info_cache = {
            'version': model_registry.version,
            'metrics': metrics,
            'confusion_matrix_path': confusion_path if os.path.exists(confusion_path) else None,
            'feature_importance_path': importance_path if os.path.exists(importance_path) else None
        }

        logger.info(f"Loaded model metadata: version {model_registry.version}, {len(metrics)} metrics")

    except Exception as e:
        logger.error(f"Error loading model metadata: {e}")
        model_info_cache = None


def get_model_info() -> str:
    """Get model information as HTML."""
    if not model_info_cache:
        return "<p>Model information not available. Try refreshing the page.</p>"

    metrics = model_info_cache.get('metrics', {})
    version = model_info_cache.get('version', 'unknown')
    model_name = pipeline.model_name if pipeline else "phishing_detector"

    # Check if we have any metrics
    if not metrics:
        return "<p>No metrics found in model artifacts.</p>"

    # Helper function to safely format metric values
    def format_metric(key, default='N/A'):
        value = metrics.get(key, default)
        if value == default:
            return default
        if isinstance(value, float):
            return f"{value:.4f}"
        if isinstance(value, int):
            return str(value)
        return str(value)

    # Get model type from metrics or use default
    model_type = metrics.get('model_type', 'Unknown')

    html = f"""
    <div style="padding: 20px; background-color: #f5f5f5; border-radius: 10px; margin: 10px;">
        <h3>🤖 Model Information</h3>
        <p><strong>Model Type:</strong> {model_type}</p>
        <p><strong>Model Name:</strong> {model_name}</p>
        <p><strong>Version:</strong> {version}</p>
        <hr>
        <h4>📊 Test Performance</h4>
        <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 15px; margin: 15px 0;">
            <div style="padding: 10px; background-color: white; border-radius: 5px;">
                <strong>Accuracy:</strong> <span style="font-size: 1.2em; color: #2196F3;">{format_metric('test_accuracy')}</span>
            </div>
            <div style="padding: 10px; background-color: white; border-radius: 5px;">
                <strong>Precision:</strong> <span style="font-size: 1.2em; color: #4CAF50;">{format_metric('test_precision')}</span>
            </div>
            <div style="padding: 10px; background-color: white; border-radius: 5px;">
                <strong>Recall:</strong> <span style="font-size: 1.2em; color: #FF9800;">{format_metric('test_recall')}</span>
            </div>
            <div style="padding: 10px; background-color: white; border-radius: 5px;">
                <strong>F1 Score:</strong> <span style="font-size: 1.2em; color: #9C27B0;">{format_metric('test_f1_score')}</span>
            </div>
            <div style="padding: 10px; background-color: white; border-radius: 5px; grid-column: span 2;">
                <strong>ROC-AUC:</strong> <span style="font-size: 1.2em; color: #F44336;">{format_metric('test_roc_auc')}</span>
            </div>
        </div>
        <hr>
        <h4>🎯 Training Details</h4>
        <p><strong>CV Folds:</strong> {format_metric('cv_folds')}</p>
        <p><strong>Search Iterations:</strong> {format_metric('search_iterations')}</p>
        <p><strong>Best CV Score:</strong> {format_metric('best_cv_score')}</p>
        <p><strong>Number of Features:</strong> {format_metric('n_features')}</p>
    """

    # Add feature names if available
    if 'feature_names' in metrics:
        html += f"""
        <hr>
        <h4>📝 Features Used</h4>
        <p style="font-size: 0.9em; line-height: 1.6;">{metrics['feature_names']}</p>
        """

    html += "</div>"
    return html


def gradio_interface(url: str) -> Tuple[str, str, str]:
    """
    Gradio interface function.

    Args:
        url: URL to analyze

    Returns:
        Tuple of (result_html, confidence_html, details_html)
    """
    if not url or not url.strip():
        return "Please enter a URL", "", ""

    # Clean URL
    url = url.strip()

    # Add http:// if no protocol specified
    if not url.startswith(('http://', 'https://')):
        url = 'https://' + url

    # Run prediction using pipeline
    result = pipeline.predict_url(url)

    # Check for errors
    if "error" in result:
        return (
            f'<h2 style="color: orange;">ERROR</h2>',
            "",
            f"<p><strong>Error:</strong> {result['error']}</p>"
        )

    # Format result with color
    prediction = result["prediction"]
    confidence = result["confidence"]

    if prediction == "PHISHING":
        result_html = f'<h2 style="color: red;">PHISHING</h2>'
        color = "red"
    elif prediction == "LEGITIMATE":
        result_html = f'<h2 style="color: green;">LEGITIMATE</h2>'
        color = "green"
    else:
        result_html = f'<h2 style="color: orange;">UNKNOWN</h2>'
        color = "orange"

    confidence_html = f'<h3 style="color: {color};">Confidence: {confidence * 100:.2f}%</h3>'

    # Format details
    details_html = f"""
    <h4>Prediction Details:</h4>
    <ul>
        <li><strong>Phishing Probability:</strong> {result['phishing_probability'] * 100:.2f}%</li>
        <li><strong>Legitimate Probability:</strong> {result['legitimate_probability'] * 100:.2f}%</li>
        <li><strong>URLScan UUID:</strong> {result.get('scan_uuid', 'N/A')}</li>
    </ul>

    <h4>Extracted Features:</h4>
    <ul>
        <li><strong>Domain Age (days):</strong> {result['features'].get('domain_age_days', 'N/A')}</li>
        <li><strong>Secure Percentage:</strong> {result['features'].get('secure_percentage', 'N/A')}%</li>
        <li><strong>Has Umbrella Rank:</strong> {'Yes' if result['features'].get('has_umbrella_rank') else 'No'}</li>
        <li><strong>Umbrella Rank:</strong> {result['features'].get('umbrella_rank', 'N/A')}</li>
        <li><strong>Has TLS:</strong> {'Yes' if result['features'].get('has_tls') else 'No'}</li>
        <li><strong>TLS Valid Days:</strong> {result['features'].get('tls_valid_days', 'N/A')}</li>
        <li><strong>URL Length:</strong> {result['features'].get('url_length', 'N/A')}</li>
        <li><strong>Subdomain Count:</strong> {result['features'].get('subdomain_count', 'N/A')}</li>
    </ul>
    """

    return result_html, confidence_html, details_html


def create_gradio_app():
    """Create and configure the Gradio interface."""

    # Custom CSS for better styling
    css = """
    .output-box {
        padding: 20px;
        border-radius: 10px;
        margin: 10px 0;
    }
    """

    with gr.Blocks(css=css, title="Phishing URL Detection") as demo:
        gr.Markdown(
            """
            # Phishing URL Detection

            Enter a URL to check if it's a phishing website or legitimate.
            This app uses URLScan.io to analyze the website and a machine learning model
            trained on URLScan features to predict if it's phishing.

            **Note:** Scanning a URL can take up to 90 seconds as we wait for URLScan.io to complete the analysis.
            """
        )

        # Create tabs for URL Checker and Model Info
        with gr.Tabs():
            # Tab 1: URL Checker
            with gr.Tab("🔍 URL Checker"):
                with gr.Row():
                    with gr.Column(scale=3):
                        url_input = gr.Textbox(
                            label="URL to Check",
                            placeholder="Enter URL (e.g., example.com or https://example.com)",
                            lines=1
                        )
                    with gr.Column(scale=1):
                        submit_btn = gr.Button("Check URL", variant="primary", size="lg")

                with gr.Row():
                    result_output = gr.HTML(label="Prediction")

                with gr.Row():
                    confidence_output = gr.HTML(label="Confidence")

                with gr.Row():
                    details_output = gr.HTML(label="Details")

                # Example URLs
                gr.Markdown("### Example URLs to Try:")
                gr.Examples(
                    examples=[
                        ["https://google.com"],
                        ["http://001983878188731stea8a1a0.myclickfunnels.com/31acc20172"],
                        ["https://www.alphaspel.se/"],
                    ],
                    inputs=url_input,
                )

                # Connect button to function
                submit_btn.click(
                    fn=gradio_interface,
                    inputs=url_input,
                    outputs=[result_output, confidence_output, details_output]
                )

            # Tab 2: Model Information
            with gr.Tab("📊 Model Info"):
                gr.HTML(value=get_model_info(), label="Model Statistics")

                gr.Markdown("### Model Evaluation Visualizations")

                with gr.Row():
                    with gr.Column():
                        gr.Markdown("#### Confusion Matrix")
                        if model_info_cache and model_info_cache.get('confusion_matrix_path'):
                            gr.Image(value=model_info_cache['confusion_matrix_path'], label="Confusion Matrix")
                        else:
                            gr.Markdown("*Confusion matrix image not available*")

                    with gr.Column():
                        gr.Markdown("#### Feature Importance")
                        if model_info_cache and model_info_cache.get('feature_importance_path'):
                            gr.Image(value=model_info_cache['feature_importance_path'], label="Feature Importance")
                        else:
                            gr.Markdown("*Feature importance image not available*")

        gr.Markdown(
            """
            ---
            **Disclaimer:** This tool is for educational and research purposes only.
            The predictions are not 100% accurate and should not be the sole basis for security decisions.
            """
        )

    return demo


def main():
    """Main function to run the Gradio app."""
    logger.info("Starting Phishing Detection App...")

    # Initialize app (load model and URLScan client)
    logger.info("Initializing app...")
    if not initialize_app():
        logger.error("Failed to initialize app. Exiting.")
        return

    # Create and launch Gradio app
    logger.info("Creating Gradio interface...")
    demo = create_gradio_app()

    logger.info("Launching Gradio app...")
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )


if __name__ == "__main__":
    main()