model_api.py

"""
model_api.py  -  LinkForensics | Flask REST API
Mirrors predict_cli_1.py logic exactly.

Required folder structure (same as your CLI):
    artifacts/
        models/   RUN7_LGB_LEX_TFIDF_BINARY.joblib
        tfidf/    tfidf_vectorizer.joblib + svd_tfidf.joblib
        lex/      lexical_feature_columns_binary.json

Run:
    python model_api.py
Listens on http://localhost:5001
"""

import os, re, json, traceback, math
import numpy as np
import joblib
from urllib.parse import urlparse
from math import log2
from flask import Flask, request, jsonify
from flask_cors import CORS

app = Flask(__name__)
CORS(app)

# ── Artifact paths (mirrors predict_cli_1.py) ────────────────────────────────
BASE_DIR   = os.path.dirname(os.path.abspath(__file__))
ART        = os.path.join(BASE_DIR, "artifacts")
MODELS_DIR = os.path.join(ART, "models")
TFIDF_DIR  = os.path.join(ART, "tfidf")
LEX_DIR    = os.path.join(ART, "lex")

BIN_LEX_JSON   = os.path.join(LEX_DIR,    "lexical_feature_columns_binary.json")
TFIDF_VEC_PATH = os.path.join(TFIDF_DIR,  "tfidf_vectorizer.joblib")
SVD_PATH       = os.path.join(TFIDF_DIR,  "svd_tfidf.joblib")
BIN_MODEL_PATH = os.path.join(MODELS_DIR, "RUN7_LGB_LEX_TFIDF_BINARY.joblib")
WHITELIST_PATH = os.path.join(LEX_DIR,    "whitelist.txt")

# ── Load all artifacts ───────────────────────────────────────────────────────
try:
    with open(BIN_LEX_JSON, "r", encoding="utf-8") as f:
        BIN_LEX_COLS = json.load(f)
    print(f"[OK]  Lex columns: {len(BIN_LEX_COLS)} features")
except Exception as e:
    BIN_LEX_COLS = []
    print(f"[ERR] Lex columns: {e}")

try:
    tfidf_vec = joblib.load(TFIDF_VEC_PATH)
    print("[OK]  TF-IDF vectorizer loaded")
except Exception as e:
    tfidf_vec = None
    print(f"[ERR] TF-IDF: {e}")

try:
    svd = joblib.load(SVD_PATH)
    print("[OK]  SVD loaded")
except Exception as e:
    svd = None
    print(f"[ERR] SVD: {e}")

try:
    lgb_bin = joblib.load(BIN_MODEL_PATH)
    print(f"[OK]  Model loaded -> {type(lgb_bin).__name__}")
    print(f"[OK]  Model expects {lgb_bin.num_feature()} features")
except Exception as e:
    lgb_bin = None
    print(f"[ERR] Model: {e}")

try:
    with open(WHITELIST_PATH, "r", encoding="utf-8") as f:
        WHITELIST = set(line.strip().lower() for line in f if line.strip())
    print(f"[OK]  Whitelist loaded: {len(WHITELIST)} domains")
except Exception as e:
    WHITELIST = set()
    print(f"[WARN] Whitelist not loaded: {e}")


# ═══════════════════════════════════════════════════════════════════════════
#  FEATURE EXTRACTION  (exact copy from features_lex_1.py)
# ═══════════════════════════════════════════════════════════════════════════

SUSPICIOUS_WORDS = [
    "login","secure","account","update","free","verify","bank","confirm",
    "signin","password","webscr","invoice","billing"
]
SHORTENERS = ["bit.ly","tinyurl.com","goo.gl","t.co","ow.ly","is.gd","buff.ly"]
IP_REGEX   = re.compile(r"(?:\d{1,3}\.){3}\d{1,3}")


def has_ip_address(host):
    return int(bool(IP_REGEX.search(host or "")))


def count_subdirs(path):
    if not path:
        return 0
    return len([p for p in path.split("/") if p])


