Spaces:
Paused
Paused
| import re | |
| import math | |
| import pickle | |
| import requests | |
| import numpy as np | |
| import pandas as pd | |
| from datetime import datetime | |
| from joblib import load | |
| import tensorflow as tf | |
| # ── Load Models ──────────────────────────────────────────────── | |
| xgb_model = load('models/xgb_model_v3.joblib') | |
| cnn_model = tf.keras.models.load_model('models/cnn_model_v3.keras') | |
| with open('models/label_encoder_v3.pkl', 'rb') as f: le = pickle.load(f) | |
| with open('models/xgb_feature_cols_v3.pkl', 'rb') as f: XGB_COLUMNS = pickle.load(f) | |
| with open('models/char2idx_v3.pkl', 'rb') as f: char2idx = pickle.load(f) | |
| with open('models/ensemble_weights_v3.pkl', 'rb') as f: ens_weights = pickle.load(f) | |
| with open('models/tranco_ranks_v3.pkl', 'rb') as f: TRANCO_RANKS = pickle.load(f) | |
| XGB_WEIGHT = ens_weights['xgb'] | |
| CNN_WEIGHT = ens_weights['cnn'] | |
| MAX_URL_LEN = 200 | |
| CONFIDENCE_THRESHOLD = 85.0 | |
| # ── Trusted Root Domains (override — no retraining needed) ──── | |
| TRUSTED_ROOT_DOMAINS = { | |
| 'google.com', 'microsoft.com', 'apple.com', 'amazon.com', | |
| 'youtube.com', 'facebook.com', 'instagram.com', 'twitter.com', | |
| 'x.com', 'github.com', 'gitlab.com', 'stackoverflow.com', | |
| 'linkedin.com', 'reddit.com', 'wikipedia.org', 'mozilla.org', | |
| 'cloudflare.com', 'anthropic.com', 'openai.com', 'netflix.com', | |
| 'spotify.com', 'dropbox.com', 'notion.so', 'slack.com', | |
| 'zoom.us', 'discord.com', 'twitch.tv', 'pinterest.com', | |
| 'huggingface.co', 'kaggle.com', 'pypi.org', 'npmjs.com', | |
| 'tensorflow.org', 'pytorch.org', 'python.org', 'nodejs.org', | |
| 'amazonaws.com', 'googleusercontent.com', 'gstatic.com', | |
| 'googleapis.com', 'googlevideo.com', 'githubusercontent.com', | |
| 'ytimg.com', 'fbcdn.net', 'akamai.net', 'fastly.net', | |
| # Indian | |
| 'irctc.co.in', 'sbi.co.in', 'onlinesbi.sbi', 'gov.in', | |
| 'nic.in', 'paytm.com', 'phonepe.com', 'flipkart.com', | |
| 'zomato.com', 'swiggy.com', 'naukri.com', | |
| } | |
| # ── TLD Lists ────────────────────────────────────────────────── | |
| HIGHLY_TRUSTED_TLDS = ['.gov','.gov.in','.gov.uk','.gov.au', | |
| '.edu','.ac.in','.ac.uk','.mil'] | |
| TRUSTED_TLDS = ['.com','.org','.net','.co.in','.co.uk', | |
| '.com.au','.io','.dev','.app','.ai'] | |
| SUSPICIOUS_TLDS = ['.tk','.ml','.ga','.cf','.gq','.xyz','.top', | |
| '.club','.online','.site','.website','.space', | |
| '.click','.link','.download','.win','.bid', | |
| '.loan','.zip','.date','.faith','.racing'] | |
| # ── Cache ────────────────────────────────────────────────────── | |
| _domain_cache = {} | |
| # ── Helpers ──────────────────────────────────────────────────── | |
| def _char_entropy(s): | |
| if not s: return 0 | |
| freq = [s.count(c)/len(s) for c in set(s)] | |
| return -sum(p*math.log2(p) for p in freq) | |
| def _url_to_sequence(url): | |
| return [char2idx.get(c,1) for c in str(url).lower()[:MAX_URL_LEN]] | |
| def _pad(seq): | |
| arr = np.zeros((1, MAX_URL_LEN), dtype=np.int32) | |
| arr[0, :min(len(seq), MAX_URL_LEN)] = seq[:MAX_URL_LEN] | |
| return arr | |
| def _get_tranco_score(domain): | |
| rank = TRANCO_RANKS.get(domain) | |
| if rank is None: return 0.0 | |
| return round(1 - (rank / 100000), 4) | |
| # ── Feature Extractor (v3 — matches training) ────────────────── | |
| def extract_features(url): | |
| f = {} | |
| url = str(url).strip() | |
| f['IsHttps'] = 1 if url.startswith('https://') else 0 | |
| f['IsHttp'] = 1 if url.startswith('http://') else 0 | |
| f['NumDots'] = url.count('.') | |
| f['NumDash'] = url.count('-') | |
| f['NumUnderscore'] = url.count('_') | |
| f['NumPercent'] = url.count('%') | |
| f['NumAmpersand'] = url.count('&') | |
| f['NumHash'] = url.count('#') | |
| f['NumNumericChars'] = sum(c.isdigit() for c in url) | |
| f['UrlLength'] = len(url) | |
| f['NumQueryComponents'] = url.count('=') | |
| f['HasExe'] = 1 if '.exe' in url.lower() else 0 | |
| f['HasPhp'] = 1 if '.php' in url.lower() else 0 | |
| f['HasHtml'] = 1 if '.html' in url.lower() else 0 | |
| f['HasDll'] = 1 if '.dll' in url.lower() else 0 | |
| f['HasSh'] = 1 if '.sh' in url.lower() else 0 | |
| try: | |
| domain = url.split('/')[2] if '//' in url else url.split('/')[0] | |
| except: | |
| domain = url | |
| parts = domain.split('.') | |
| registered_domain = '.'.join(parts[-2:]) if len(parts) >= 2 else domain | |
| subdomain = '.'.join(parts[:-2]) if len(parts) >= 2 else '' | |
| core_domain = registered_domain.split('.')[0] | |
| tld = '.' + parts[-1] if parts else '' | |
| co_tld = '.' + '.'.join(parts[-2:]) if len(parts) >= 2 else tld | |
| tri_tld = '.' + '.'.join(parts[-3:]) if len(parts) >= 3 else co_tld | |
| f['IsHighlyTrustedTLD'] = 1 if any(t in [tld,co_tld,tri_tld] | |
| for t in HIGHLY_TRUSTED_TLDS) else 0 | |
| f['IsTrustedTLD'] = 1 if any(t in [tld,co_tld] | |
| for t in TRUSTED_TLDS) else 0 | |
| f['IsSuspiciousTLD'] = 1 if any(t in [tld,co_tld] | |
| for t in SUSPICIOUS_TLDS) else 0 | |
| f['TLDTrustScore'] = (2 if f['IsHighlyTrustedTLD'] else | |
| (1 if f['IsTrustedTLD'] else | |
| (-1 if f['IsSuspiciousTLD'] else 0))) | |
| f['TrancoRankScore'] = _get_tranco_score(registered_domain) | |
| f['InTrancoTop100k'] = 1 if TRANCO_RANKS.get(registered_domain) else 0 | |
| f['InTrancoTop10k'] = 1 if (TRANCO_RANKS.get(registered_domain) | |
| or 999999) <= 10000 else 0 | |
| f['InTrancoTop1k'] = 1 if (TRANCO_RANKS.get(registered_domain) | |
| or 999999) <= 1000 else 0 | |
| f['HostnameLength'] = len(domain) | |
| f['NumDashInHostname'] = domain.count('-') | |
| f['SubdomainLevel'] = max(domain.count('.')-1, 0) | |
| f['CoreDomainLength'] = len(core_domain) | |
| f['DigitRatioInDomain'] = sum(c.isdigit() for c in domain) / max(len(domain),1) | |
| f['IsURLShortener'] = 1 if any(s in domain for s in | |
| ['bit.ly','tinyurl','t.co','ow.ly']) else 0 | |
| f['HasConsecutiveDigits'] = 1 if re.search(r'\d{4,}', domain) else 0 | |
| f['DomainInSubdomains'] = 1 if any(b in subdomain for b in | |
| ['paypal','google','facebook','apple', | |
| 'microsoft','amazon','sbi','bank']) else 0 | |
| f['RandomString'] = 1 if _char_entropy(core_domain) > 3.5 else 0 | |
| f['ConsonantRatio'] = sum(c in 'bcdfghjklmnpqrstvwxyz' | |
| for c in domain.lower()) / max(len(domain),1) | |
| f['VowelRatio'] = sum(c in 'aeiou' | |
| for c in domain.lower()) / max(len(domain),1) | |
| f['HasPortNumber'] = 1 if re.search(r':\d+', url) else 0 | |
| f['NumSubdomains'] = len(subdomain.split('.')) if subdomain else 0 | |
| f['IpAddress'] = 1 if re.match(r'\d+\.\d+\.\d+\.\d+', | |
| domain) else 0 | |
| try: | |
| path = '/' + '/'.join(url.split('/')[3:]) | |
| except: | |
| path = '/' | |
| f['PathLength'] = len(path) | |
| f['PathLevel'] = path.count('/') | |
| f['DomainInPaths'] = 1 if any(b in path for b in | |
| ['paypal','google','facebook', | |
| 'apple','microsoft','amazon','sbi']) else 0 | |
| query = url.split('?')[1] if '?' in url else '' | |
| f['QueryLength'] = len(query) | |
| f['AtSymbolInDomain'] = 1 if '@' in domain else 0 | |
| f['AtSymbolInPath'] = 1 if '@' in path else 0 | |
| f['TildeSymbol'] = 1 if '~' in url else 0 | |
| f['DoubleSlashInPath'] = 1 if '//' in path else 0 | |
| f['HttpsInHostname'] = 1 if 'https' in domain else 0 | |
| sensitive = ['secure','login','verify','account','update', | |
| 'banking','confirm','password','signin','webscr','paypal'] | |
| f['NumSensitiveWords'] = sum(w in url.lower() for w in sensitive) | |
| f['EmbeddedBrandName'] = f['DomainInSubdomains'] | |
| f['UrlEntropy'] = _char_entropy(url) | |
| return f | |
| # ── Risk Signals (for UI explanation) ───────────────────────── | |
| def get_risk_signals(url, label): | |
| signals = [] | |
| try: | |
| domain = url.split('/')[2] if '//' in url else url.split('/')[0] | |
| except: | |
| domain = url | |
| parts = domain.split('.') | |
| registered_domain = '.'.join(parts[-2:]) if len(parts) >= 2 else domain | |
| subdomain = '.'.join(parts[:-2]) if len(parts) >= 2 else '' | |
| tld = '.' + parts[-1] if parts else '' | |
| co_tld = '.' + '.'.join(parts[-2:]) if len(parts) >= 2 else tld | |
| if not url.startswith('https://'): | |
| signals.append({'icon': '🔓', 'text': 'No HTTPS encryption', 'risk': 'high'}) | |
| if any(t in [tld,co_tld] for t in SUSPICIOUS_TLDS): | |
| signals.append({'icon': '⚠️', 'text': f'Suspicious TLD: {tld}', 'risk': 'high'}) | |
| if re.match(r'\d+\.\d+\.\d+\.\d+', domain): | |
| signals.append({'icon': '🖥️', 'text': 'IP address used instead of domain', 'risk': 'high'}) | |
| if re.search(r':\d+', url): | |
| signals.append({'icon': '🔌', 'text': 'Non-standard port detected', 'risk': 'medium'}) | |
| if any(b in subdomain for b in ['paypal','google','facebook','apple', | |
| 'microsoft','amazon','sbi','bank']): | |
| signals.append({'icon': '🎭', 'text': 'Brand name in subdomain (spoofing)', 'risk': 'high'}) | |
| if len(url) > 150: | |
| signals.append({'icon': '📏', 'text': f'Unusually long URL ({len(url)} chars)', 'risk': 'medium'}) | |
| sensitive = ['secure','login','verify','account','update', | |
| 'banking','confirm','password'] | |
| found = [w for w in sensitive if w in url.lower()] | |
| if found: | |
| signals.append({'icon': '🎣', 'text': f'Sensitive words: {", ".join(found)}', 'risk': 'medium'}) | |
| if TRANCO_RANKS.get(registered_domain): | |
| rank = TRANCO_RANKS[registered_domain] | |
| signals.append({'icon': '✅', | |
| 'text': f'Domain ranked #{rank:,} in Tranco top 100k', | |
| 'risk': 'safe'}) | |
| if '.exe' in url.lower() or '.dll' in url.lower() or '.sh' in url.lower(): | |
| signals.append({'icon': '☠️', 'text': 'Executable file in URL', 'risk': 'high'}) | |
| return signals | |
| # ── RDAP Domain Age Lookup ───────────────────────────────────── | |
| def rdap_lookup(domain): | |
| if domain in _domain_cache: | |
| return _domain_cache[domain] | |
| result = {'age_days': -1, 'verdict': 'unknown', 'reason': 'Lookup failed'} | |
| try: | |
| resp = requests.get( | |
| f"https://rdap.org/domain/{domain}", | |
| timeout=6, headers={'User-Agent': 'Mozilla/5.0'} | |
| ) | |
| if resp.status_code == 200: | |
| created = None | |
| for event in resp.json().get('events', []): | |
| if event.get('eventAction') == 'registration': | |
| created = event.get('eventDate','')[:10] | |
| break | |
| if created: | |
| age = (datetime.now() - datetime.strptime(created, '%Y-%m-%d')).days | |
| result['age_days'] = age | |
| if age < 30: | |
| result.update(verdict='phishing', | |
| reason=f'Only {age} days old — brand new domain') | |
| elif age < 180: | |
| result.update(verdict='phishing', | |
| reason=f'Only {age} days old — recently created') | |
| elif age > 730: | |
| result.update(verdict='benign', | |
| reason=f'{age//365} year old domain — established') | |
| else: | |
| result.update(verdict='uncertain', | |
| reason=f'{age} days old — moderate age') | |
| else: | |
| result['reason'] = 'No creation date in RDAP' | |
| except Exception as e: | |
| result['reason'] = f'RDAP error: {str(e)[:60]}' | |
| _domain_cache[domain] = result | |
| return result | |
| # ── Main Prediction Pipeline ─────────────────────────────────── | |
| DANGEROUS_LABELS = {'phishing', 'malware'} | |
| def predict_url(url): | |
| try: | |
| domain = url.split('/')[2] if '//' in url else url.split('/')[0] | |
| domain = domain.lower().replace('www.', '') | |
| except: | |
| domain = url | |
| parts = domain.split('.') | |
| registered_domain = '.'.join(parts[-2:]) if len(parts) >= 2 else domain | |
| # ── Override 1: Mega trusted domains (Google, MS, Apple etc) ─ | |
| # Handles deep subdomains: colab.research.google.com etc. | |
| if registered_domain in TRUSTED_ROOT_DOMAINS: | |
| rank = TRANCO_RANKS.get(registered_domain, 0) | |
| rank_text = (f'Ranked #{rank:,} in Tranco top 100k' | |
| if rank else 'Globally recognized domain') | |
| return { | |
| 'verdict': 'SAFE', | |
| 'verdict_class': 'safe', | |
| 'class_probs': {'benign': 99.5, 'phishing': 0.3, 'malware': 0.2}, | |
| 'risk_signals': [ | |
| {'icon': '✅', | |
| 'text': f'Verified trusted domain: {registered_domain}', | |
| 'risk': 'safe'}, | |
| {'icon': '✅', | |
| 'text': rank_text, | |
| 'risk': 'safe'}, | |
| ], | |
| 'xgb': {'result': 'BENIGN', 'label': 'benign', 'confidence': 99.5}, | |
| 'cnn': {'result': 'BENIGN', 'label': 'benign'}, | |
| 'ensemble': { | |
| 'result': 'BENIGN', | |
| 'label': 'benign', | |
| 'confidence': 99.5, | |
| 'source': 'TrustedDomain', | |
| 'reason': f'{registered_domain} is a globally verified trusted domain', | |
| }, | |
| 'rdap': None, | |
| } | |
| # ── Override 2: Tranco top 100k + https + trusted TLD ──────── | |
| # Fixes false positives on legitimate short/random-looking domains | |
| # that are in Tranco (e.g. lgeapi.com, dtjm6.com) | |
| rank = TRANCO_RANKS.get(registered_domain) | |
| if rank and rank <= 100000: | |
| try: | |
| host = url.split('/')[2] if '//' in url else url | |
| except: | |
| host = url | |
| is_https = url.startswith('https://') | |
| has_trusted = any(host.endswith(t) | |
| for t in TRUSTED_TLDS + HIGHLY_TRUSTED_TLDS) | |
| no_executable = not any(x in url.lower() | |
| for x in ['.exe','.dll','.sh','.bat','.bin']) | |
| no_ip = not re.match(r'\d+\.\d+\.\d+\.\d+', host) | |
| no_port = not re.search(r':\d+', host) | |
| if is_https and has_trusted and no_executable and no_ip and no_port: | |
| rank_text = (f'Ranked #{rank:,} in Tranco top 100k — ' | |
| f'legitimate high-traffic domain') | |
| return { | |
| 'verdict': 'SAFE', | |
| 'verdict_class': 'safe', | |
| 'class_probs': {'benign': 95.0, 'phishing': 3.0, 'malware': 2.0}, | |
| 'risk_signals': [ | |
| {'icon': '✅', 'text': rank_text, 'risk': 'safe'}, | |
| {'icon': '✅', 'text': 'HTTPS encryption present', 'risk': 'safe'}, | |
| {'icon': '✅', 'text': 'Trusted TLD confirmed', 'risk': 'safe'}, | |
| ], | |
| 'xgb': {'result': 'BENIGN', 'label': 'benign', 'confidence': 95.0}, | |
| 'cnn': {'result': 'BENIGN', 'label': 'benign'}, | |
| 'ensemble': { | |
| 'result': 'BENIGN', | |
| 'label': 'benign', | |
| 'confidence': 95.0, | |
| 'source': 'TrancoRank', | |
| 'reason': rank_text, | |
| }, | |
| 'rdap': None, | |
| } | |
| # ── XGBoost ─────────────────────────────────────────────────── | |
| xgb_in = pd.DataFrame([extract_features(url)]).reindex( | |
| columns=XGB_COLUMNS, fill_value=0) | |
| xgb_prob = xgb_model.predict_proba(xgb_in)[0] | |
| xgb_label = le.classes_[int(np.argmax(xgb_prob))] | |
| xgb_conf = round(float(max(xgb_prob)) * 100, 1) | |
| # ── CNN ─────────────────────────────────────────────────────── | |
| cnn_prob = cnn_model.predict(_pad(_url_to_sequence(url)), verbose=0)[0] | |
| cnn_label = le.classes_[int(np.argmax(cnn_prob))] | |
| # ── Ensemble ────────────────────────────────────────────────── | |
| final_prob = (xgb_prob * XGB_WEIGHT) + (cnn_prob * CNN_WEIGHT) | |
| ensemble_label = le.classes_[int(np.argmax(final_prob))] | |
| ensemble_conf = round(float(max(final_prob)) * 100, 1) | |
| class_probs = {le.classes_[i]: round(float(final_prob[i])*100, 1) | |
| for i in range(len(le.classes_))} | |
| # ── RDAP fallback for low confidence ────────────────────────── | |
| rdap_info = None | |
| source = 'Ensemble' | |
| reason = '' | |
| if ensemble_conf < CONFIDENCE_THRESHOLD: | |
| rdap_info = rdap_lookup(domain) | |
| if rdap_info['verdict'] == 'benign': | |
| ensemble_label = 'benign' | |
| ensemble_conf = min(ensemble_conf + 20, 95.0) | |
| source = 'RDAP' | |
| reason = rdap_info['reason'] | |
| elif rdap_info['verdict'] == 'phishing': | |
| ensemble_label = 'phishing' | |
| ensemble_conf = min(ensemble_conf + 15, 95.0) | |
| source = 'RDAP' | |
| reason = rdap_info['reason'] | |
| verdict = 'DANGEROUS' if ensemble_label in DANGEROUS_LABELS else 'SAFE' | |
| verdict_cls = 'dangerous' if ensemble_label in DANGEROUS_LABELS else 'safe' | |
| risk_signals = get_risk_signals(url, ensemble_label) | |
| return { | |
| 'verdict': verdict, | |
| 'verdict_class': verdict_cls, | |
| 'class_probs': class_probs, | |
| 'risk_signals': risk_signals, | |
| 'xgb': {'result': xgb_label.upper(), 'label': xgb_label, 'confidence': xgb_conf}, | |
| 'cnn': {'result': cnn_label.upper(), 'label': cnn_label}, | |
| 'ensemble': { | |
| 'result': ensemble_label.upper(), | |
| 'label': ensemble_label, | |
| 'confidence': ensemble_conf, | |
| 'source': source, | |
| 'reason': reason, | |
| }, | |
| 'rdap': rdap_info, | |
| } |