Building a Custom Risk Scorer

from dataclasses import dataclass

@dataclass
class ScoringRule:
    """A single rule that contributes to the risk score."""
    name: str
    weight: float  # 0.0 - 1.0
    description: str

    def evaluate(self, ctx) -> float:
        """Return a raw score in [0.0, 1.0] for this rule."""
        raise NotImplementedError

from attesta import ActionContext

class EnvironmentRule(ScoringRule):
    """Production actions are inherently riskier."""
    def evaluate(self, ctx: ActionContext) -> float:
        if ctx.environment == "production":
            return 0.8
        if ctx.environment == "staging":
            return 0.4
        return 0.1

class DataClassificationRule(ScoringRule):
    """Score based on data sensitivity hints."""
    def evaluate(self, ctx: ActionContext) -> float:
        classification = ctx.hints.get("data_classification", "public")
        scores = {
            "public": 0.1,
            "internal": 0.3,
            "confidential": 0.6,
            "restricted": 0.9,
        }
        return scores.get(classification, 0.5)

class BlastRadiusRule(ScoringRule):
    """Score based on how many systems are affected."""
    def evaluate(self, ctx: ActionContext) -> float:
        affected = ctx.hints.get("affected_systems", 1)
        if affected >= 10:
            return 0.9
        if affected >= 5:
            return 0.6
        if affected >= 2:
            return 0.3
        return 0.1

class PolicyRiskScorer:
    """Rule-based scorer driven by organizational security policies."""

    def __init__(self):
        self.rules = [
            EnvironmentRule(
                name="environment",
                weight=0.3,
                description="Production environment risk",
            ),
            DataClassificationRule(
                name="data_classification",
                weight=0.4,
                description="Data sensitivity classification",
            ),
            BlastRadiusRule(
                name="blast_radius",
                weight=0.3,
                description="Number of affected systems",
            ),
        ]

    @property
    def name(self) -> str:
        return "policy"

    def score(self, ctx) -> float:
        total_weight = sum(r.weight for r in self.rules)
        weighted_sum = sum(
            r.evaluate(ctx) * r.weight for r in self.rules
        )
        return max(0.0, min(1.0, weighted_sum / total_weight))

import hashlib
from attesta import ActionContext

def extract_features(ctx: ActionContext) -> dict:
    """Extract ML features from an action context."""
    fn_tokens = ctx.function_name.lower().split("_")
    arg_str = " ".join(str(a) for a in ctx.args)

    return {
        "has_destructive_verb": any(
            t in {"delete", "drop", "destroy", "purge"} for t in fn_tokens
        ),
        "has_mutating_verb": any(
            t in {"write", "update", "create", "send", "deploy"} for t in fn_tokens
        ),
        "is_production": ctx.environment == "production",
        "arg_length": len(arg_str),
        "has_pii_hint": bool(ctx.hints.get("pii")),
        "has_financial_hint": bool(ctx.hints.get("financial")),
        "function_name_hash": int(
            hashlib.md5(ctx.function_name.encode()).hexdigest()[:8], 16
        ),
    }

import pickle
from pathlib import Path
from attesta import ActionContext

class MLRiskScorer:
    """Risk scorer backed by a trained scikit-learn model."""

    def __init__(self, model_path: str | Path):
        with open(model_path, "rb") as f:
            self._model = pickle.load(f)
        self._feature_names = [
            "has_destructive_verb",
            "has_mutating_verb",
            "is_production",
            "arg_length",
            "has_pii_hint",
            "has_financial_hint",
            "function_name_hash",
        ]

    @property
    def name(self) -> str:
        return "ml-classifier"

    def score(self, ctx: ActionContext) -> float:
        features = extract_features(ctx)
        feature_vector = [
            [features[name] for name in self._feature_names]
        ]

        # predict_proba returns [[p_low, p_high]]
        proba = self._model.predict_proba(feature_vector)[0]
        risk_score = proba[1]  # probability of "high risk" class

        return max(0.0, min(1.0, risk_score))

from attesta.core.risk import DefaultRiskScorer

class RobustMLScorer:
    """ML scorer with automatic fallback to the default heuristic."""

    def __init__(self, model_path: str):
        self._ml = MLRiskScorer(model_path)
        self._fallback = DefaultRiskScorer()

    @property
    def name(self) -> str:
        return "ml-with-fallback"

    def score(self, ctx) -> float:
        try:
            return self._ml.score(ctx)
        except Exception:
            # Log the error in production
            return self._fallback.score(ctx)

from attesta.core.risk import CompositeRiskScorer, DefaultRiskScorer

scorer = CompositeRiskScorer(
    scorers=[
        (RobustMLScorer("models/risk_v2.pkl"), 0.6),
        (DefaultRiskScorer(), 0.4),
    ]
)