YanushanthR/URLs_Classifier

1import streamlit as st2import pandas as pd3import joblib4from urllib.parse import urlparse5from math import log26import re7import os8 9# Load the trained model and feature names10model_path = "RF_trained_model.pkl"11feature_names_path = "feature_names.pkl"12 13if not os.path.exists(model_path) or not os.path.exists(feature_names_path):14    st.error("Model or feature names file not found. Please check the paths.")15    st.stop()16 17model = joblib.load(model_path)18feature_names = joblib.load(feature_names_path)19 20# Feature extraction function (same as used for training)21def extract_features(url):22    if not isinstance(url, str) or pd.isna(url):23        return {24            "url_length": 0,25            "num_digits": 0,26            "num_special_chars": 0,27            "has_https": 0,28            "num_subdomains": 0,29            "has_suspicious_keywords": 0,30            "url_entropy": 0,31            "domain_length": 0,32            "path_length": 0,33            "presence_of_ip": 0,34            "tld": "",35            "num_query_params": 0,36            "has_encoded_chars": 0,37            "path_depth": 0,38            "has_suspicious_substrings": 0,39            "has_malicious_file_extension": 040        }41 42    parsed_url = urlparse(url)43    path = parsed_url.path44    query = parsed_url.query45 46    # Calculate entropy47    try:48        probabilities = [float(url.count(c)) / len(url) for c in set(url)]49        url_entropy = -sum(p * log2(p) for p in probabilities)50    except ValueError:51        url_entropy = 052 53    # Check for IP address54    def contains_ip(url):55        ip_pattern = re.compile(r"(?:\\d{1,3}\\.){3}\\d{1,3}")56        return int(bool(ip_pattern.search(url)))57 58    # Extract TLD59    tld = parsed_url.netloc.split('.')[-1] if '.' in parsed_url.netloc else ""60 61    # Check for encoded characters62    def has_encoded_chars(url):63        return int("%" in url)64 65    # Check for suspicious substrings66    def has_suspicious_substrings(url):67        suspicious_substrings = ["@", "-", "_", "~"]68        return int(any(substring in url for substring in suspicious_substrings))69 70    # Check for malicious file extensions71    def has_malicious_file_extension(url):72        malicious_extensions = [".exe", ".zip", ".js", ".rar", ".bat"]73        return int(any(url.lower().endswith(ext) for ext in malicious_extensions))74 75    return {76        "url_length": len(url),77        "num_digits": sum(c.isdigit() for c in url),78        "num_special_chars": sum(not c.isalnum() for c in url),79        "has_https": int("https" in url.lower()),80        "num_subdomains": url.count('.'),81        "has_suspicious_keywords": int(any(keyword in url.lower() for keyword in ["login", "secure", "account", "update"])),82        "url_entropy": url_entropy,83        "domain_length": len(parsed_url.netloc),84        "path_length": len(path),85        "presence_of_ip": contains_ip(url),86        "tld": tld,87        "num_query_params": len(query.split('&')) if query else 0,88        "has_encoded_chars": has_encoded_chars(url),89        "path_depth": len(path.split('/')) - 1 if path else 0,90        "has_suspicious_substrings": has_suspicious_substrings(url),91        "has_malicious_file_extension": has_malicious_file_extension(url)92    }93 94# Streamlit web app95st.title("URL Classifier")96st.subheader("Enter the URL or paste the URL below")97 98# URL input bar99url_input = st.text_input("Enter URL:", "")100 101if st.button("Classify"):102    if url_input:103        # Extract features104        features = extract_features(url_input)105        features_df = pd.DataFrame([features])106 107        # Ensure compatibility with the model108        features_df = pd.get_dummies(features_df, columns=["tld"], drop_first=True)109 110        # Align with model's feature names111        for col in feature_names:112            if col not in features_df.columns:113                features_df[col] = 0  # Add missing columns with default value114 115        features_df = features_df[feature_names]  # Ensure column order matches116 117        # Predict using the loaded model118        try:119            prediction = model.predict(features_df)[0]120 121            # Display the result122            if prediction == "benign":123                st.success("The URL is classified as **Benign**.")124            else:125                st.error("The URL is classified as **Malicious**.")126        except Exception as e:127            st.error(f"An error occurred during prediction: {e}")128    else:129        st.warning("Please enter a URL to classify.")130 131