Version_3/cli.py

import pickle
import sys
import os
import re
import mailbox
import argparse
import csv
from html import unescape
from bs4 import BeautifulSoup

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# Apply the BASE_DIR to the paths so they resolve absolutely!
MODEL_PATH   = os.path.join(BASE_DIR, "models", "SVM_model.pkl")
FEATURE_PATH = os.path.join(BASE_DIR, "models", "vectorizer.pkl")

# ── text helpers ─────────────────────────────────────────────────────────────
def clean_text(text):
    if not isinstance(text, str):
        return text
    text = re.sub(r'[\x00-\x08\x0B-\x0C\x0E-\x1F\u200B-\u200F\uFEFF]', '', text)
    text = text.encode("utf-16", "surrogatepass").decode("utf-16", "ignore")
    return text[:32767]

def extract_body(msg):
    texts = []
    if msg.is_multipart():
        for part in msg.walk():
            if part.get_content_type() in ("text/plain", "text/html"):
                payload = part.get_payload(decode=True)
                if payload:
                    text = unescape(payload.decode(errors="ignore"))
                    text = BeautifulSoup(text, "html.parser").get_text(" ")
                    texts.append(text)
    else:
        payload = msg.get_payload(decode=True)
        if payload:
            text = unescape(payload.decode(errors="ignore"))
            text = BeautifulSoup(text, "html.parser").get_text(" ")
            texts.append(text)
    combined = " ".join(texts)
    combined = re.sub(r'[\r\n\t]+', ' ', combined)
    combined = re.sub(r'\s+', ' ', combined)
    return combined.strip()

# ── model ─────────────────────────────────────────────────────────────────────
def load_models():
    for path in (MODEL_PATH, FEATURE_PATH):
        if not os.path.exists(path):
            print(f"[ERROR] File not found: {path}")
            sys.exit(1)
    vectorizer = pickle.load(open(FEATURE_PATH, "rb"))
    model      = pickle.load(open(MODEL_PATH,   "rb"))
    return vectorizer, model

def sigmoid(x):
    import math
    return 1 / (1 + math.exp(-x))

def predict(text, vectorizer, model):
    """
    Returns:
      label      : "Spam" or "Ham"
      confidence : probability % if model supports it, else None
      spam_score : 0-100 spam intensity score (always available)
    """
    cleaned  = clean_text(text)
    features = vectorizer.transform([cleaned])
    pred     = model.predict(features)[0]
    label    = "Spam" if str(pred) == "0" else "Ham"

    # confidence via predict_proba (not all SVMs support this)
    confidence = None
    try:
        proba      = model.predict_proba(features)
        # index 0 = spam class (label 0), index 1 = ham class (label 1)
        confidence = round(float(proba[0][0]) * 100, 1)
    except Exception:
        pass

    # spam score via decision function (always works for SVM)
    # decision_function > 0 means ham, < 0 means spam for binary SVC
    # we flip and sigmoid-scale so higher = more spammy
    spam_score = None
    try:
        df_val = float(model.decision_function(features)[0])
        # flip: spam has negative decision value in sklearn SVC (class 0)
        spam_score = round(sigmoid(-df_val) * 100, 1)
    except Exception:
        pass

    # fall back score to confidence if decision_function unavailable
    if spam_score is None and confidence is not None:
        spam_score = confidence

    return label, confidence, spam_score

# ── display helpers ───────────────────────────────────────────────────────────
def score_bar(score, width=20):
    """Visual bar: [████████░░░░░░░░░░░░] 42.3"""
    if score is None:
        return "N/A"
    filled = round(score / 100 * width)
    bar = "█" * filled + "░" * (width - filled)
    return f"[{bar}] {score:.1f}%"

def risk_level(score):
    if score is None:
        return "Unknown"
    if score >= 80:
        return "HIGH RISK"
    if score >= 50:
        return "MEDIUM RISK"
    return "LOW RISK"

def print_result(label, confidence, spam_score):
    verdict = "*** SPAM ***" if label == "Spam" else "Ham (Safe) "
    print(f"\n  Verdict    : {verdict}")
    print(f"  Spam Score : {score_bar(spam_score)}")
    if confidence is not None:
        print(f"  Confidence : {confidence:.1f}%")
    print(f"  Risk Level : {risk_level(spam_score)}")
    print()

