backend/reranker.py

"""
backend/reranker.py
Inventory matching + contextual boost formula.
Mirrors app/reranker.py but lives here for the FastAPI backend.
"""

import pandas as pd
try:
    from thefuzz import fuzz          # type: ignore
except ImportError:
    from fuzzywuzzy import fuzz       # type: ignore


def match_inventory(entities_str, inventory_list: list[str], description: str = "") -> list[str]:
    """
    Return inventory items that appear in the CVE via:
      1. Fuzzy match against NER-extracted entities
      2. Direct substring / fuzzy match against raw description
    """
    matched: list[str] = []

    # ── 1. Entities column ──────────────────────────────────────────────
    if entities_str and not pd.isna(entities_str):
        entities = [e.strip() for e in str(entities_str).split(",") if e.strip()]
        for inv_item in inventory_list:
            for entity in entities:
                if fuzz.partial_ratio(inv_item.lower(), entity.lower()) >= 75:
                    matched.append(inv_item)
                    break

    # ── 2. Raw description ──────────────────────────────────────────────
    if description:
        desc_lower = description.lower()
        for inv_item in inventory_list:
            if inv_item not in matched:
                inv_lower = inv_item.lower()
                if inv_lower in desc_lower:
                    matched.append(inv_item)
                elif fuzz.partial_ratio(inv_lower, desc_lower) >= 85:
                    matched.append(inv_item)

    return list(set(matched))


def compute_context_score(row: dict, inventory_list: list[str], base_prob_critical: float) -> dict:
    """
    Apply boost multipliers and return context_score, matched_inventory, boost_factor.

    boost formula (from README):
        boost  = 1.0
               + 0.30 × inventory_matches
               × 1.25  (if public exploit)
               × 1.15  (if remote + unauthenticated)
               × 1.10  (if attack_vector == NETWORK)

        context_score = min(prob_critical × boost, 1.0)
    """
    boost = 1.0

    matched = match_inventory(
        row.get("entities", ""),
        inventory_list,
        row.get("description", ""),
    )

    if matched:
        boost += 0.3 * len(matched)
    if row.get("exploit_available", 0) == 1:
        boost *= 1.25
    if row.get("has_remote", 0) == 1 and row.get("has_unauth", 0) == 1:
        boost *= 1.15
    if str(row.get("attack_vector", "")).upper() == "NETWORK":
        boost *= 1.10

    context_score = min(float(base_prob_critical) * boost, 1.0)
    return {
        "context_score":     round(context_score, 4),
        "matched_inventory": matched,
        "boost_factor":      round(boost, 3),
    }