app.py

from flask import Flask, request, jsonify, render_template
import logging
from urllib.parse import urlparse
from url_analyzer import URLAnalyzer
import sqlite3
from typing import Dict, Optional
import os

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)
url_analyzer = URLAnalyzer()

def get_db():
    """Get database connection."""
    db = sqlite3.connect('data/whitelist.db')
    db.row_factory = sqlite3.Row
    return db

def check_whitelist(domain: str) -> Dict:
    """Check if a domain is in the whitelist.
    
    Args:
        domain: Domain to check
        
    Returns:
        Dict with whitelist status and details
    """
    try:
        # Create a new connection for this request
        db = sqlite3.connect('data/whitelist.db')
        db.row_factory = sqlite3.Row
        cursor = db.cursor()
        
        # Check umbrella table
        cursor.execute("""
            SELECT rank, last_updated 
            FROM umbrella 
            WHERE domain = ?
        """, (domain,))
        result = cursor.fetchone()
        
        # Close connection
        db.close()
        
        if result:
            return {
                "status": "whitelisted",
                "rank": result['rank'],
                "last_updated": result['last_updated']
            }
        return {"status": "not_whitelisted"}
        
    except Exception as e:
        logger.error(f"Failed to check whitelist for {domain}: {str(e)}")
        return {"status": "error", "error": str(e)}

@app.route('/')
def index():
    """Render main page."""
    return render_template('index.html')

@app.route('/api/check', methods=['POST'])
def check_url():
    """Check if URL is phishing."""
    try:
        data = request.get_json()
        url = data.get('url')
        detailed = data.get('detailed', False)
        
        if not url:
            return jsonify({"error": "URL is required"}), 400
            
        # Parse domain for whitelist check
        domain = urlparse(url).netloc
        if domain.startswith('www.'):
            domain = domain[4:]
            
        # Check whitelist first
        whitelist_result = check_whitelist(domain)
        
        # If whitelisted, return early
        if whitelist_result["status"] == "whitelisted":
            return jsonify({
                "url": url,
                "is_phishing": False,
                "confidence": 1.0,
                "details": {
                    "whitelist_status": "whitelisted",
                    "whitelist_rank": whitelist_result["rank"],
                    "whitelist_updated": whitelist_result["last_updated"]
                }
            })
            
        # If not whitelisted, analyze URL
        analysis = url_analyzer.analyze_url(url)
        
        response = {
            "url": url,
            "is_phishing": analysis["is_phishing"],
            "confidence": analysis["confidence"]
        }
        
        if detailed:
            response["details"] = {
                "whitelist_status": whitelist_result["status"],
                "model_score": analysis["confidence"],
                "features": analysis["features"],
                "risk_factors": analysis["analysis"]["risk_factors"]
            }
            
        return jsonify(response)
        
    except Exception as e:
        logger.error(f"Error processing request: {str(e)}")
        return jsonify({"error": str(e)}), 500

@app.route('/api/batch', methods=['POST'])
def check_batch():
    """Check multiple URLs for phishing."""
    try:
        data = request.get_json()
        urls = data.get('urls', [])
        
        if not urls:
            return jsonify({"error": "URLs list is required"}), 400
            
        results = []
        for url in urls:
            # Create a mock request for check_url
            with app.test_request_context(json={'url': url, 'detailed': True}):
                # Call check_url directly
                response = check_url()
                # Get the JSON data from the response
                if isinstance(response, tuple):
                    result = response[0].get_json()
                else:
                    result = response.get_json()
                results.append(result)
            
        return jsonify({
            "results": results,
            "total": len(results),
            "phishing_count": sum(1 for r in results if r["is_phishing"])
        })
        
    except Exception as e:
        logger.error(f"Error processing batch request: {str(e)}")
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    # Ensure data directory exists
    os.makedirs('data', exist_ok=True)
    
    # Start Flask app
    app.run(debug=True)