ZealPyae/V2LinkDetection

1from fastapi import FastAPI2from pydantic import BaseModel3import joblib4import numpy as np5from urllib.parse import urlparse6from tld import get_tld7import re8 9# FastAPI instance10app = FastAPI()11 12# Load your trained model13model = joblib.load("rf_model.pkl")  # Ensure you save your RandomForest model as rf_model.pkl14 15# Define the request body16class URLRequest(BaseModel):17    url: str18 19# Feature extraction functions20def having_ip_address(url):21    match = re.search(22        '(([01]?\\d\\d?|2[0-4]\\d|25[0-5])\\.([01]?\\d\\d?|2[0-4]\\d|25[0-5])\\.([01]?\\d\\d?|2[0-4]\\d|25[0-5])\\.'23        '([01]?\\d\\d?|2[0-4]\\d|25[0-5])\\/)|'  # IPv424        '((0x[0-9a-fA-F]{1,2})\\.(0x[0-9a-fA-F]{1,2})\\.(0x[0-9a-fA-F]{1,2})\\.(0x[0-9a-fA-F]{1,2})\\/)'  # IPv4 in hexadecimal25        '(?:[a-fA-F0-9]{1,4}:){7}[a-fA-F0-9]{1,4}', url)  # Ipv626    return 1 if match else 027 28def abnormal_url(url):29    hostname = urlparse(url).hostname30    match = re.search(str(hostname), url)31    return 1 if match else 032 33def count_dot(url):34    return url.count('.')35 36def count_www(url):37    return url.count('www')38 39def count_atrate(url):40    return url.count('@')41 42def no_of_dir(url):43    return urlparse(url).path.count('/')44 45def no_of_embed(url):46    return urlparse(url).path.count('//')47 48def shortening_service(url):49    match = re.search('bit\.ly|goo\.gl|shorte\.st|go2l\.ink|x\.co|ow\.ly|t\.co|tinyurl|tr\.im|is\.gd|cli\.gs|'50                      'yfrog\.com|migre\.me|ff\.im|tiny\.cc|url4\.eu|twit\.ac|su\.pr|twurl\.nl|snipurl\.com|'51                      'short\.to|BudURL\.com|ping\.fm|post\.ly|Just\.as|bkite\.com|snipr\.com|fic\.kr|loopt\.us|'52                      'doiop\.com|short\.ie|kl\.am|wp\.me|rubyurl\.com|om\.ly|to\.ly|bit\.do|t\.co|lnkd\.in|'53                      'db\.tt|qr\.ae|adf\.ly|goo\.gl|bitly\.com|cur\.lv|tinyurl\.com|ow\.ly|bit\.ly|ity\.im|'54                      'q\.gs|is\.gd|po\.st|bc\.vc|twitthis\.com|u\.to|j\.mp|buzurl\.com|cutt\.us|u\.bb|yourls\.org|'55                      'x\.co|prettylinkpro\.com|scrnch\.me|filoops\.info|vzturl\.com|qr\.net|1url\.com|tweez\.me|v\.gd|'56                      'tr\.im|link\.zip\.net', url)57    return 1 if match else 058 59def count_https(url):60    return url.count('https')61 62def count_http(url):63    return url.count('http')64 65def count_per(url):66    return url.count('%')67 68def count_ques(url):69    return url.count('?')70 71def count_hyphen(url):72    return url.count('-')73 74def count_equal(url):75    return url.count('=')76 77def url_length(url):78    return len(str(url))79 80def hostname_length(url):81    return len(urlparse(url).netloc)82 83def suspicious_words(url):84    match = re.search('PayPal|login|signin|bank|account|update|free|lucky|service|bonus|ebayisapi|webscr', url)85    return 1 if match else 086 87def digit_count(url):88    return sum(1 for i in url if i.isnumeric())89 90def letter_count(url):91    return sum(1 for i in url if i.isalpha())92 93def fd_length(url):94    urlpath = urlparse(url).path95    try:96        return len(urlpath.split('/')[1])97    except:98        return 099 100def tld_length(tld):101    try:102        return len(tld)103    except:104        return -1105 106# Extract features from URL107def main(url):108    status = []109    status.append(having_ip_address(url))110    status.append(abnormal_url(url))111    status.append(count_dot(url))112    status.append(count_www(url))113    status.append(count_atrate(url))114    status.append(no_of_dir(url))115    status.append(no_of_embed(url))116    status.append(shortening_service(url))117    status.append(count_https(url))118    status.append(count_http(url))119    status.append(count_per(url))120    status.append(count_ques(url))121    status.append(count_hyphen(url))122    status.append(count_equal(url))123    status.append(url_length(url))124    status.append(hostname_length(url))125    status.append(suspicious_words(url))126    status.append(digit_count(url))127    status.append(letter_count(url))128    status.append(fd_length(url))129    tld = get_tld(url, fail_silently=True)130    status.append(tld_length(tld))131    return status132 133def get_prediction_from_url(test_url):134    features_test = main(test_url)135    features_test = np.array(features_test).reshape((1, -1))136    pred = model.predict(features_test)137    if int(pred[0]) == 0:138        return "SAFE"139    elif int(pred[0]) == 1:140        return "DEFACEMENT"141    elif int(pred[0]) == 2:142        return "PHISHING"143    elif int(pred[0]) == 3:144        return "MALWARE"145 146# Define prediction endpoint147@app.post("/predict/")148def predict(request: URLRequest):149    prediction = get_prediction_from_url(request.url)150    return {"prediction": prediction}151    152if __name__ == "__main__":153    import uvicorn154    uvicorn.run(app, host="0.0.0.0", port=8000)