neighbours.py

"""
Stage 3 - where this window sits in the corpus Part A measured.

Part A's bake-off picked `nomic-ai/nomic-embed-text-v1.5` over MiniLM and BGE, embedding the
`payload_window` column with the `search_document: ` prefix its model card specifies for indexing,
unit-normalised, 768 dimensions. Those four facts are recorded in `part_a_results.json` under
`winner`, and `check_provenance()` asserts them against that file rather than trusting this
docstring - because a query embedded with a different model, or without the prefix, lands in a
different space and every neighbour returned is meaningless, silently, with no error anywhere.

What it is worth: **precision@5 of 35.6%** against a 6.8% random baseline. Fewer than 2 of the 5
files returned are the same kind of attack as the query. That is far better than chance and it is
not good, which is why this is labelled "nearest files in the corpus" and never "the same attack".
"""

import json
import re
import urllib.request

import numpy as np
import pandas as pd

REPO = "Cyber-security-final-project/Evaluation_of_OpenSource_Models_for_PDF_Injection_Recognition"
BASE = f"https://huggingface.co/datasets/{REPO}/resolve/main/"
INDEX_URL = BASE + "Part_A_Outputs/corpus_embeddings.parquet"
RESULTS_URL = BASE + "Part_A_Outputs/part_a_results.json"

# Part A's winner. Every field here is asserted against the results file before use.
EMBED_REPO = "nomic-ai/nomic-embed-text-v1.5"
PREFIX = "search_document: "
DIMS = 768

# The corpus filename encodes its own attack family ("javascript_injection_OWASP_0004.pdf"), and
# the clean controls are hash-named. That is a label on a known corpus file, not a model input -
# nothing here ever shows a filename to a model.
FAMILY_RE = re.compile(r"^([a-z_]+?)_(?:AMTSO|WICAR|OWASP|AtomicRedTeam|Metasploit|Glasswall|"
                       r"mindcrypt|RanSim|RANSIM|custom)_\d+\.pdf$")

_index = None
_model = None


def family_of(file_id: str) -> str:
    """The attack family a corpus file carries, read off its name. Clean files are hash-named."""
    m = FAMILY_RE.match(file_id or "")
    return m.group(1) if m else "clean (control)"


def check_provenance() -> dict:
    """Assert this module embeds queries the way the index was built. Cheap, and load-bearing."""
    with urllib.request.urlopen(RESULTS_URL) as r:
        winner = json.loads(r.read().decode())["winner"]
    assert winner["repo"] == EMBED_REPO, f"index built by {winner['repo']}, app uses {EMBED_REPO}"
    assert winner["prefix"] == PREFIX, "prefix differs from the one the index was built with"
    assert winner["dims"] == DIMS, "dimension mismatch"
    assert winner["normalised"], "index is not unit-normalised; the dot product is not cosine"
    assert winner["input_column"] == "payload_window", "index was not built on payload windows"
    return winner


def load_index():
    """The 1,100 x 768 index Part A exported, straight from the dataset repo. ~3 MB."""
    global _index
    if _index is None:
        df = pd.read_parquet(INDEX_URL)
        _index = {"ids": df["file_id"].to_numpy(),
                  "matrix": df.drop(columns="file_id").to_numpy(dtype=np.float32)}
    return _index


def load_model():
    """Nomic-embed-v1.5 on the CPU. ~550 MB, and a 3,000-character window embeds in a second."""
    global _model
    if _model is None:
        from sentence_transformers import SentenceTransformer
        _model = SentenceTransformer(EMBED_REPO, trust_remote_code=True, device="cpu")
    return _model


def embed(text: str) -> np.ndarray:
    """Embed with the prefix and normalisation Part A used. Anything else is a different space."""
    return load_model().encode([PREFIX + text], normalize_embeddings=True,
                               show_progress_bar=False)[0]


def neighbours(text: str, k: int = 5) -> list:
    """
    The k corpus files nearest one window.

    Vectors are unit-normalised, so the dot product is the cosine similarity and the whole lookup
    against 1,100 files is one matrix-vector product - no index structure needed at this size.
    """
    idx = load_index()
    sims = idx["matrix"] @ embed(text)
    top = np.argsort(-sims)[:k]
    return [{"file_id": str(idx["ids"][i]),
             "family": family_of(str(idx["ids"][i])),
             "similarity": round(float(sims[i]), 3)} for i in top]


NEIGHBOUR_COLUMNS = ["corpus file", "its attack family", "cosine similarity"]


def neighbour_rows(text: str, k: int = 5) -> list:
    return [[n["file_id"], n["family"], n["similarity"]] for n in neighbours(text, k)]