def url_entropy(raw):
    if not raw:
        return 0.0
    freq = {}
    for ch in raw:
        freq[ch] = freq.get(ch, 0) + 1
    n = len(raw)
    ent = 0.0
    for c in freq.values():
        p = c / n
        ent -= p * log2(p)
    return float(ent)


def abnormal_url_score(raw, host, path):
    raw_l  = (raw  or "").lower()
    host_l = (host or "").lower()
    score  = 0.0
    if not host_l:                                     score += 2.0
    if "@" in raw_l:                                   score += 1.0
    if host_l.count(".") >= 4:                         score += 1.0
    if len(raw_l) >= 75:                               score += 1.0
    if has_ip_address(host_l):                         score += 1.5
    if any(w in raw_l for w in SUSPICIOUS_WORDS):      score += 0.5
    if any(s in host_l for s in SHORTENERS):           score += 1.0
    if raw_l.count("//") > 1:                          score += 0.5
    return float(score)


def canonicalize_url(u):
    raw    = (u or "").strip()
    parsed = urlparse(raw if "://" in raw else "http://" + raw)
    scheme = parsed.scheme.lower()
    host   = (parsed.netloc or "").lower()
    if host.startswith("www."):  host = host[4:]
    if host.endswith(":80"):     host = host[:-3]
    if host.endswith(":443"):    host = host[:-4]
    path  = parsed.path  or ""
    query = parsed.query or ""
    if path != "/" and path.endswith("/"):
        path = path[:-1]
    text = host + path
    if query:
        text += "?" + query
    return {"raw": raw, "scheme": scheme, "host": host,
            "path": path, "query": query, "text": text}


def extract_lex_features(raw_url):
    obj  = canonicalize_url(raw_url)
    raw  = obj["raw"]
    text = obj["text"]
    host = obj["host"]
    path = obj["path"]

    url_len  = len(text)
    host_len = len(host)
    tld_len  = len(host.split(".")[-1]) if "." in host else 0
    digits   = sum(ch.isdigit() for ch in text)
    letters  = sum(ch.isalpha() for ch in text)
    parts    = path.split("/")
    first_dir = parts[1] if len(parts) > 1 else ""
    abn_score = abnormal_url_score(raw, host, path)

    return {
        "url_length":         float(url_len),
        "hostname_length":    float(host_len),
        "count_dot":          float(text.count(".")),
        "count_slash":        float(text.count("/")),
        "count_double_slash": float(text.count("//")),
        "count_www":          float(text.lower().count("www")),
        "count_at":           float(text.count("@")),
        "count_percent":      float(text.count("%")),
        "count_qm":           float(text.count("?")),
        "count_hyphen":       float(text.count("-")),
        "count_equal":        float(text.count("=")),
        "digit_count":        float(digits),
        "letter_count":       float(letters),
        "first_dir_length":   float(len(first_dir)),
        "tld_length":         float(tld_len),
        "digit_ratio":        float(digits / (url_len + 1e-9)),
        "special_ratio":      float(sum((not ch.isalnum()) for ch in text) / (url_len + 1e-9)),
        "has_https":          float(int(raw.lower().startswith("https://"))),
        "has_http":           float(int(raw.lower().startswith("http://"))),
        "suspicious_word":    float(int(any(w in raw.lower() for w in SUSPICIOUS_WORDS))),
        "is_shortened":       float(int(any(s in host for s in SHORTENERS))),
        "abnormal_score":     float(abn_score),
        "abnormal_url":       float(int(abn_score >= 1)),
        "entropy":            float(url_entropy(text)),
        "has_ip":             float(has_ip_address(host)),
        "subdir_count":       float(count_subdirs(path)),
    }


def dict_to_vector(feat_dict, columns):
    return np.array([feat_dict[c] for c in columns], dtype=np.float32).reshape(1, -1)


def make_tfidf_svd(text):
    X = tfidf_vec.transform([text])
    X = svd.transform(X)
    return X.astype(np.float32)


def safe_predict_proba(model, X):
    if hasattr(model, "predict_proba"):
        return model.predict_proba(X)
    pred = np.asarray(model.predict(X))
    if pred.ndim == 1:
        p1 = pred.reshape(-1, 1)
        return np.hstack([1.0 - p1, p1])
    return pred