# ── modes ─────────────────────────────────────────────────────────────────────
def mode_interactive(vectorizer, model):
    print("\n=== Spam Email Detector ===")
    print("Paste your email text. Press Enter twice to classify.")
    print("Type 'quit' to exit.\n")
    while True:
        print("─" * 44)
        lines = []
        blank_count = 0
        while True:
            try:
                line = input()
            except EOFError:
                break
            if line.strip().lower() in ("quit", "exit"):
                print("Goodbye.")
                sys.exit(0)
            if line.strip() == "":
                blank_count += 1
                if blank_count >= 2:
                    break
            else:
                blank_count = 0
            lines.append(line)

        text = "\n".join(lines).strip()
        if not text:
            print("[!] No text entered. Try again.\n")
            continue

        label, confidence, spam_score = predict(text, vectorizer, model)
        print_result(label, confidence, spam_score)

def mode_single(text, vectorizer, model):
    label, confidence, spam_score = predict(text, vectorizer, model)
    print_result(label, confidence, spam_score)

def mode_mbox(mbox_path, vectorizer, model, output_csv):
    if not os.path.exists(mbox_path):
        print(f"[ERROR] File not found: {mbox_path}")
        sys.exit(1)
    print(f"Loading: {mbox_path}")
    mbox     = mailbox.mbox(mbox_path)
    messages = list(mbox)
    print(f"Found {len(messages)} emails. Classifying...\n")

    results    = []
    spam_count = 0
    scores     = []

    for i, msg in enumerate(messages, 1):
        subject = clean_text(msg.get("Subject", "(no subject)"))
        body    = extract_body(msg)
        label, confidence, spam_score = predict(body, vectorizer, model)

        if label == "Spam":
            spam_count += 1
        if spam_score is not None:
            scores.append(spam_score)

        conf_str  = f"{confidence:.1f}%" if confidence is not None else "N/A"
        score_str = f"{spam_score:.1f}%"  if spam_score  is not None else "N/A"
        risk      = risk_level(spam_score)
        marker    = "SPAM" if label == "Spam" else "Ham "

        results.append({
            "Index":      i,
            "Subject":    subject,
            "Prediction": label,
            "SpamScore":  score_str,
            "Confidence": conf_str,
            "RiskLevel":  risk,
        })
        print(f"  [{i:>4}] {marker}  Score:{score_str:>6}  {risk:<11}  {subject[:50]}")

    ham_count = len(results) - spam_count
    avg_score = round(sum(scores) / len(scores), 1) if scores else 0

    print(f"\n{'─'*44}")
    print(f"  Total emails : {len(results)}")
    print(f"  Spam         : {spam_count}")
    print(f"  Ham          : {ham_count}")
    print(f"  Avg Spam Score : {avg_score}%")

    if output_csv:
        with open(output_csv, "w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=["Index","Subject","Prediction","SpamScore","Confidence","RiskLevel"])
            writer.writeheader()
            writer.writerows(results)
        print(f"\n  Saved to: {output_csv}")

# ── entry point ───────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(
        description="Spam Email Detector — CMD",
        formatter_class=argparse.RawTextHelpFormatter,
        epilog="""
Examples:
  python cli.py                                          # interactive
  python cli.py --text "You won a prize! Click now."    # inline text
  python cli.py --file email.txt                        # from file
  python cli.py --mbox inbox.mbox --output results.csv  # batch mbox
"""
    )
    parser.add_argument("--text",   help="Email text to classify (inline)")
    parser.add_argument("--file",   help="Path to a .txt file with email content")
    parser.add_argument("--mbox",   help="Path to an .mbox file for batch classification")
    parser.add_argument("--output", help="CSV file to save mbox results (optional)")
    args = parser.parse_args()

    print("Loading models...", end=" ", flush=True)
    vectorizer, model = load_models()
    print("OK\n")

    if args.mbox:
        mode_mbox(args.mbox, vectorizer, model, args.output)
    elif args.text:
        mode_single(args.text, vectorizer, model)
    elif args.file:
        if not os.path.exists(args.file):
            print(f"[ERROR] File not found: {args.file}")
            sys.exit(1)
        with open(args.file, "r", encoding="utf-8", errors="ignore") as f:
            text = f.read()
        mode_single(text, vectorizer, model)
    else:
        mode_interactive(vectorizer, model)

if __name__ == "__main__":
    main()