import gradio as gr import joblib import json import numpy as np import re from urllib.parse import urlparse import os from huggingface_hub import hf_hub_download # Define the model and username MODEL_NAME = "XGBoost" HF_USERNAME = "Devishetty100" CUSTOM_MODEL_NAME = "NeoGuardianAI" REPO_ID = f"{HF_USERNAME}/{CUSTOM_MODEL_NAME.lower()}" # List of trusted domains that should always be considered safe TRUSTED_DOMAINS = [ 'huggingface.co', 'github.com', 'google.com', 'microsoft.com', 'apple.com', 'amazon.com', 'facebook.com', 'twitter.com', 'linkedin.com', 'youtube.com', 'wikipedia.org' ] # Load model files (either from local files or Hugging Face Hub) def load_model_files(): try: print(f"Attempting to download model from Hugging Face Hub: {REPO_ID}") # Try to list files in the repository to see what's available try: from huggingface_hub import list_repo_files files = list_repo_files(repo_id=REPO_ID) print(f"Files available in the repository: {files}") except Exception as list_error: print(f"Error listing repository files: {list_error}") # Use lowercase 'xgboost' instead of MODEL_NAME.lower() to match the actual filename model_path = hf_hub_download(repo_id=REPO_ID, filename="xgboost_model.joblib") print(f"Downloaded model file to: {model_path}") scaler_path = hf_hub_download(repo_id=REPO_ID, filename="scaler.joblib") feature_names_path = hf_hub_download(repo_id=REPO_ID, filename="feature_names.json") # Load the model and preprocessing components model = joblib.load(model_path) scaler = joblib.load(scaler_path) # Load feature names with open(feature_names_path, 'r') as f: feature_names = json.load(f) print("Successfully downloaded model from Hugging Face Hub.") return model, scaler, feature_names except Exception as hub_error: print(f"Error downloading from Hugging Face Hub: {hub_error}") # If downloading fails, try to load from local files try: print("Attempting to load model from local files...") # Try with the correct lowercase name model = joblib.load("xgboost_model.joblib") print("Successfully loaded xgboost_model.joblib") scaler = joblib.load("scaler.joblib") with open("feature_names.json", 'r') as f: feature_names = json.load(f) print("Successfully loaded model from local files.") return model, scaler, feature_names except Exception as local_error: print(f"Error loading from local files: {local_error}") raise RuntimeError("Failed to load model from both Hugging Face Hub and local files.") # Extract features from URL def extract_features(url): """Extract features from a URL for model prediction.""" features = {} # Basic URL properties features['length_url'] = len(url) # Parse URL parsed_url = urlparse(url) hostname = parsed_url.netloc path = parsed_url.path # Hostname features features['length_hostname'] = len(hostname) features['ip'] = 1 if re.match(r'\d+\.\d+\.\d+\.\d+', hostname) else 0 # Count special characters features['nb_dots'] = url.count('.') features['nb_hyphens'] = url.count('-') features['nb_at'] = url.count('@') features['nb_qm'] = url.count('?') features['nb_and'] = url.count('&') features['nb_or'] = url.count('|') features['nb_eq'] = url.count('=') features['nb_underscore'] = url.count('_') features['nb_tilde'] = url.count('~') features['nb_percent'] = url.count('%') features['nb_slash'] = url.count('/') features['nb_star'] = url.count('*') features['nb_colon'] = url.count(':') features['nb_comma'] = url.count(',') features['nb_semicolumn'] = url.count(';') features['nb_dollar'] = url.count('$') features['nb_space'] = url.count(' ') # Other URL features features['nb_www'] = 1 if 'www' in hostname else 0 features['nb_com'] = 1 if '.com' in hostname else 0 features['nb_dslash'] = url.count('//') features['http_in_path'] = 1 if 'http' in path else 0 features['https_token'] = 1 if 'https' in url and 'http://' not in url else 0 # Ratio features digits_count = sum(c.isdigit() for c in url) features['ratio_digits_url'] = digits_count / len(url) if len(url) > 0 else 0 features['ratio_digits_host'] = sum(c.isdigit() for c in hostname) / len(hostname) if len(hostname) > 0 else 0 # Punycode features['punycode'] = 1 if 'xn--' in hostname else 0 # Port features['port'] = 1 if ':' in hostname and any(c.isdigit() for c in hostname.split(':')[1]) else 0 # TLD features tlds = ['.com', '.org', '.net', '.edu', '.gov', '.mil', '.int'] features['tld_in_path'] = 1 if any(tld in path for tld in tlds) else 0 features['tld_in_subdomain'] = 1 if hostname.count('.') > 1 and any(tld in hostname.split('.')[0] for tld in tlds) else 0 # Subdomain features features['abnormal_subdomain'] = 1 if hostname.count('.') > 2 else 0 features['nb_subdomains'] = hostname.count('.') # Other suspicious features features['prefix_suffix'] = 1 if '-' in hostname else 0 features['random_domain'] = 1 if len(hostname) > 12 and sum(c.isdigit() for c in hostname) > 4 else 0 # Shortening service shortening_services = ['bit.ly', 'goo.gl', 'tinyurl.com', 't.co', 'tr.im', 'is.gd', 'cli.gs', 'ow.ly', 'yfrog.com', 'migre.me', 'ff.im', 'tiny.cc', 'url4.eu', 'twit.ac', 'su.pr', 'twurl.nl', 'snipurl.com', 'short.to', 'budurl.com', 'ping.fm', 'post.ly', 'just.as', 'bkite.com', 'snipr.com', 'fic.kr', 'loopt.us', 'doiop.com', 'twitthis.com', 'htxt.it', 'ak.im', 'shar.es', 'kl.am', 'wp.me', 'rubyurl.com', 'om.ly', 'to.ly', 'bit.do', 't.co', 'lnkd.in', 'db.tt', 'qr.ae', 'adf.ly', 'goo.gl', 'bitly.com', 'cur.lv', 'tinyurl.com', 'ow.ly', 'bit.ly', 'ity.im', 'q.gs', 'is.gd', 'po.st', 'bc.vc', 'twitthis.com', 'u.to', 'j.mp', 'buzurl.com', 'cutt.us', 'u.bb', 'yourls.org', 'x.co', 'prettylinkpro.com', 'scrnch.me', 'filoops.info', 'vzturl.com', 'qr.net', '1url.com', 'tweez.me', 'v.gd', 'tr.im', 'link.zip.net'] features['shortening_service'] = 1 if any(service in hostname for service in shortening_services) else 0 # Path features features['path_extension'] = 1 if '.' in path.split('/')[-1] else 0 # Fill in remaining features with default values # These would normally be computed with more complex analysis for feature in ['nb_redirection', 'nb_external_redirection', 'length_words_raw', 'char_repeat', 'shortest_words_raw', 'shortest_word_host', 'shortest_word_path', 'longest_words_raw', 'longest_word_host', 'longest_word_path', 'avg_words_raw', 'avg_word_host', 'avg_word_path', 'phish_hints', 'domain_in_brand', 'brand_in_subdomain', 'brand_in_path', 'suspecious_tld', 'statistical_report', 'nb_hyperlinks', 'ratio_intHyperlinks', 'ratio_extHyperlinks', 'ratio_nullHyperlinks', 'nb_extCSS', 'ratio_intRedirection', 'ratio_extRedirection', 'ratio_intErrors', 'ratio_extErrors', 'login_form', 'external_favicon', 'links_in_tags', 'submit_email', 'ratio_intMedia', 'ratio_extMedia', 'sfh', 'iframe', 'popup_window', 'safe_anchor', 'onmouseover', 'right_clic', 'empty_title', 'domain_in_title', 'domain_with_copyright', 'whois_registered_domain', 'domain_registration_length', 'domain_age', 'web_traffic', 'dns_record', 'google_index', 'page_rank']: if feature not in features: features[feature] = 0 return features # Load model and components try: model, scaler, feature_names = load_model_files() print("Model loaded successfully!") except Exception as e: print(f"Error loading model: {e}") # Create dummy model and components for demo purposes print("Using dummy model for demonstration purposes.") import numpy as np from sklearn.ensemble import RandomForestClassifier # Create a dummy model model = RandomForestClassifier(n_estimators=10) model.fit(np.array([[0, 0]]), np.array([0])) model.predict_proba = lambda x: np.array([[0.5, 0.5]]) # Create dummy scaler and feature names scaler = lambda x: x scaler.transform = lambda x: x feature_names = ['length_url', 'length_hostname'] def predict_url(url): """Predict if a URL is phishing or legitimate.""" if not url or not url.strip(): return "Please enter a URL", 0.0, "N/A" try: # Check if the URL belongs to a trusted domain parsed_url = urlparse(url) domain = parsed_url.netloc # Remove 'www.' prefix if present if domain.startswith('www.'): domain = domain[4:] # Check if the domain or any parent domain is in the trusted list is_trusted = False domain_parts = domain.split('.') for i in range(len(domain_parts) - 1): check_domain = '.'.join(domain_parts[i:]) if check_domain in TRUSTED_DOMAINS: is_trusted = True break if is_trusted: return "Legitimate (Trusted Domain)", 1.0, "✅ SAFE" # Extract features url_features = extract_features(url) # Ensure features are in the correct order features_array = [] for feature in feature_names: if feature in url_features: features_array.append(url_features[feature]) else: features_array.append(0) # Default value if feature is missing # Scale features scaled_features = scaler.transform([features_array]) # Make prediction prediction = model.predict(scaled_features)[0] probability = model.predict_proba(scaled_features)[0][1] # Prepare return values prediction_text = "Phishing" if prediction == 1 else "Legitimate" confidence = float(probability) if prediction == 1 else float(1 - probability) status = "⚠️ UNSAFE" if prediction == 1 else "✅ SAFE" # Return three separate values for the three output components return prediction_text, confidence, status except Exception as e: error_msg = f"Error: {str(e)}" return error_msg, 0.0, "Error" # Create Gradio interface def create_interface(): with gr.Blocks(title="NeoGuardianAI - URL Phishing Detection", theme=gr.themes.Soft()) as demo: gr.Markdown( """ # NeoGuardianAI - URL Phishing Detection This app uses a machine learning model to detect if a URL is legitimate or phishing. Enter a URL below to check if it's safe or potentially malicious. """ ) with gr.Row(): url_input = gr.Textbox(label="Enter URL", placeholder="https://example.com") submit_btn = gr.Button("Check URL", variant="primary") with gr.Row(): status_output = gr.Textbox(label="Status") prediction_output = gr.Textbox(label="Prediction") confidence_output = gr.Textbox(label="Confidence") submit_btn.click( fn=predict_url, inputs=url_input, outputs=[ prediction_output, confidence_output, status_output ] ) gr.Markdown( """ ## How it works This model was trained on the [pirocheto/phishing-url](https://huggingface.co/datasets/pirocheto/phishing-url) dataset from Hugging Face. The model extracts various features from the URL and uses a machine learning algorithm to classify it as legitimate or phishing. **Note**: While this model is highly accurate, it's not perfect. Always exercise caution when visiting unfamiliar websites. ## API Usage You can also use this model via the Hugging Face Inference API: ```python import requests API_URL = "https://api-inference.huggingface.co/models/Devishetty100/neoguardianai" headers = {"Authorization": "Bearer YOUR_API_TOKEN"} def query(url): payload = {"inputs": url} response = requests.post(API_URL, headers=headers, json=payload) return response.json() # Example result = query("https://example.com") print(result) ``` """ ) return demo # Launch the app if __name__ == "__main__": demo = create_interface() demo.launch()