Upload 4 files

Dockerfile

@@ -0,0 +1,34 @@
+# Use an official lightweight Python image
+FROM python:3.10-slim
+
+# Create a non-root user that Hugging Face requires (User ID 1000)
+RUN useradd -m -u 1000 user
+USER user
+
+# Set up the environment paths for the new user
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+# Set the working directory
+WORKDIR $HOME/app
+
+# Copy your requirements and install them
+COPY --chown=user requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy all your application code and models into the container
+COPY --chown=user . .
+
+# Switch to root temporarily to create the /tmp/ folders and grant open permissions
+USER root
+RUN mkdir -p /tmp/tag_and_trail_downloads /tmp/tag_and_trail_uploads && \
+    chmod -R 777 /tmp/tag_and_trail_downloads /tmp/tag_and_trail_uploads
+
+# Switch back to the safe user
+USER user
+
+# Expose Hugging Face's mandatory port
+EXPOSE 7860
+
+# Command to boot up your ML Brain
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,67 @@
+import os
+import shutil
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from pydantic import BaseModel
+
+# Import the master pipelines we just built
+from ml_pipeline import process_url_pipeline, process_media_pipeline
+
+app = FastAPI(title="Tag & Trail ML Brain")
+
+# Pydantic model for incoming URL requests
+class URLRequest(BaseModel):
+    text: str
+
+@app.post("/predict_url")
+async def predict_url(req: URLRequest):
+    """Receives a URL, runs ML, and extracts metadata if safe."""
+    if not req.text:
+        raise HTTPException(status_code=400, detail="No text/URL provided.")
+    
+    try:
+        result = process_url_pipeline(req.text)
+        
+        # Ensure we return a "prediction" key so your helpers.py can read it perfectly
+        if "class" in result:
+            result["prediction"] = result["class"].lower()
+            
+        return result
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"URL Pipeline Error: {str(e)}")
+
+@app.post("/predict_pdf")
+async def predict_pdf(file: UploadFile = File(...)):
+    """Receives a PDF/Media file, converts, runs ML, and extracts metadata if safe."""
+    
+    # Hugging Face Spaces allows writing to /tmp/ for ephemeral storage
+    temp_dir = "/tmp/tag_and_trail_downloads"
+    os.makedirs(temp_dir, exist_ok=True)
+    
+    temp_path = os.path.join(temp_dir, file.filename)
+    
+    try:
+        # Save the incoming file from Twilio/helpers.py
+        with open(temp_path, "wb") as buffer:
+            shutil.copyfileobj(file.file, buffer)
+            
+        # Run your master media pipeline
+        result = process_media_pipeline(
+            file_path=temp_path, 
+            mime=file.content_type, 
+            original_url_or_name=file.filename
+        )
+        
+        return result
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Media Pipeline Error: {str(e)}")
+        
+    finally:
+        # ALWAYS clean up to prevent memory/storage leaks in Hugging Face
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
+
+# Root endpoint just for quick health checks
+@app.get("/")
+def health_check():
+    return {"status": "Tag & Trail ML API is running smoothly!"}

ml_pipeline.py

