Spaces:
Running on Zero
Running on Zero
| from transformers import pipeline | |
| import logging | |
| import threading | |
| from src.controls import resolve_control | |
| logger = logging.getLogger(__name__) | |
| MODELS = [ | |
| {"name": "facebook/bart-large-mnli", "weight": 0.5}, | |
| {"name": "MoritzLaurer/mDeBERTa-v3-base-xnli-multilingual-nli-2mil7", "weight": 0.3}, | |
| {"name": "cross-encoder/nli-deberta-v3-base", "weight": 0.2}, | |
| ] | |
| CATEGORY_MAP = { | |
| "identity and access management": "IAM", | |
| "network security": "Network", | |
| "container security": "Container", | |
| "secrets management": "Secrets", | |
| "data protection": "Data", | |
| } | |
| CATEGORIES = list(CATEGORY_MAP.keys()) | |
| CATEGORY_RULES = [ | |
| ("IAM", [ | |
| "iam role", "wildcard permission", "action: *", "resource: *", | |
| "mfa", "identity", "rbac", "service account", "cluster-admin", | |
| "access key", "assume role", "trust policy", "permission boundary", | |
| ]), | |
| ("Network", [ | |
| "security group", "0.0.0.0/0", "ingress", "egress", "network policy", | |
| "cidr", "load balancer", "waf", "dns", "vpc", "subnet", | |
| "nacl", "port", "firewall", | |
| ]), | |
| ("Container", [ | |
| "kubernetes", "container", "image", "runs as root", "privileged", | |
| "helm", "docker", "pod", "deployment", "workload", | |
| ]), | |
| ("Secrets", [ | |
| "secret committed", "private key", "hardcoded password", "password is hardcoded", | |
| "hardcoded secret", "credential", "api key", "token exposed", | |
| ]), | |
| ("Data", [ | |
| "storage bucket", "customer records", "encryption", "unencrypted", | |
| "backup", "rds", "database", "pii", "phi", | |
| ]), | |
| ] | |
| SEVERITY_RULES = [ | |
| ("critical", [ | |
| "publicly accessible", "0.0.0.0/0", "action: *", "resource: *", | |
| "admin access", "root user", "secret committed", "private key committed", | |
| "cluster-admin", "hardcoded password", "password is hardcoded", | |
| ]), | |
| ("high", [ | |
| "privileged container", "runs as root", "unencrypted", "no encryption", | |
| "wildcard permission", "hardcoded secret", "api key committed", "no waf", | |
| ]), | |
| ("medium", [ | |
| "no network policy", "missing mfa", "old image", "outdated image", | |
| "no rotation", "security group allows", "latest tag", | |
| ]), | |
| ] | |
| REMEDIATIONS = { | |
| "IAM": "Apply least privilege, remove wildcard permissions, require MFA for privileged identities, and review access through policy analysis.", | |
| "Network": "Restrict ingress and egress to required ports and trusted CIDRs; add segmentation and Kubernetes NetworkPolicies where applicable.", | |
| "Container": "Use a minimal patched image, run as a non-root user, drop Linux capabilities, enable a read-only filesystem, and scan the image in CI.", | |
| "Secrets": "Revoke and rotate the exposed value, remove it from source history, use a managed secret store, and add secret scanning to CI.", | |
| "Data": "Enable encryption at rest and in transit, restrict data access, enable audit logging, and verify backup protection.", | |
| } | |
| OWNERS = { | |
| "IAM": "Platform Security", | |
| "Network": "Cloud Platform", | |
| "Container": "Container Platform", | |
| "Secrets": "Platform Security", | |
| "Data": "Data Platform", | |
| } | |
| def determine_severity(finding): | |
| normalized = finding.lower() | |
| for severity, signals in SEVERITY_RULES: | |
| if any(signal in normalized for signal in signals): | |
| return severity | |
| return "low" | |
| def determine_category(finding): | |
| normalized = finding.lower() | |
| for category, signals in CATEGORY_RULES: | |
| if any(signal in normalized for signal in signals): | |
| return category | |
| return None | |
| def determine_control(finding, category=None): | |
| """Return a stable control identifier while preserving legacy categories.""" | |
| resolved = resolve_control(finding, category or determine_category(finding)) | |
| return resolved["id"] if resolved else "unclassified" | |
| class EnsembleClassifier: | |
| def __init__(self, confidence_threshold=0.55): | |
| self.confidence_threshold = confidence_threshold | |
| self._models = {} | |
| self._model_lock = threading.Lock() | |
| def _get_model(self, model_name): | |
| if model_name not in self._models: | |
| with self._model_lock: | |
| if model_name not in self._models: | |
| try: | |
| self._models[model_name] = pipeline( | |
| "zero-shot-classification", | |
| model=model_name, | |
| ) | |
| except Exception as e: | |
| logger.warning("Failed to load model %s: %s", model_name, e) | |
| return None | |
| return self._models[model_name] | |
| def classify_with_model(self, finding, model_name): | |
| classifier = self._get_model(model_name) | |
| if classifier is None: | |
| return None | |
| try: | |
| result = classifier(finding, CATEGORIES, multi_label=False) | |
| label = result["labels"][0] | |
| score = float(result["scores"][0]) | |
| return { | |
| "category": CATEGORY_MAP.get(label, "IAM"), | |
| "raw_label": label, | |
| "confidence": score, | |
| } | |
| except Exception as e: | |
| logger.warning("Model %s classification failed: %s", model_name, e) | |
| return None | |
| def ensemble_classify(self, finding): | |
| rule_category = determine_category(finding) | |
| if rule_category: | |
| return { | |
| "category": rule_category, | |
| "confidence": 1.0, | |
| "method": "rule_match", | |
| "model_results": [], | |
| } | |
| model_results = [] | |
| category_scores = {} | |
| total_weight = 0 | |
| for model_cfg in MODELS: | |
| result = self.classify_with_model(finding, model_cfg["name"]) | |
| if result: | |
| result["model"] = model_cfg["name"] | |
| result["weight"] = model_cfg["weight"] | |
| model_results.append(result) | |
| cat = result["category"] | |
| weighted_score = result["confidence"] * model_cfg["weight"] | |
| category_scores[cat] = category_scores.get(cat, 0) + weighted_score | |
| total_weight += model_cfg["weight"] | |
| if not model_results: | |
| return { | |
| "category": "IAM", | |
| "confidence": 0.0, | |
| "method": "fallback", | |
| "model_results": [], | |
| } | |
| best_category = max(category_scores, key=category_scores.get) | |
| ensemble_confidence = category_scores[best_category] / total_weight if total_weight else 0 | |
| confidence_gap = 0 | |
| sorted_scores = sorted(category_scores.values(), reverse=True) | |
| if len(sorted_scores) >= 2: | |
| confidence_gap = sorted_scores[0] - sorted_scores[1] | |
| return { | |
| "category": best_category, | |
| "confidence": round(ensemble_confidence, 3), | |
| "confidence_gap": round(confidence_gap, 3), | |
| "method": "ensemble", | |
| "model_results": model_results, | |
| } | |
| def classify(self, finding): | |
| rule_category = determine_category(finding) | |
| severity = determine_severity(finding) | |
| if rule_category: | |
| control = resolve_control(finding, rule_category) | |
| return { | |
| "category": rule_category, | |
| "control_id": control["id"] if control else "unclassified", | |
| "owner": control["owner"] if control else OWNERS[rule_category], | |
| "validation_step": control["validation"] if control else "Review the finding with the owning team.", | |
| "severity": severity, | |
| "confidence": 1.0, | |
| "method": "rule_match", | |
| "model_results": [], | |
| "ensemble_agreement": True, | |
| } | |
| ensemble = self.ensemble_classify(finding) | |
| agreement = len(set(r["category"] for r in ensemble["model_results"])) == 1 | |
| control = resolve_control(finding, ensemble["category"]) | |
| return { | |
| "category": ensemble["category"], | |
| "control_id": control["id"] if control else "unclassified", | |
| "owner": control["owner"] if control else OWNERS.get(ensemble["category"], "Platform Security"), | |
| "validation_step": control["validation"] if control else "Review the finding with the owning team.", | |
| "severity": severity, | |
| "confidence": ensemble["confidence"], | |
| "method": ensemble["method"], | |
| "model_results": ensemble["model_results"], | |
| "ensemble_agreement": agreement, | |
| "confidence_gap": ensemble.get("confidence_gap", 0), | |
| } | |