update code with advance model

pages/malicious_url.py

@@ -1,98 +1,448 @@
-# pages/malicious_url.py
-
-import streamlit as st
-import numpy as np
-import pandas as pd
-import tensorflow as tf
-from urllib.parse import urlparse
-import re
-import joblib
-
-@st.cache_resource
-def load_model_and_scaler():
-    model = tf.keras.models.load_model("models/malicious_url_model.h5")
-    scaler = joblib.load("models/scaler.pkl")
-    return model, scaler
-
-model, scaler = load_model_and_scaler()
-
-def extract_features(url):
-    try:
-        parsed_url = urlparse(str(url))
-        features = {
-            'url_length': len(str(url)),
-            'hostname_length': len(parsed_url.hostname) if parsed_url.hostname else 0,
-            'path_length': len(parsed_url.path) if parsed_url.path else 0,
-            'query_length': len(parsed_url.query) if parsed_url.query else 0,
-            'fragment_length': len(parsed_url.fragment) if parsed_url.fragment else 0,
-            'num_dots': str(url).count('.'),
-            'num_hyphens': str(url).count('-'),
-            'num_at': str(url).count('@'),
-            'num_question': str(url).count('?'),
-            'num_ampersand': str(url).count('&'),
-            'num_equals': str(url).count('='),
-            'num_exclamation': str(url).count('!'),
-            'num_slash': str(url).count('/'),
-            'num_plus': str(url).count('+'),
-            'num_asterisk': str(url).count('*'),
-            'num_underscore': str(url).count('_'),
-            'num_hash': str(url).count('#'),
-            'num_dollar': str(url).count('$'),
-            'num_percent': str(url).count('%'),
-            'is_https': 1 if parsed_url.scheme == 'https' else 0,
-            'has_http_in_hostname': 1 if parsed_url.hostname and 'http' in parsed_url.hostname else 0,
-            'hostname_is_ip': 1 if parsed_url.hostname and re.match(r'^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$', parsed_url.hostname) else 0,
-            'path_depth': str(url).count('/') - 2 if url and urlparse(str(url)).scheme in ['http', 'https'] and str(url).count('/') > 2 else 0
-        }
-        return pd.Series(features)
-    except Exception:
-        return pd.Series({
-            'url_length': 0, 'hostname_length': 0, 'path_length': 0,
-            'query_length': 0, 'fragment_length': 0, 'num_dots': 0,
-            'num_hyphens': 0, 'num_at': 0, 'num_question': 0,
-            'num_ampersand': 0, 'num_equals': 0, 'num_exclamation': 0,
-            'num_slash': 0, 'num_plus': 0, 'num_asterisk': 0,
-            'num_underscore': 0, 'num_hash': 0, 'num_dollar': 0,
-            'num_percent': 0, 'is_https': 0, 'has_http_in_hostname': 0,
-            'hostname_is_ip': 0, 'path_depth': 0
-        })
-
-X_columns = [
-    'url_length', 'hostname_length', 'path_length', 'query_length',
-    'fragment_length', 'num_dots', 'num_hyphens', 'num_at',
-    'num_question', 'num_ampersand', 'num_equals', 'num_exclamation',
-    'num_slash', 'num_plus', 'num_asterisk', 'num_underscore',
-    'num_hash', 'num_dollar', 'num_percent', 'is_https',
-    'has_http_in_hostname', 'hostname_is_ip', 'path_depth'
-]
-
-def app():
-    st.title("🔗 Malicious URL Detector")
-    st.markdown("Enter a URL below to check if it's likely malicious.")
-
-    url_input = st.text_input(
-        "🔗 Enter a URL:",
-        placeholder="e.g., https://example.com", 
-        help="Type any URL you want to analyze"
-    )
-
-    if st.button("🔍 Analyze URL"):
-        if not url_input.strip():
-            st.warning("Please enter a valid URL.")
-        else:
-            with st.spinner("Analyzing..."):
-                features = extract_features(url_input)
-                df_new = pd.DataFrame([features])
-                X_new = df_new[X_columns]
-                X_new.fillna(-1, inplace=True)
-                X_scaled = scaler.transform(X_new)
-                prediction = model.predict(X_scaled)
-                prob = float(prediction[0][0])
-
-                if prob > 0.5:
-                    st.error(f"⚠️ This URL is likely **malicious**. Confidence: `{prob:.4f}`")
-                else:
-                    st.success(f"✅ This URL appears to be **safe**. Confidence: `{1 - prob:.4f}`")
-
-    st.markdown("---")
-    st.markdown("💡 *Model trained on URL-based features like length, special characters, domain patterns, etc.*")
+import streamlit as st
+import re
+import math
+import time
+import os
+import joblib
+import numpy as np
+import pandas as pd
+import torch
+import tensorflow as tf
+from collections import Counter
+from tensorflow.keras.models import load_model
+from tensorflow.keras.preprocessing.sequence import pad_sequences
+import tldextract
+from rapidfuzz import fuzz, process
+
+# Set page config
+st.set_page_config(
+    page_title="URL Threat Detector",
+    page_icon="🛡️",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+# Disable GPU usage for TensorFlow and PyTorch
+os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+tf.config.set_visible_devices([], 'GPU')
+
+# Global configuration
+MAX_LEN = 200
+FIXED_FEATURE_COLS = [
+    'url_length', 'domain_length', 'subdomain_count', 'path_depth',
+    'param_count', 'has_ip', 'has_executable', 'has_double_extension',
+    'hex_encoded', 'digit_ratio', 'special_char_ratio', 'entropy',
+    'is_safe_domain', 'is_uncommon_tld'
+]
+
+# Enhanced domain and TLD lists
+SAFE_DOMAINS = {
+    'google.com', 'google.co.in', 'google.co.uk', 'google.fr', 'google.de',
+    'amazon.com', 'amazon.in', 'amazon.co.uk', 'amazon.de', 'amazon.fr',
+    'wikipedia.org', 'github.com', 'python.org', 'irs.gov', 'adobe.com',
+    'steampowered.com', 'imdb.com', 'weather.com', 'archive.org', 'cdc.gov',
+    'microsoft.com', 'apple.com', 'youtube.com', 'facebook.com', 'twitter.com',
+    'linkedin.com', 'instagram.com', 'netflix.com', 'reddit.com', 'stackoverflow.com',
+    'google.com', 'amazon.in', 'linkedin.com'
+}
+
+COMMON_TLDS = {
+    'com', 'org', 'net', 'gov', 'edu', 'mil', 'co', 'io', 'ai', 'in',
+    'uk', 'us', 'ca', 'au', 'de', 'fr', 'es', 'it', 'nl', 'jp', 'cn',
+    'br', 'mx', 'ru', 'ch', 'se', 'no', 'dk', 'fi', 'be', 'at', 'nz'
+}
+
+# Initialize tldextract
+tld_extractor = tldextract.TLDExtract()
+
+@st.cache_resource
+def load_char_mapping():
+    char_to_idx_path = 'char_to_idx.pkl'
+    if not os.path.exists(char_to_idx_path):
+        st.error(f"Character mapping file not found: {char_to_idx_path}")
+        return None
+    return joblib.load(char_to_idx_path)
+
+@st.cache_resource
+def load_all_models():
+    """Load models with CPU optimization"""
+    models = {}
+    model_dir = "models"
+
+    if not os.path.exists(model_dir):
+        os.makedirs(model_dir)
+        st.warning(f"Created model directory: {model_dir}")
+
+    # Hybrid models
+    hybrid_models = {
+        'hybrid': 'hybrid_model.h5',
+        'hybrid_fold1': 'best_hybrid_fold1.h5',
+        'hybrid_fold2': 'best_hybrid_fold2.h5'
+    }
+
+    for name, file in hybrid_models.items():
+        path = os.path.join(model_dir, file)
+        if os.path.exists(path):
+            try:
+                models[name] = load_model(path)
+                st.success(f"Loaded {name}")
+            except Exception as e:
+                st.error(f"Error loading {name}: {str(e)}")
+        else:
+            st.warning(f"Model file not found: {path}")
+
+    # Traditional models
+    traditional_models = {
+        'random_forest': 'random_forest_model.pkl',
+        'xgboost': 'xgboost_model.pkl',
+    }
+
+    for name, file in traditional_models.items():
+        path = os.path.join(model_dir, file)
+        if os.path.exists(path):
+            try:
+                models[name] = joblib.load(path)
+                st.success(f"Loaded {name}")
+            except Exception as e:
+                st.error(f"Error loading {name}: {str(e)}")
+        else:
+            st.warning(f"Model file not found: {path}")
+
+    return models
+
+def normalize_url(url):
+    """Safer URL normalization"""
+    try:
+        is_https = url.lower().startswith('https://')
+        url = url.lower()
+
+        prefixes = ['http://', 'ftp://', 'www.', 'ww2.', 'web.']
+        for prefix in prefixes:
+            if url.startswith(prefix):
+                url = url[len(prefix):]
+
+        if is_https:
+            url = "https://" + url
+
+        url = url.split('#')[0]
+
+        if '?' in url:
+            base, query = url.split('?', 1)
+            if not any(sd in base for sd in SAFE_DOMAINS):
+                params = [p for p in query.split('&') if '=' in p]
+                essential_params = [p for p in params if any(
+                    kw in p for kw in ['id=', 'ref=', 'token='])]
+                url = base + ('?' + '&'.join(essential_params)
+        return re.sub(r'/{2,}', '/', url)
+    except Exception:
+        return url
+
+def extract_url_components(url):
+    """Robust URL parsing"""
+    try:
+        extracted = tld_extractor(url)
+        subdomain = extracted.subdomain
+        domain = extracted.domain
+        suffix = extracted.suffix
+
+        path = ""
+        query = ""
+        if "/" in url:
+            path_start = url.find("/", url.find("//") + 2) if "//" in url else url.find("/")
+            if path_start != -1:
+                path_query = url[path_start:]
+                if "?" in path_query:
+                    path, query = path_query.split("?", 1)
+                else:
+                    path = path_query
+
+        if not domain and subdomain:
+            domain_parts = subdomain.split('.')
+            if len(domain_parts) > 1:
+                domain = domain_parts[-1]
+                subdomain = '.'.join(domain_parts[:-1])
+
+        return {
+            'subdomain': subdomain,
+            'domain': domain,
+            'suffix': suffix,
+            'path': path,
+            'query': query
+        }
+    except:
+        return {
+            'subdomain': '',
+            'domain': '',
+            'suffix': '',
+            'path': '',
+            'query': ''
+        }
+
+def calculate_entropy(s):
+    """Compute Shannon entropy"""
+    if not s:
+        return 0
+    try:
+        p, lns = Counter(s), float(len(s))
+        return -sum(count/lns * math.log(count/lns, 2) for count in p.values())
+    except:
+        return 0
+
+def fuzzy_domain_match(domain):
+    """Safe domain matching"""
+    if domain in SAFE_DOMAINS:
+        return True
+
+    domain_parts = domain.split('.')
+    if len(domain_parts) > 2:
+        base_domain = '.'.join(domain_parts[-2:])
+        if base_domain in SAFE_DOMAINS:
+            return True
+
+    best_match, score, _ = process.extractOne(domain, SAFE_DOMAINS, scorer=fuzz.WRatio)
+    return score > 85
+
+def extract_robust_features(url):
+    """Feature extraction optimized for CPU"""
+    try:
+        clean_url = re.sub(r'[^\x00-\x7F]+', '', str(url))
+        normalized = normalize_url(clean_url)
+        components = extract_url_components(clean_url)
+        full_domain = f"{components['domain']}.{components['suffix']}" if components['suffix'] else components['domain']
+
+        # Structural features
+        url_length = len(clean_url)
+        domain_length = len(components['domain'])
+        subdomain_count = len(components['subdomain'].split('.')) if components['subdomain'] else 0
+        path_depth = components['path'].count('/') if components['path'] else 0
+        param_count = len(components['query'].split('&')) if components['query'] else 0
+
+        # Security features
+        has_ip = 1 if re.match(r'^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$', components['domain']) else 0
+        has_executable = 1 if re.search(r'\.(exe|js|jar|bat|sh|py|dll)$', components['path'], re.I) else 0
+        has_double_extension = 1 if re.search(r'\.\w+\.\w+$', components['path'], re.I) else 0
+        hex_encoded = 1 if re.search(r'%[0-9a-f]{2}', normalized, re.I) else 0
+
+        # Lexical features
+        digit_count = sum(c.isdigit() for c in normalized)
+        special_chars = sum(not (c.isalnum() or c in ' ./-') for c in normalized)
+
+        digit_ratio = digit_count / url_length if url_length > 0 else 0
+        special_char_ratio = special_chars / url_length if url_length > 0 else 0
+        entropy = calculate_entropy(normalized)
+
+        # Domain reputation
+        is_safe_domain = 1 if fuzzy_domain_match(full_domain) else 0
+        is_uncommon_tld = 1 if components['suffix'] and components['suffix'] not in COMMON_TLDS else 0
+
+        if url.startswith('https://') and full_domain in SAFE_DOMAINS:
+            is_safe_domain = 1
+
+        return {
+            'url_length': url_length,
+            'domain_length': domain_length,
+            'subdomain_count': subdomain_count,
+            'path_depth': path_depth,
+            'param_count': param_count,
+            'has_ip': has_ip,
+            'has_executable': has_executable,
+            'has_double_extension': has_double_extension,
+            'hex_encoded': hex_encoded,
+            'digit_ratio': digit_ratio,
+            'special_char_ratio': special_char_ratio,
+            'entropy': entropy,
+            'is_safe_domain': is_safe_domain,
+            'is_uncommon_tld': is_uncommon_tld
+        }
+    except Exception as e:
+        st.error(f"Feature extraction error: {str(e)}")
+        return {col: 0 for col in FIXED_FEATURE_COLS}
+
+def preprocess_url(url, char_to_idx):
+    """URL preprocessing for CPU"""
+    try:
+        clean_url = re.sub(r'[^\x00-\x7F]+', '', str(url))
+        normalized = normalize_url(clean_url)
+        features = extract_robust_features(clean_url)
+        feature_vector = np.array([features.get(col, 0) for col in FIXED_FEATURE_COLS]).reshape(1, -1)
+
+        char_seq = [char_to_idx.get(c, 0) for c in normalized]
+        char_seq = pad_sequences([char_seq], maxlen=MAX_LEN, padding='post', truncating='post')
+
+        return char_seq, feature_vector, features
+    except Exception as e:
+        st.error(f"Preprocessing error: {str(e)}")
+        return np.zeros((1, MAX_LEN)), np.zeros((1, len(FIXED_FEATURE_COLS))), {}
+
+def weighted_ensemble_predict(models, char_seq, feature_vector, features):
+    """Ensemble prediction for CPU"""
+    predictions = []
+    weights = {
+        'hybrid': 0.25,
+        'hybrid_fold1': 0.20,
+        'hybrid_fold2': 0.20,
+        'xgboost': 0.35
+    }
+
+    if features.get('is_safe_domain', 0) == 1:
+        return 0.01, [('safe_domain_override', 0.01)]
+
+    for model_name, model in models.items():
+        if model_name in weights:
+            try:
+                if 'hybrid' in model_name:
+                    proba = model.predict([char_seq, feature_vector], verbose=0)[0][0]
+                else:
+                    adjusted_features = feature_vector[:, :14] if feature_vector.shape[1] > 14 else feature_vector
+                    proba = model.predict_proba(adjusted_features)[0][1]
+                predictions.append((model_name, proba))
+            except Exception as e:
+                st.error(f"Prediction error in {model_name}: {str(e)}")
+
+    if predictions:
+        weighted_sum = sum(p * weights.get(name, 0) for name, p in predictions)
+        total_weight = sum(weights.get(name, 0) for name, _ in predictions)
+        avg_proba = weighted_sum / total_weight if total_weight > 0 else sum(p for _, p in predictions) / len(predictions)
+    else:
+        avg_proba = 0.5
+
+    return avg_proba, predictions
+
+def analyze_single_url(url, char_to_idx, models):
+    """Analyze a single URL"""
+    with st.spinner(f"Analyzing URL: {url[:50]}..."):
+        start_time = time.time()
+        
+        char_seq, feature_vector, features = preprocess_url(url, char_to_idx)
+        ensemble_proba, model_predictions = weighted_ensemble_predict(
+            models, char_seq, feature_vector, features)
+        
+        processing_time = time.time() - start_time
+        
+        st.subheader("Analysis Results")
+        col1, col2 = st.columns([1, 2])
+        
+        with col1:
+            if ensemble_proba >= 0.5:
+                st.error(f"🔴 **Threat Detected!** (Probability: {ensemble_proba:.4f})")
+            else:
+                st.success(f"🟢 **Safe URL** (Probability: {ensemble_proba:.4f})")
+            
+            st.metric("Processing Time", f"{processing_time*1000:.2f} ms")
+            
+            st.subheader("Key Features")
+            st.json({
+                "URL Length": features.get('url_length', 0),
+                "Domain Length": features.get('domain_length', 0),
+                "Subdomains": features.get('subdomain_count', 0),
+                "Path Depth": features.get('path_depth', 0),
+                "Parameters": features.get('param_count', 0),
+                "Contains IP": bool(features.get('has_ip', 0)),
+                "Contains Executable": bool(features.get('has_executable', 0)),
+                "Double Extension": bool(features.get('has_double_extension', 0)),
+                "Hex Encoded": bool(features.get('hex_encoded', 0)),
+                "Safe Domain": bool(features.get('is_safe_domain', 0)),
+                "Uncommon TLD": bool(features.get('is_uncommon_tld', 0)),
+                "Entropy": features.get('entropy', 0)
+            })
+        
+        with col2:
+            st.subheader("Model Predictions")
+            model_df = pd.DataFrame(model_predictions, columns=['Model', 'Probability'])
+            model_df['Prediction'] = model_df['Probability'].apply(
+                lambda x: "MALICIOUS" if x >= 0.5 else "SAFE")
+            
+            st.bar_chart(model_df.set_index('Model')['Probability'])
+            
+            st.write("Detailed Model Results:")
+            for model_name, proba in model_predictions:
+                pred = "MALICIOUS" if proba >= 0.5 else "SAFE"
+                st.write(f"- **{model_name}**: {proba:.4f} ({pred})")
+
+def analyze_batch_urls(urls, char_to_idx, models):
+    """Analyze multiple URLs"""
+    results = []
+    progress_bar = st.progress(0)
+    status_text = st.empty()
+    
+    for i, url in enumerate(urls):
+        status_text.text(f"Processing {i+1}/{len(urls)}: {url[:50]}...")
+        progress_bar.progress((i + 1) / len(urls))
+        
+        try:
+            char_seq, feature_vector, features = preprocess_url(url, char_to_idx)
+            ensemble_proba, _ = weighted_ensemble_predict(models, char_seq, feature_vector, features)
+            results.append({
+                'URL': url,
+                'Threat Probability': ensemble_proba,
+                'Classification': "MALICIOUS" if ensemble_proba >= 0.5 else "SAFE"
+            })
+        except Exception as e:
+            st.error(f"Error processing {url}: {str(e)}")
+    
+    if results:
+        results_df = pd.DataFrame(results)
+        st.dataframe(results_df)
+        
+        csv = results_df.to_csv(index=False).encode('utf-8')
+        st.download_button(
+            "Download Results",
+            csv,
+            "url_analysis_results.csv",
+            "text/csv",
+            key='download-csv'
+        )
+
+def main():
+    st.title("🛡️ URL Threat Detector (CPU Version)")
+    st.markdown("""
+    This tool analyzes URLs using machine learning models to detect potential threats.
+    Optimized for CPU-only environments.
+    """)
+    
+    with st.sidebar:
+        st.header("About")
+        st.markdown("""
+        - **Models**: Hybrid CNN+MLP, XGBoost
+        - **Features**: URL structure, lexical patterns, domain reputation
+        - **Environment**: CPU-only
+        """)
+        
+        st.header("Example URLs")
+        st.code("https://paypal-security-alert.com/login")
+        st.code("https://github.com/features/actions")
+    
+    # Load resources
+    with st.spinner("Loading models..."):
+        char_to_idx = load_char_mapping()
+        models = load_all_models()
+    
+    if not char_to_idx or not models:
+        st.error("Failed to load required resources. Please check the model files.")
+        return
+    
+    # URL input
+    st.subheader("Single URL Analysis")
+    url_input = st.text_input("Enter URL to analyze:", 
+                             placeholder="https://example.com")
+    
+    if st.button("Analyze URL") and url_input:
+        analyze_single_url(url_input, char_to_idx, models)
+    
+    # Batch analysis
+    st.subheader("Batch Analysis")
+    uploaded_file = st.file_uploader("Upload a text file with URLs (one per line)", 
+                                   type=['txt', 'csv'])
+    
+    if uploaded_file is not None:
+        urls = [line.decode('utf-8').strip() for line in uploaded_file if line.strip()]
+        if urls and st.button("Analyze All URLs"):
+            analyze_batch_urls(urls, char_to_idx, models)
+
+if __name__ == "__main__":
+    # Configure TensorFlow logging
+    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+    tf.get_logger().setLevel('ERROR')
+    
+    # Run the app
+    main()