@@ -0,0 +1,447 @@
+import os
+import io
+import re
+import json
+import math
+import uuid
+import fitz
+import docx
+import joblib
+import requests
+import numpy as np
+import pandas as pd
+
+from pptx import Presentation
+from openpyxl import load_workbook
+from PIL import Image
+from bs4 import BeautifulSoup
+from PyPDF2 import PdfReader
+
+# TensorFlow Imports
+from tensorflow.keras.models import load_model
+from tensorflow.keras.preprocessing.text import tokenizer_from_json
+from tensorflow.keras.preprocessing.sequence import pad_sequences
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
+os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
+
+BASE_DIR = os.path.dirname(__file__)
+
+# Hugging Face usually provides a writeable /tmp/ directory for ephemeral files
+DOWNLOAD_FOLDER = "/tmp/tag_and_trail_downloads"
+os.makedirs(DOWNLOAD_FOLDER, exist_ok=True)
+
+# ==============================================================================
+# 1. URL ML MODEL (LSTM + CNN)
+# ==============================================================================
+
+tokenizer_path = os.path.join(BASE_DIR, "models", "tokenizer.json")
+with open(tokenizer_path, 'r') as f:
+    tokenizer_data = json.load(f)
+    tokenizer = tokenizer_from_json(tokenizer_data)
+
+lstm_model = load_model(os.path.join(BASE_DIR, 'models', 'tag_and_trail_lstm_best.keras'))
+cnn_model = load_model(os.path.join(BASE_DIR, 'models', 'tag_trail_url_cnn_model.keras'))
+
+def tag_and_trail_inference(raw_url):
+    url = raw_url.lower().strip()
+    sequences = tokenizer.texts_to_sequences([url])
+    padded_data = pad_sequences(sequences, maxlen=200, padding='post')
+    
+    lstm_preds = lstm_model.predict(padded_data, verbose=0)
+    cnn_preds = cnn_model.predict(padded_data, verbose=0)
+    
+    final_probabilities = (lstm_preds + cnn_preds) / 2
+    
+    classes = ['Safe', 'Defacement', 'Malware', 'Phishing']
+    predicted_idx = np.argmax(final_probabilities)
+    confidence = float(np.max(final_probabilities))
+    
+    return {
+        "url": raw_url,
+        "class": classes[predicted_idx],
+        "confidence": confidence,
+        "raw_scores": final_probabilities.tolist()
+    }
+
+
+# ==============================================================================
+# 2. PDF ML MODEL (Feature Extraction + Joblib)
+# ==============================================================================
+
+FEATURE_NAMES = ["pdfsize", "pages", "isEncrypted", "JS", "OpenAction", "launch", "AA", "EmbeddedFile", "ObjStm", "entropy"]
+SUSPICIOUS_KEYS = [b"/JS", b"/JavaScript", b"/OpenAction", b"/Launch", b"/AA", b"/EmbeddedFile", b"/ObjStm"]
+FLAG_MAPPING = {3: "JavaScript", 4: "OpenAction", 5: "Launch", 6: "AdditionalActions", 7: "EmbeddedFile", 8: "ObjectStream"}
+
+MODEL_PATH = os.path.join(BASE_DIR, "models", "model.joblib")
+try:
+    bundle = joblib.load(MODEL_PATH)
+    pdf_model = bundle["model"]
+    FEATURE_COLUMNS = bundle["features"]
+    scaler = bundle.get("scaler", None)
+    print("PDF model loaded successfully")
+except Exception as e:
+    print("PDF model load error:", e)
+    pdf_model = None
+    scaler = None
+    FEATURE_COLUMNS = FEATURE_NAMES
+
+def shannon_entropy(data: bytes):
+    if not data: return 0.0
+    freq = {}
+    for b in data: freq[b] = freq.get(b, 0) + 1
+    entropy = 0.0
+    length = len(data)
+    for count in freq.values():
+        p = count / length
+        entropy -= p * math.log2(p)
+    return entropy
+
+def extract_features(pdf_path):
+    try:
+        pdfsize = os.path.getsize(pdf_path) / 1024.0
+        with open(pdf_path, "rb") as f:
+            raw = f.read(1024 * 1024)
+        entropy = shannon_entropy(raw)
+        reader = PdfReader(pdf_path, strict=False)
+        pages = len(reader.pages)
+        is_encrypted = 1 if reader.is_encrypted else 0
+        counts = {k.decode('latin-1'): raw.count(k) for k in SUSPICIOUS_KEYS}
+        features = [
+            pdfsize, float(pages), float(is_encrypted),
+            float(counts.get('/JS', 0) + counts.get('/JavaScript', 0)),
+            float(1 if counts.get('/OpenAction', 0) > 0 else 0),
+            float(1 if counts.get('/Launch', 0) > 0 else 0),
+            float(1 if counts.get('/AA', 0) > 0 else 0),
+            float(1 if counts.get('/EmbeddedFile', 0) > 0 else 0),
+            float(1 if counts.get('/ObjStm', 0) > 0 else 0),
+            entropy
+        ]
+        features[6] = min(features[6], 0.3)  # AA
+        features[8] = min(features[8], 0.3)  # ObjStm
+        return features
+    except Exception:
+        return [0.0] * len(FEATURE_NAMES)
+
+def check_pdf_encryption(pdf_path):
+    try:
+        reader = PdfReader(pdf_path, strict=False)
+        return reader.is_encrypted
+    except:
+        return False
+
+def extract_flags(features):
+    flags = []
+    for index, name in FLAG_MAPPING.items():
+        if features[index] > 0: flags.append(name)
+    return flags
+
+def predict_pdf(pdf_path):
+    if not os.path.exists(pdf_path):
+        return {"prediction": "UNKNOWN", "malicious_probability": 0.0, "risk_level": "Error", "flags": []}
+    
+    is_encrypted = check_pdf_encryption(pdf_path)
+    features = extract_features(pdf_path)
+    flags = extract_flags(features)
+    
+    if pdf_model is None:
+        return {"prediction": "UNKNOWN", "malicious_probability": 0.0, "risk_level": "Error", "flags": flags}
+    try:
+        X = pd.DataFrame([features], columns=FEATURE_COLUMNS)
+        if scaler: X_scaled = scaler.transform(X)
+        else: X_scaled = X.values
+        
+        prob_malicious = float(pdf_model.predict_proba(X_scaled)[0][1])
+
+        js = features[3]
+        openaction = features[4]
+        launch = features[5]
+        aa = features[6]
+        embedded = features[7]
+
+        if (js == 0 and openaction == 0 and launch == 0 and embedded == 0 and not is_encrypted):
+            prob_malicious = 0.1
+        strong_flags = js + openaction + launch + embedded
+        if strong_flags <=1 and prob_malicious > 0.5:  
+            prob_malicious = 0.3
+
+        if prob_malicious >= 0.85: risk_level = "HIGH"
+        elif prob_malicious >= 0.6: risk_level = "MEDIUM"
+        elif prob_malicious >= 0.3: risk_level = "LOW"
+        else: risk_level = "SAFE"
+
+        return {
+            "prediction": "Malicious" if prob_malicious > 0.5 else "Safe",
+            "malicious_probability": round(prob_malicious, 4),
+            "risk_level": risk_level,
+            "flags": flags,
+            "is_encrypted": is_encrypted
+        }
+    except Exception as e:
+        return {"prediction": "UNKNOWN", "malicious_probability": 0.0, "risk_level": "Error", "flags": flags, "error": str(e)}
+
+
+# ==============================================================================
+# 3. CONVERTERS
+# ==============================================================================
+
+def convert_image_to_pdf(image_path: str) -> str:
+    img = Image.open(image_path)
+    if img.mode != "RGB": img = img.convert("RGB")
+    img_bytes = io.BytesIO()
+    img.save(img_bytes, format="PDF")
+    pdf_path = os.path.join(DOWNLOAD_FOLDER, f"{uuid.uuid4()}.pdf")
+    with open(pdf_path, "wb") as f: f.write(img_bytes.getvalue())
+    return pdf_path
+
+def convert_text_to_pdf(text_path: str) -> str:
+    with open(text_path, "r", encoding="utf-8", errors="ignore") as f: text = f.read()
+    doc = fitz.open()
+    page = doc.new_page()
+    page.insert_text((50, 50), text, fontsize=11)
+    pdf_path = os.path.join(DOWNLOAD_FOLDER, f"{uuid.uuid4()}.pdf")
+    doc.save(pdf_path)
+    doc.close()
+    return pdf_path
+
+def convert_docx_to_pdf(docx_path: str) -> str:
+    doc = docx.Document(docx_path)
+    full_text = [para.text for para in doc.paragraphs]
+    text = "\n".join(full_text)
+    tmp_path = docx_path + ".tmp.txt"
+    with open(tmp_path, "w", encoding="utf-8") as f: f.write(text)
+    pdf_path = convert_text_to_pdf(tmp_path)
+    os.remove(tmp_path)
+    return pdf_path
+
+def convert_pptx_to_pdf(pptx_path: str) -> str:
+    prs = Presentation(pptx_path)
+    text_runs = []
+    for slide in prs.slides:
+        for shape in slide.shapes:
+            if hasattr(shape, "text"): text_runs.append(shape.text)
+    text = "\n".join(text_runs)
+    tmp_path = pptx_path + ".tmp.txt"
+    with open(tmp_path, "w", encoding="utf-8") as f: f.write(text)
+    pdf_path = convert_text_to_pdf(tmp_path)
+    os.remove(tmp_path)
+    return pdf_path
+
+def convert_xlsx_to_pdf(xlsx_path: str) -> str:
+    wb = load_workbook(xlsx_path, read_only=True, data_only=True)
+    text_lines = []
+    for sheet in wb.worksheets:
+        for row in sheet.iter_rows(values_only=True):
+            row_str = " | ".join(str(cell) if cell is not None else "" for cell in row)
+            if row_str.strip(): text_lines.append(row_str)
+    text = "\n".join(text_lines)
+    tmp_path = xlsx_path + ".tmp.txt"
+    with open(tmp_path, "w", encoding="utf-8") as f: f.write(text)
+    pdf_path = convert_text_to_pdf(tmp_path)
+    os.remove(tmp_path)
+    return pdf_path
+
+def convert_to_pdf(file_path: str, mime: str) -> str:
+    if mime.startswith("image/"): return convert_image_to_pdf(file_path)
+    elif mime == "text/plain": return convert_text_to_pdf(file_path)
+    elif mime == "application/vnd.openxmlformats-officedocument.wordprocessingml.document": return convert_docx_to_pdf(file_path)
+    elif mime == "application/msword": raise ValueError("Unsupported mime type: .doc (old format) – please convert to .docx")
+    elif mime == "application/vnd.openxmlformats-officedocument.presentationml.presentation": return convert_pptx_to_pdf(file_path)
+    elif mime == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": return convert_xlsx_to_pdf(file_path)
+    else: return file_path  # If already PDF or unknown, pass through
+
+
+# ==============================================================================
+# 4. METADATA EXTRACTION
+# ==============================================================================
+
+def clean_text(text: str, keep_punctuation: bool = True) -> str:
+    if not text: return ""
+    text = text.lower()
+    text = re.sub(r"\[\d+\]", "", text)
+    if keep_punctuation: text = re.sub(r"[^a-z0-9\s.,!?;:'\"()-]", " ", text)
+    else: text = re.sub(r"[^a-z0-9\s]", " ", text)
+    text = re.sub(r"\s+", " ", text)
+    return text.strip()
+
+def title_from_url(url: str) -> str:
+    from urllib.parse import urlparse
+    parsed = urlparse(url)
+    filename = os.path.splitext(os.path.basename(parsed.path))[0]
+    if not filename: return ""
+    parts = re.split(r"[-_/]", filename)
+    words = [p for p in parts if p and not p.isdigit()]
+    return clean_text(" ".join(words), keep_punctuation=False)
+
+def is_bad_title(title: str) -> bool:
+    if not title: return True
+    title_lower = title.lower()
+    bad_phrases = ["error", "not found", "404", "403", "access denied", "forbidden", "captcha", "challenge", "just a moment", "enable javascript", "redirecting", "loading"]
+    return any(phrase in title_lower for phrase in bad_phrases)
+
+def extract_first_paragraph_html(soup):
+    for tag in soup(["script", "style", "nav", "footer", "header", "aside"]): tag.decompose()
+    content_parents = soup.find_all(['article', 'main'])
+    if not content_parents:
+        content_parents = soup.find_all('div', class_=re.compile(r"(content|post|entry|article)", re.I))
+    if not content_parents: content_parents = [soup]
+    best_paragraph = None
+    max_len = 0
+    for parent in content_parents:
+        for p in parent.find_all('p', recursive=True):
+            text = clean_text(p.get_text(" ", strip=True))
+            length = len(text)
+            if length <= 80: continue
+            if text.startswith("by ") and length < 120: continue
+            if length > 300: return text
+            if length > max_len:
+                max_len = length
+                best_paragraph = text
+    if best_paragraph: return best_paragraph
+    for p in soup.find_all('p'):
+        text = clean_text(p.get_text(" ", strip=True))
+        if len(text) > 80: return text
+    return None
+
+def extract_first_paragraph_pdf(pdf_bytes, max_pages=5):
+    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
+    try:
+        if doc.page_count == 0: return None
+        pages_to_scan = min(doc.page_count, max_pages)
+        for page_num in range(pages_to_scan):
+            page = doc.load_page(page_num)
+            text = page.get_text()
+            for block in text.split("\n\n"):
+                block = block.strip()
+                if not block: continue
+                cleaned = clean_text(block)
+                length = len(cleaned)
+                if length <= 80: continue
+                digit_ratio = sum(c.isdigit() for c in cleaned) / length
+                if digit_ratio > 0.3: continue
+                return cleaned
+            lines = [line.strip() for line in text.split("\n") if line.strip()]
+            combined = ""
+            for line in lines:
+                combined += " " + line
+                if len(combined) > 200:
+                    cleaned = clean_text(combined)
+                    length = len(cleaned)
+                    if length <= 80: continue
+                    digit_ratio = sum(c.isdigit() for c in cleaned) / length
+                    if digit_ratio > 0.3: continue
+                    return cleaned
+        return None
+    finally:
+        doc.close()
+
+def extract_metadata_from_url(url: str) -> dict:
+    metadata = {"title": None, "description": None}
+    try:
+        response = requests.get(url, timeout=(5, 10), headers={"User-Agent": "Mozilla/5.0"}, allow_redirects=True)
+        response.raise_for_status()
+        response.encoding = response.apparent_encoding
+        content_type = response.headers.get("Content-Type", "").lower()
+        if not content_type.startswith("text/html"):
+            metadata["title"] = title_from_url(url)
+            return metadata
+        soup = BeautifulSoup(response.text, "html.parser")
+        title_candidates = []
+        og_title = soup.find("meta", property="og:title")
+        if og_title and og_title.get("content"): title_candidates.append(clean_text(og_title["content"]))
+        if soup.title and soup.title.string: title_candidates.append(clean_text(soup.title.string))
+        h1 = soup.find("h1")
+        if h1: title_candidates.append(clean_text(h1.get_text()))
+        title_candidates.append(title_from_url(url))
+        for candidate in title_candidates:
+            if candidate and not is_bad_title(candidate):
+                if " | " in candidate: candidate = candidate.split(" | ")[0].strip()
+                if " - " in candidate: candidate = candidate.split(" - ")[0].strip()
+                metadata["title"] = candidate
+                break
+        desc = None
+        meta_desc = soup.find("meta", attrs={"name": re.compile(r"^description$", re.I)})
+        if meta_desc and meta_desc.get("content"): desc = clean_text(meta_desc["content"])
+        if not desc:
+            og_desc = soup.find("meta", property="og:description")
+            if og_desc and og_desc.get("content"): desc = clean_text(og_desc["content"])
+        if not desc:
+            twitter_desc = soup.find("meta", attrs={"name": re.compile(r"^twitter:description$", re.I)})
+            if twitter_desc and twitter_desc.get("content"): desc = clean_text(twitter_desc["content"])
+        first_para = extract_first_paragraph_html(soup)
+        if desc and first_para: metadata["description"] = f"{desc} | {first_para}"
+        elif desc: metadata["description"] = desc
+        elif first_para: metadata["description"] = first_para
+        return metadata
+    except requests.RequestException:
+        metadata["title"] = title_from_url(url)
+        return metadata
+
+def extract_metadata_from_pdf(file_path: str, url: str) -> dict:
+    metadata = {"title": None, "description": None}
+    try:
+        doc = fitz.open(file_path)
+        try:
+            pdf_title = doc.metadata.get("title")
+            cleaned_title = clean_text(pdf_title) if pdf_title else None
+            if cleaned_title and not is_bad_title(cleaned_title): metadata["title"] = cleaned_title
+            else: metadata["title"] = title_from_url(url)
+        finally:
+            doc.close()
+        with open(file_path, "rb") as f: pdf_bytes = f.read()
+        first_para = extract_first_paragraph_pdf(pdf_bytes)
+        if first_para: metadata["description"] = first_para
+        return metadata
+    except Exception:
+        metadata["title"] = title_from_url(url)
+        return metadata
+
+
+# ==============================================================================
+# 5. THE MASTER PIPELINES (Safety First Rule Enforced)
+# ==============================================================================
+
+def process_url_pipeline(url: str) -> dict:
+    """Predicts threat first. If safe, extracts metadata for tagging."""
+    result = tag_and_trail_inference(url)
+    
+    # If the model finds it malicious, abort metadata extraction to save compute/safety
+    if result["class"] in ["Malware", "Phishing"]:
+        result["metadata_skipped"] = "Threat detected. Extraction aborted."
+        return result
+    
+    # If safe, extract the data
+    metadata = extract_metadata_from_url(url)
+    result.update(metadata)
+    
+    # (Here is where you will eventually call your HF Tagging API with the metadata)
+    
+    return result
+
+def process_media_pipeline(file_path: str, mime: str, original_url_or_name: str) -> dict:
+    """Converts to PDF (if needed), predicts threat first. If safe, extracts metadata."""
+    
+    # Step 1: Conversion (The PDF model requires a PDF)
+    if mime != "application/pdf":
+        try:
+            pdf_path = convert_to_pdf(file_path, mime)
+        except Exception as e:
+            return {"error": f"Conversion failed: {str(e)}"}
+    else:
+        pdf_path = file_path
+
+    # Step 2: Machine Learning Prediction
+    result = predict_pdf(pdf_path)
+    
+    # Step 3: Safety Check - Only extract if it's Safe
+    if result["prediction"] == "Safe" and result["risk_level"] in ["SAFE", "LOW"]:
+        metadata = extract_metadata_from_pdf(pdf_path, original_url_or_name)
+        result.update(metadata)
+    else:
+        result["metadata_skipped"] = "Threat detected. Extraction aborted."
+
+    # Clean up the converted PDF if we created a new one
+    if pdf_path != file_path and os.path.exists(pdf_path):
+        os.remove(pdf_path)
+
+    return result