# ═══════════════════════════════════════════════════════════════════════════
#  PREDICTION  (mirrors predict_one() from predict_cli_1.py)
# ═══════════════════════════════════════════════════════════════════════════

def predict(url, threshold=0.5):
    if lgb_bin is None:   return {"error": "Model not loaded"}
    if tfidf_vec is None: return {"error": "TF-IDF vectorizer not loaded"}
    if svd is None:       return {"error": "SVD not loaded"}
    if not BIN_LEX_COLS:  return {"error": "Lex columns not loaded"}

    # ── Whitelist check (skip model for known-safe domains) ──────────────────
    try:
        from urllib.parse import urlparse as _up
        _host = (_up(url if "://" in url else "http://"+url).hostname or "").lower()
        _host = _host.replace("www.", "")
        if _host in WHITELIST or any(_host.endswith("."+d) for d in WHITELIST):
            return {"url": url, "result": "SAFE", "binary": "benign",
                    "malicious_prob": 0.0, "attack_type": None,
                    "attack_prob": None, "model_type": "whitelist"}
    except Exception:
        pass

    try:
        canon    = canonicalize_url(url)
        tf_text  = canon["text"]

        feat_raw = extract_lex_features(url)
        xlex_bin = dict_to_vector(feat_raw, BIN_LEX_COLS)       # (1, 26)
        xtf      = make_tfidf_svd(tf_text)                       # (1, 300)
        X_bin    = np.hstack([xlex_bin, xtf]).astype(np.float32) # (1, 326)

        proba_bin = safe_predict_proba(lgb_bin, X_bin)            # (1, 2)
        p_mal = float(proba_bin[0, 1]) if proba_bin.shape[1] >= 2 else float(proba_bin[0, 0])

        is_unsafe = p_mal >= threshold
        bin_label = "malicious" if is_unsafe else "benign"

        return {
            "url":             url,
            "result":          "UNSAFE" if is_unsafe else "SAFE",
            "binary":          bin_label,
            "malicious_prob":  round(p_mal, 4),
            "attack_type":     "phishing" if is_unsafe else None,
            "attack_prob":     round(p_mal, 4) if is_unsafe else None,
            "model_type":      type(lgb_bin).__name__,
        }

    except Exception as exc:
        return {"error": str(exc), "trace": traceback.format_exc()}


# ═══════════════════════════════════════════════════════════════════════════
#  ROUTES
# ═══════════════════════════════════════════════════════════════════════════

@app.route("/predict", methods=["POST", "GET"])
def predict_endpoint():
    if request.method == "POST":
        data = request.get_json(force=True, silent=True) or {}
        url  = data.get("url", "").strip()
    else:
        url = request.args.get("url", "").strip()

    if not url:
        return jsonify({"error": "No URL provided"}), 400

    result = predict(url)
    status = 500 if "error" in result else 200
    return jsonify(result), status


@app.route("/inspect", methods=["GET"])
def inspect_endpoint():
    return jsonify({
        "model_loaded":       lgb_bin   is not None,
        "tfidf_loaded":       tfidf_vec is not None,
        "svd_loaded":         svd       is not None,
        "lex_cols_count":     len(BIN_LEX_COLS),
        "model_type":         type(lgb_bin).__name__ if lgb_bin else None,
        "model_num_features": lgb_bin.num_feature()  if lgb_bin else None,
    })


@app.route("/health", methods=["GET"])
def health():
    all_ok = all([lgb_bin, tfidf_vec, svd, BIN_LEX_COLS])
    return jsonify({"status": "ok" if all_ok else "degraded"})


if __name__ == "__main__":
    print("\n" + "="*50)
    print("  LinkForensics  Model API")
    print("  http://localhost:5001")
    print("  POST /predict  { url: '...' }")
    print("  GET  /inspect")
    print("="*50 + "\n")
    app.run(host="0.0.0.0", port=5001, debug=False)