""" CyberForge AI - ML Training & Inference Platform Hugging Face Spaces deployment with: 1) Gradio UI for notebook execution, training, and inference 2) FastAPI REST endpoints for the Heroku backend (mlService.js) """ import gradio as gr import pandas as pd import numpy as np import json import os import subprocess import sys from pathlib import Path from datetime import datetime import logging from typing import Dict, List, Any, Optional, Tuple from urllib.parse import urlparse # ML Libraries from sklearn.model_selection import train_test_split, cross_val_score from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier, IsolationForest from sklearn.linear_model import LogisticRegression from sklearn.preprocessing import StandardScaler, LabelEncoder from sklearn.metrics import classification_report, confusion_matrix, accuracy_score, f1_score import joblib # Hugging Face Hub from huggingface_hub import HfApi, hf_hub_download, upload_file # FastAPI for REST endpoints from fastapi import FastAPI, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse # Gemini AI (new SDK: google-genai) try: from google import genai GEMINI_AVAILABLE = True except ImportError: GEMINI_AVAILABLE = False logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # ============================================================================ # CONFIGURATION # ============================================================================ APP_DIR = Path(__file__).parent.absolute() MODELS_DIR = APP_DIR / "trained_models" MODELS_DIR.mkdir(exist_ok=True) DATASETS_DIR = APP_DIR / "datasets" DATASETS_DIR.mkdir(exist_ok=True) NOTEBOOKS_DIR = APP_DIR / "notebooks" KNOWLEDGE_BASE_DIR = APP_DIR / "knowledge_base" KNOWLEDGE_BASE_DIR.mkdir(exist_ok=True) TRAINING_DATA_DIR = APP_DIR / "training_data" TRAINING_DATA_DIR.mkdir(exist_ok=True) # Environment GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "") GEMINI_MODEL = os.environ.get("GEMINI_MODEL", "gemini-2.5-flash") HF_TOKEN = os.environ.get("HF_TOKEN", "") HF_MODEL_REPO = os.environ.get("HF_MODEL_REPO", "Che237/cyberforge-models") logger.info(f"APP_DIR: {APP_DIR}") logger.info(f"NOTEBOOKS_DIR: {NOTEBOOKS_DIR}") logger.info(f"NOTEBOOKS_DIR exists: {NOTEBOOKS_DIR.exists()}") # Model types available for training MODEL_TYPES = { "Random Forest": RandomForestClassifier, "Gradient Boosting": GradientBoostingClassifier, "Logistic Regression": LogisticRegression, "Isolation Forest (Anomaly)": IsolationForest, } SECURITY_TASKS = [ "Malware Detection", "Phishing Detection", "Network Intrusion Detection", "Anomaly Detection", "Botnet Detection", "Web Attack Detection", "Spam Detection", "Vulnerability Assessment", "DNS Tunneling Detection", "Cryptomining Detection", ] # ============================================================================ # GEMINI SERVICE (for REST API) # ============================================================================ class GeminiService: """Google Gemini AI for cybersecurity chat and analysis""" SYSTEM_PROMPT = """You are CyberForge AI, an advanced cybersecurity expert. You specialize in: - Real-time threat detection and analysis - Malware and phishing identification - Network security assessment - Browser security monitoring - Risk assessment and mitigation When analyzing security queries, provide: 1. Risk Level (Critical/High/Medium/Low) 2. Threat Types identified 3. Confidence Score (0.0-1.0) 4. Detailed technical analysis 5. Specific actionable recommendations Always be precise, professional, and actionable.""" def __init__(self): self.client = None self.ready = False self.custom_knowledge = {} self.training_examples = [] def initialize(self): if not GEMINI_AVAILABLE: logger.warning("google-genai not installed, Gemini unavailable") return if not GEMINI_API_KEY: logger.warning("GEMINI_API_KEY not set, Gemini unavailable") return try: self.client = genai.Client(api_key=GEMINI_API_KEY) resp = self.client.models.generate_content( model=GEMINI_MODEL, contents="Test. Respond with OK." ) if resp.text: self.ready = True logger.info(f"✅ Gemini initialized (model: {GEMINI_MODEL})") self._load_knowledge() except Exception as e: logger.error(f"❌ Gemini init failed: {e}") self.ready = False def _load_knowledge(self): try: for f in KNOWLEDGE_BASE_DIR.glob("*.json"): with open(f) as fh: self.custom_knowledge[f.stem] = json.load(fh) logger.info(f"Loaded {len(self.custom_knowledge)} knowledge files") except Exception as e: logger.warning(f"Knowledge load error: {e}") try: for f in TRAINING_DATA_DIR.glob("*.json"): with open(f) as fh: data = json.load(fh) if isinstance(data, list): self.training_examples.extend(data) else: self.training_examples.append(data) logger.info(f"Loaded {len(self.training_examples)} training examples") except Exception as e: logger.warning(f"Training data load error: {e}") def analyze(self, query: str, context: Dict = None, history: List = None) -> Dict: if not self.ready: return self._fallback(query) try: knowledge_str = json.dumps(self.custom_knowledge, indent=1)[:2000] if self.custom_knowledge else "None" examples_str = "\n".join( f"Q: {ex.get('input','')}\nA: {ex.get('output','')}" for ex in self.training_examples[-3:] ) if self.training_examples else "None" context_str = f"\nCONTEXT:\n{json.dumps(context, indent=1)[:3000]}\n" if context else "" prompt = f"""{self.SYSTEM_PROMPT} KNOWLEDGE BASE (summary): {knowledge_str} TRAINING EXAMPLES: {examples_str} {context_str} USER QUERY: {query} Provide a comprehensive cybersecurity analysis:""" response = self.client.models.generate_content( model=GEMINI_MODEL, contents=prompt, config={"temperature": 0.3, "max_output_tokens": 2048}, ) text = response.text if response.text else "No response generated" text_lower = text.lower() if "critical" in text_lower: risk_level, risk_score = "Critical", 9.0 elif "high" in text_lower: risk_level, risk_score = "High", 7.0 elif "medium" in text_lower: risk_level, risk_score = "Medium", 5.0 else: risk_level, risk_score = "Low", 3.0 return { "response": text, "confidence": 0.85, "risk_level": risk_level, "risk_score": risk_score, "insights": [f"Analysis performed by Gemini ({GEMINI_MODEL})"], "recommendations": [], "model_used": GEMINI_MODEL, "timestamp": datetime.utcnow().isoformat(), } except Exception as e: logger.error(f"Gemini analysis error: {e}") return self._fallback(query) def _fallback(self, query: str) -> Dict: """ML-powered analysis when Gemini is unavailable""" import re global ml_loader url_pattern = re.compile(r'https?://[^\s"\'<>]+') urls = url_pattern.findall(query) risk_level = "Unknown" risk_score = 0.0 insights = [] response_lines = ["## CyberForge ML Security Analysis\n"] try: loader = ml_loader n_loaded = len(loader.models) except Exception: loader = None n_loaded = 0 if urls and loader and n_loaded > 0: threat_models_fired = [] all_scores = [] for url in urls[:3]: features = extract_url_features(url) url_threats = [] url_scores = [] for model_name in ["phishing_detection", "malware_detection", "web_attack_detection"]: pred = loader.predict(model_name, features) if pred.get("prediction", 0) == 1: label = model_name.replace("_detection", "").replace("_", " ").title() url_threats.append(label) url_scores.append(pred.get("confidence", 0.5)) threat_models_fired.append(model_name) avg = sum(url_scores) / len(url_scores) if url_scores else 0.15 all_scores.append(avg) lvl = "HIGH" if avg > 0.6 else "MEDIUM" if avg > 0.35 else "LOW" threats_str = ", ".join(url_threats) if url_threats else "None detected" display_url = url if len(url) <= 70 else url[:67] + "..." response_lines.append(f"**URL:** `{display_url}`") response_lines.append(f"- Risk: **{lvl}** | Threats: {threats_str} | Score: {avg:.0%}\n") overall = sum(all_scores) / len(all_scores) if all_scores else 0.2 if overall > 0.65: risk_level, risk_score = "High", 7.5 elif overall > 0.4: risk_level, risk_score = "Medium", 5.0 else: risk_level, risk_score = "Low", 2.5 insights = [f"{n_loaded} ML models active"] + [ f"Threat model triggered: {m.replace('_detection', '').replace('_', ' ').title()}" for m in set(threat_models_fired) ] if risk_level == "High": response_lines.append("### Recommendation\n⚠️ **Block immediately.** Phishing or malware indicators detected — do not visit this URL.") elif risk_level == "Medium": response_lines.append("### Recommendation\n⚡ **Exercise caution.** Validate with additional threat intelligence before accessing.") else: response_lines.append("### Recommendation\n✅ **URL appears structurally safe** based on ML analysis.") else: query_lower = query.lower() malware_kws = ["malware", "virus", "ransomware", "trojan", "spyware", "backdoor", "worm"] phishing_kws = ["phishing", "credential", "fake login", "spoof", "scam", "social engineering"] safe_kws = ["safe", "legitimate", "trusted", "secure", "verify"] if any(k in query_lower for k in malware_kws): risk_level, risk_score = "High", 7.0 response_lines.append( "**Malware indicators detected in your query.**\n\n" "Recommended actions:\n" "- Isolate the affected system immediately\n" "- Run a full endpoint detection scan\n" "- Review recently installed software and browser extensions\n" "- Check startup processes and scheduled tasks for persistence\n" "- Rotate credentials if any exposure is suspected" ) elif any(k in query_lower for k in phishing_kws): risk_level, risk_score = "High", 7.0 response_lines.append( "**Phishing threat indicators detected.**\n\n" "Recommended actions:\n" "- Do not submit credentials to the suspected site\n" "- Verify sender identity through a secondary channel\n" "- Report to your IT security team immediately\n" "- Enable MFA on all accounts that may be affected\n" "- Review email headers for spoofing indicators" ) elif any(k in query_lower for k in safe_kws): risk_level, risk_score = "Low", 2.0 response_lines.append("No immediate threat indicators detected. Continue standard monitoring procedures.") else: response_lines.append( f"CyberForge ML is operational with **{n_loaded}/4 models** loaded.\n\n" "For best results, include a URL or specific threat indicators in your query.\n\n" "**Available analysis capabilities:**\n" "- URL threat analysis (phishing, malware, web attacks)\n" "- Network anomaly detection\n" "- Real-time threat event monitoring\n\n" "*AI chat (Gemini) is currently unavailable. Provide a URL for full ML-based analysis.*" ) insights = [f"{n_loaded} ML models active"] response_lines.append("\n---\n*Powered by CyberForge ML models — Gemini AI offline.*") return { "response": "\n".join(response_lines), "confidence": 0.65 if urls else 0.4, "risk_level": risk_level, "risk_score": risk_score, "insights": insights, "recommendations": [], "model_used": "cyberforge-ml-fallback", "timestamp": datetime.utcnow().isoformat(), } # Singleton gemini_service = GeminiService() # ============================================================================ # ML MODEL LOADER (loads .pkl from HF Hub for REST predictions) # ============================================================================ class MLModelLoader: MODEL_NAMES = [ "phishing_detection", "malware_detection", "anomaly_detection", "web_attack_detection", ] def __init__(self): self.models: Dict[str, Any] = {} self.scalers: Dict[str, Any] = {} self.ready = False def initialize(self): loaded = 0 # All directories where models might exist (notebooks save to ../models) search_dirs = [ MODELS_DIR, # trained_models/ APP_DIR / "models", # models/ (where notebooks output) APP_DIR.parent / "models", # one level up fallback ] for name in self.MODEL_NAMES: if name in self.models: continue try: # 1. Try HuggingFace Hub first model_file = f"{name}/best_model.pkl" scaler_file = f"{name}/scaler.pkl" try: model_path = hf_hub_download( repo_id=HF_MODEL_REPO, filename=model_file, token=HF_TOKEN or None, cache_dir=str(MODELS_DIR), ) scaler_path = hf_hub_download( repo_id=HF_MODEL_REPO, filename=scaler_file, token=HF_TOKEN or None, cache_dir=str(MODELS_DIR), ) self.models[name] = joblib.load(model_path) self.scalers[name] = joblib.load(scaler_path) loaded += 1 logger.info(f"✅ Loaded model from Hub: {name}") continue except Exception: pass # 2. Try all local search directories for sdir in search_dirs: if name in self.models: break for model_fname in [f"{name}/best_model.pkl", f"{name}/model.pkl", f"{name}_model.pkl"]: candidate = sdir / model_fname if candidate.exists(): self.models[name] = joblib.load(candidate) # Try to find matching scaler for scaler_fname in [f"{name}/scaler.pkl", f"{name}_scaler.pkl"]: sc = sdir / scaler_fname if sc.exists(): self.scalers[name] = joblib.load(sc) break loaded += 1 logger.info(f"✅ Loaded model from {sdir.name}/{model_fname}: {name}") break except Exception as e: logger.warning(f"Error loading model {name}: {e}") # Sweep all search dirs for any .pkl files not yet loaded for sdir in search_dirs: if not sdir.exists(): continue for pkl in sdir.glob("*.pkl"): stem = pkl.stem.replace("_model", "").replace("_best", "") if stem not in self.models: try: self.models[stem] = joblib.load(pkl) loaded += 1 logger.info(f"✅ Loaded model sweep: {stem} from {sdir.name}") except Exception: pass self.ready = loaded > 0 logger.info(f"ML Models: {loaded} loaded — {list(self.models.keys())}") def predict(self, model_name: str, features: Dict) -> Dict: if model_name not in self.models: return self._heuristic_predict(model_name, features) try: model = self.models[model_name] scaler = self.scalers.get(model_name) X = np.array([list(features.values())]) if scaler: X = scaler.transform(X) prediction = int(model.predict(X)[0]) confidence = 0.5 probabilities = {} if hasattr(model, "predict_proba"): proba = model.predict_proba(X)[0] confidence = float(max(proba)) probabilities = {str(i): round(float(p), 4) for i, p in enumerate(proba)} # Derive a consistent threat_score: probability of class-1 (threat) # so that the /analyze-url aggregation has a uniform field across all models if "1" in probabilities: threat_score = probabilities["1"] elif prediction == 1: threat_score = confidence else: threat_score = 1.0 - confidence return { "model": model_name, "prediction": prediction, "prediction_label": "threat" if prediction == 1 else "benign", "confidence": round(confidence, 4), "threat_score": round(threat_score, 4), "probabilities": probabilities, "inference_source": "ml_model", "timestamp": datetime.utcnow().isoformat(), } except Exception as e: logger.error(f"Prediction error {model_name}: {e}") return self._heuristic_predict(model_name, features) def feature_importance(self, model_name: str) -> Dict: """Return feature importances for tree-based models (explainability).""" feature_names = [ "url_length", "hostname_length", "path_length", "is_https", "has_ip_address", "has_suspicious_tld", "subdomain_count", "has_port", "query_params_count", "has_at_symbol", "has_double_slash", "special_char_count", ] if model_name not in self.models: return {"error": f"Model '{model_name}' not loaded", "available": list(self.models.keys())} model = self.models[model_name] try: if hasattr(model, "feature_importances_"): importances = model.feature_importances_ pairs = sorted( zip(feature_names, [round(float(v), 4) for v in importances]), key=lambda x: x[1], reverse=True, ) return { "model": model_name, "method": "gini_importance", "feature_importances": dict(pairs), "top_features": [p[0] for p in pairs[:5]], } elif hasattr(model, "coef_"): coefs = model.coef_[0] if model.coef_.ndim > 1 else model.coef_ pairs = sorted( zip(feature_names, [round(float(v), 4) for v in coefs]), key=lambda x: abs(x[1]), reverse=True, ) return { "model": model_name, "method": "logistic_coefficients", "feature_importances": dict(pairs), "top_features": [p[0] for p in pairs[:5]], } else: return { "model": model_name, "method": "not_available", "note": f"Model type {type(model).__name__} does not expose feature importances", } except Exception as e: return {"model": model_name, "error": str(e)} def _heuristic_predict(self, model_name: str, features: Dict) -> Dict: score = 0.0 reasons = [] is_https = features.get("is_https", features.get("has_https", 1)) has_ip = features.get("has_ip_address", 0) suspicious_tld = features.get("has_suspicious_tld", 0) url_length = features.get("url_length", 0) special_chars = features.get("special_char_count", 0) if not is_https: score += 0.3; reasons.append("No HTTPS") if has_ip: score += 0.25; reasons.append("IP address in URL") if suspicious_tld: score += 0.2; reasons.append("Suspicious TLD") if url_length > 100: score += 0.15; reasons.append("Very long URL") if special_chars > 10: score += 0.1; reasons.append("Many special characters") is_threat = score >= 0.5 return { "model": model_name, "prediction": 1 if is_threat else 0, "prediction_label": "threat" if is_threat else "benign", "confidence": min(score + 0.3, 0.95) if is_threat else max(0.6, 1.0 - score), "threat_score": score, "reasons": reasons, "inference_source": "heuristic", "timestamp": datetime.utcnow().isoformat(), } ml_loader = MLModelLoader() # ============================================================================ # URL FEATURE EXTRACTION # ============================================================================ def extract_url_features(url: str) -> Dict: parsed = urlparse(url) hostname = parsed.hostname or "" return { "url_length": len(url), "hostname_length": len(hostname), "path_length": len(parsed.path or ""), "is_https": 1 if parsed.scheme == "https" else 0, "has_ip_address": 1 if all(p.isdigit() for p in hostname.split(".")) and len(hostname.split(".")) == 4 else 0, "has_suspicious_tld": 1 if any(hostname.endswith(t) for t in [".xyz", ".tk", ".ml", ".ga", ".cf", ".top", ".buzz"]) else 0, "subdomain_count": max(0, len(hostname.split(".")) - 2), "has_port": 1 if parsed.port else 0, "query_params_count": len(parsed.query.split("&")) if parsed.query else 0, "has_at_symbol": 1 if "@" in url else 0, "has_double_slash": 1 if "//" in (parsed.path or "") else 0, "special_char_count": sum(1 for c in url if c in "!@#$%^&*()+={}[]|\\:;<>?,"), } # ============================================================================ # TRANSFORMER MODEL LOADER (Phase 3 — real HF transformer models) # ============================================================================ # # Loads pretrained BERT-based classifiers from the Hugging Face Hub. # Models are loaded lazily on first request to keep cold-start fast. # A 7B Security LLM is NOT loaded locally (too big for free tier) — # it's accessed via the HF Inference API on demand. class TransformerModelLoader: """Loads pretrained Transformer classifiers from the HF Hub on demand.""" # Model registry — name → HF repo + task description # NOTE: DGA detector uses an inline entropy heuristic, not a transformer # (YangYang-Research/dga-detection has a non-standard model config that # isn't loadable via transformers.pipeline()). REGISTRY = { "url_phishing_bert": { "repo": "elftsdmr/malware-url-detect", "task": "text-classification", "labels": ["benign", "malicious"], "desc": "BERT-based URL phishing/malware classifier", }, } # Model used via HF Inference API. ZySec-AI/SecurityLLM isn't served on the # free Inference API tier (404). Mistral-7B-Instruct-v0.3 is widely available # and works well for cyber Q&A. Override via SECURITY_LLM_MODEL env var. SECURITY_LLM_REPO = os.environ.get("SECURITY_LLM_MODEL", "mistralai/Mistral-7B-Instruct-v0.3") def __init__(self): self.pipelines = {} # name → transformers.Pipeline (lazy-loaded) self.load_errors = {} # name → last error message self.transformers_available = False try: import transformers # noqa: F401 import torch # noqa: F401 self.transformers_available = True logger.info("✅ transformers + torch available") except ImportError as e: logger.warning(f"⚠️ transformers/torch not installed: {e}") def _ensure(self, name: str): """Load a pipeline on first use. Returns the pipeline or None on failure.""" if name in self.pipelines: return self.pipelines[name] if not self.transformers_available: self.load_errors[name] = "transformers/torch not installed" return None if name not in self.REGISTRY: self.load_errors[name] = f"Unknown model: {name}" return None spec = self.REGISTRY[name] try: from transformers import pipeline logger.info(f"⏳ Loading {name} from {spec['repo']}...") pipe = pipeline(spec["task"], model=spec["repo"], device=-1, top_k=None) self.pipelines[name] = pipe logger.info(f"✅ Loaded {name}") return pipe except Exception as e: err = f"{type(e).__name__}: {str(e)[:200]}" self.load_errors[name] = err logger.error(f"❌ Failed to load {name}: {err}") return None def predict_url_phishing(self, url: str) -> Dict: """Classify a URL as benign or malicious using elftsdmr/malware-url-detect.""" pipe = self._ensure("url_phishing_bert") if pipe is None: return self._unavailable("url_phishing_bert") try: # Strip the protocol — model was trained on bare URLs text = url.replace("https://", "").replace("http://", "")[:512] result = pipe(text) scores = result[0] if isinstance(result[0], list) else result return self._format_classification(scores, "url_phishing_bert") except Exception as e: return self._error("url_phishing_bert", e) def predict_dga(self, domain: str) -> Dict: """Detect DGA-generated domains using a Shannon-entropy + character-pattern heuristic. Real DGA domains have high entropy, low pronounceability, no real word substrings.""" import math d = domain.lower().strip().split('.')[0][:60] # SLD only, ignore TLD if not d: return {"model": "dga_detector", "error": "empty domain"} # Shannon entropy of character distribution from collections import Counter freq = Counter(d) n = len(d) entropy = -sum((c / n) * math.log2(c / n) for c in freq.values()) # Vowel ratio — DGAs typically have very few vowels vowels = sum(1 for c in d if c in 'aeiou') vowel_ratio = vowels / n if n else 0 # Digit ratio — DGAs often mix digits in digits = sum(1 for c in d if c.isdigit()) digit_ratio = digits / n if n else 0 # Length signal — DGAs are usually 10-25 chars length_signal = 1.0 if 12 <= n <= 30 else 0.5 if 8 <= n <= 40 else 0.2 # Consonant runs — DGAs often have 4+ consonants in a row max_consonant_run = 0 run = 0 for c in d: if c.isalpha() and c not in 'aeiou': run += 1 max_consonant_run = max(max_consonant_run, run) else: run = 0 # Score combination — empirically tuned thresholds score = 0.0 if entropy > 3.5: score += 0.35 if vowel_ratio < 0.25: score += 0.20 if digit_ratio > 0.15: score += 0.15 if max_consonant_run >= 4: score += 0.20 score *= length_signal score = min(1.0, score) is_dga = score >= 0.45 return { "model": "dga_detector", "prediction": "dga" if is_dga else "legit", "is_threat": is_dga, "confidence": round(score * 100, 2) if is_dga else round((1 - score) * 100, 2), "threat_score": round(score, 4), "features": { "entropy": round(entropy, 3), "vowel_ratio": round(vowel_ratio, 3), "digit_ratio": round(digit_ratio, 3), "max_consonant_run": max_consonant_run, "length": n, }, "inference_source": "entropy-heuristic", } def security_chat(self, query: str, max_tokens: int = 512) -> Dict: """Cybersecurity Q&A via ZySec-AI/SecurityLLM hosted on HF Inference API. Falls back to Gemini when the LLM is rate-limited or HF token is missing.""" if not HF_TOKEN: return { "model": "security-llm", "response": None, "error": "HF_TOKEN env var required to call ZySec-AI/SecurityLLM via Inference API", "fallback_available": gemini_service.ready, } try: import requests url = f"https://api-inference.huggingface.co/models/{self.SECURITY_LLM_REPO}" headers = {"Authorization": f"Bearer {HF_TOKEN}", "Content-Type": "application/json"} payload = { "inputs": query[:2000], "parameters": {"max_new_tokens": max_tokens, "temperature": 0.3, "return_full_text": False}, "options": {"wait_for_model": True}, } r = requests.post(url, headers=headers, json=payload, timeout=45) if r.status_code == 200: data = r.json() text = data[0].get("generated_text") if isinstance(data, list) and data else (data.get("generated_text") if isinstance(data, dict) else str(data)) return { "model": "security-llm", "source": "huggingface_inference_api", "response": text, "model_id": self.SECURITY_LLM_REPO, } return { "model": "security-llm", "error": f"HF Inference API HTTP {r.status_code}: {r.text[:200]}", "fallback_available": gemini_service.ready, } except Exception as e: return { "model": "security-llm", "error": f"{type(e).__name__}: {str(e)[:200]}", "fallback_available": gemini_service.ready, } def status(self) -> Dict: return { "transformers_available": self.transformers_available, "loaded": list(self.pipelines.keys()), "available": list(self.REGISTRY.keys()) + ["security_llm (via HF Inference API)"], "load_errors": self.load_errors, } @staticmethod def _format_classification(scores, model_name) -> Dict: """Normalize HF text-classification output to the cyberforge schema.""" if not scores: return {"model": model_name, "error": "Empty scores"} # scores is a list of {label, score} dicts. Find the threat label. threat_labels = {"malicious", "phishing", "malware", "dga", "label_1", "1"} # Top prediction top = max(scores, key=lambda s: s["score"]) if isinstance(scores, list) else scores is_threat = str(top["label"]).lower() in threat_labels # Threat score: probability of the threat class (not just the top class) threat_score = top["score"] if is_threat else 1.0 - top["score"] return { "model": model_name, "prediction": top["label"], "is_threat": is_threat, "confidence": round(top["score"] * 100, 2), "threat_score": round(threat_score, 4), "all_scores": scores, "inference_source": "huggingface_transformer", } @staticmethod def _unavailable(model_name) -> Dict: return {"model": model_name, "error": "Model unavailable — see /api/v2/status"} @staticmethod def _error(model_name, e: Exception) -> Dict: return {"model": model_name, "error": f"{type(e).__name__}: {str(e)[:200]}"} transformer_loader = TransformerModelLoader() # ============================================================================ # NOTEBOOK EXECUTION (existing Gradio functionality) # ============================================================================ def get_available_notebooks() -> List[str]: if not NOTEBOOKS_DIR.exists(): return [] return sorted([f.name for f in NOTEBOOKS_DIR.glob("*.ipynb")]) def read_notebook_content(notebook_name: str) -> str: """Read and display notebook content as markdown""" notebook_path = NOTEBOOKS_DIR / notebook_name if not notebook_path.exists(): return f"Notebook not found: {notebook_name}" try: with open(notebook_path, "r") as f: nb = json.load(f) output = f"# {notebook_name}\n\n" for i, cell in enumerate(nb.get("cells", []), 1): cell_type = cell.get("cell_type", "code") source = "".join(cell.get("source", [])) if cell_type == "markdown": output += f"{source}\n\n" else: output += f"### Cell {i} (Python)\n```python\n{source}\n```\n\n" return output except Exception as e: return f"Error reading notebook: {str(e)}" def execute_notebook(notebook_name: str, progress=gr.Progress()) -> Tuple[str, str]: """Execute a notebook and return output""" notebook_path = NOTEBOOKS_DIR / notebook_name output_path = NOTEBOOKS_DIR / f"output_{notebook_name}" if not notebook_path.exists(): available = list(NOTEBOOKS_DIR.glob("*.ipynb")) if NOTEBOOKS_DIR.exists() else [] return f"Error: Notebook not found: {notebook_path}\nAvailable: {available}", "" progress(0.1, desc="Starting notebook execution...") try: cmd = [ sys.executable, "-m", "nbconvert", "--to", "notebook", "--execute", "--output", str(output_path.absolute()), "--ExecutePreprocessor.timeout=600", "--ExecutePreprocessor.kernel_name=python3", str(notebook_path.absolute()), ] progress(0.3, desc="Executing cells...") result = subprocess.run(cmd, capture_output=True, text=True, cwd=str(NOTEBOOKS_DIR), timeout=900) progress(0.8, desc="Processing output...") if result.returncode == 0: if output_path.exists(): with open(output_path, "r") as f: executed_nb = json.load(f) outputs = [] for i, cell in enumerate(executed_nb.get("cells", []), 1): if cell.get("cell_type") == "code": for out in cell.get("outputs", []): if "text" in out: outputs.append(f"Cell {i}:\n{''.join(out['text'])}") elif "data" in out and "text/plain" in out["data"]: outputs.append(f"Cell {i}:\n{''.join(out['data']['text/plain'])}") progress(1.0, desc="Complete!") return "Notebook executed successfully!", "\n\n".join(outputs) else: return "Notebook executed but output file not found", result.stdout else: return f"Execution failed:\n{result.stderr}", result.stdout except subprocess.TimeoutExpired: return "Error: Notebook execution timed out (15 min limit)", "" except Exception as e: return f"Error executing notebook: {str(e)}", "" def run_notebook_cell(notebook_name: str, cell_number: int) -> str: """Execute a single cell from a notebook""" notebook_path = NOTEBOOKS_DIR / notebook_name if not notebook_path.exists(): return f"Error: Notebook not found at {notebook_path}" try: original_cwd = os.getcwd() os.chdir(NOTEBOOKS_DIR) with open(notebook_path, "r") as f: nb = json.load(f) cells = [c for c in nb.get("cells", []) if c.get("cell_type") == "code"] if cell_number < 1 or cell_number > len(cells): os.chdir(original_cwd) return f"Error: Cell {cell_number} not found. Available: 1-{len(cells)}" cell = cells[int(cell_number) - 1] source = "".join(cell.get("source", [])) import io from contextlib import redirect_stdout, redirect_stderr namespace = {"__name__": "__main__", "__file__": str(notebook_path)} stdout_capture = io.StringIO() stderr_capture = io.StringIO() with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture): try: exec(source, namespace) except Exception as e: os.chdir(original_cwd) return f"Error: {str(e)}" os.chdir(original_cwd) output = stdout_capture.getvalue() errors = stderr_capture.getvalue() result_text = f"### Cell {int(cell_number)} Output:\n" if output: result_text += f"```\n{output}\n```\n" if errors: result_text += f"\n**Warnings/Errors:**\n```\n{errors}\n```" if not output and not errors: result_text += "*(No output)*" return result_text except Exception as e: try: os.chdir(original_cwd) except Exception: pass return f"Error: {str(e)}" # ============================================================================ # MODEL TRAINING (existing Gradio functionality) # ============================================================================ class SecurityModelTrainer: def __init__(self): self.scaler = StandardScaler() self.label_encoder = LabelEncoder() def prepare_data(self, df: pd.DataFrame, target_col: str = "label") -> Tuple: if target_col not in df.columns: raise ValueError(f"Target column '{target_col}' not found") X = df.drop(columns=[target_col]) y = df[target_col] X = X.select_dtypes(include=[np.number]).fillna(0) if y.dtype == "object": y = self.label_encoder.fit_transform(y) X_scaled = self.scaler.fit_transform(X) return train_test_split(X_scaled, y, test_size=0.2, random_state=42) def train_model(self, model_type: str, X_train, y_train): if model_type not in MODEL_TYPES: raise ValueError(f"Unknown model type: {model_type}") model_class = MODEL_TYPES[model_type] if model_type == "Isolation Forest (Anomaly)": model = model_class(contamination=0.1, random_state=42) else: model = model_class(random_state=42) model.fit(X_train, y_train) return model def evaluate_model(self, model, X_test, y_test) -> Dict: y_pred = model.predict(X_test) return { "accuracy": accuracy_score(y_test, y_pred), "f1_score": f1_score(y_test, y_pred, average="weighted", zero_division=0), } trainer = SecurityModelTrainer() def train_model_from_data(data_file, model_type: str, task: str, progress=gr.Progress()): """Train model from uploaded data""" if data_file is None: return "Please upload a CSV file", None, None progress(0.1, desc="Loading data...") try: df = pd.read_csv(data_file.name) progress(0.3, desc="Preparing data...") X_train, X_test, y_train, y_test = trainer.prepare_data(df) progress(0.5, desc=f"Training {model_type}...") model = trainer.train_model(model_type, X_train, y_train) progress(0.8, desc="Evaluating model...") metrics = trainer.evaluate_model(model, X_test, y_test) model_name = f"{task.lower().replace(' ', '_')}_{model_type.lower().replace(' ', '_')}" model_path = MODELS_DIR / f"{model_name}.pkl" joblib.dump(model, model_path) progress(1.0, desc="Complete!") result = f""" ## Training Complete! **Task:** {task} **Model:** {model_type} **Samples:** {len(df)} ### Metrics - Accuracy: {metrics['accuracy']:.4f} - F1 Score: {metrics['f1_score']:.4f} **Model saved to:** {model_path} """ return result, str(model_path), json.dumps(metrics, indent=2) except Exception as e: return f"Error: {str(e)}", None, None def run_inference(model_file, features_text: str): """Run inference with a trained model""" if model_file is None: return "Please upload a model file" try: model = joblib.load(model_file.name) features = json.loads(features_text) X = np.array([list(features.values())]) prediction = model.predict(X)[0] result = {"prediction": int(prediction), "features_used": len(features)} if hasattr(model, "predict_proba"): proba = model.predict_proba(X)[0] result["confidence"] = float(max(proba)) result["probabilities"] = {str(i): float(p) for i, p in enumerate(proba)} return json.dumps(result, indent=2) except Exception as e: return f"Error: {str(e)}" def list_trained_models(): models = list(MODELS_DIR.glob("*.pkl")) if not models: return "No trained models found" output = "## Trained Models\n\n" for model_path in models: size_kb = model_path.stat().st_size / 1024 output += f"- **{model_path.name}** ({size_kb:.1f} KB)\n" return output # ============================================================================ # FASTAPI APP (REST endpoints for Heroku backend mlService.js) # ============================================================================ api = FastAPI(title="CyberForge AI API", version="1.0.0") api.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) @api.get("/health") async def api_health(): return { "status": "healthy", "timestamp": datetime.utcnow().isoformat(), "services": { "gemini": gemini_service.ready, "ml_models": ml_loader.ready, "models_loaded": list(ml_loader.models.keys()), "transformers": transformer_loader.transformers_available, "transformer_models_loaded": list(transformer_loader.pipelines.keys()), "gradio_ui": True, }, "version": "2.0.0", "new_endpoints": [ "POST /api/v2/batch-analyze", "GET /api/v2/explain/{model_name}", "POST /api/v2/ioc-scan", "POST /api/v2/url-enrich", "POST /api/v2/chat (context-grounded security chatbot, Wave 2)", "POST /api/v2/security-chat", ], } @api.post("/analyze") async def api_analyze(request: Request): """Main analysis endpoint – called by backend mlService.chatWithAI()""" body = await request.json() query = body.get("query", "") context = body.get("context", {}) history = body.get("conversation_history", []) result = gemini_service.analyze(query, context, history) return result @api.post("/analyze-url") async def api_analyze_url(request: Request): """URL analysis – called by backend mlService.analyzeWebsite() Returns per-model predictions + unified aggregate with risk_score (0-100). All models now consistently include threat_score (probability of class-1). """ body = await request.json() url = body.get("url", "") if not url: return JSONResponse(status_code=400, content={"detail": "URL required"}) features = extract_url_features(url) predictions = {} for model_name in ml_loader.MODEL_NAMES: predictions[model_name] = ml_loader.predict(model_name, features) # threat_score is now always present on every prediction (ml_model + heuristic) scores = [p.get("threat_score", 0.2) for p in predictions.values()] avg_score = sum(scores) / len(scores) if scores else 0 max_score = max(scores) if scores else 0 # Heuristic risk factors derived from URL features risk_factors = [] if not features.get("is_https"): risk_factors.append("no_https") if features.get("has_ip_address"): risk_factors.append("ip_in_url") if features.get("has_suspicious_tld"): risk_factors.append("suspicious_tld") if features.get("url_length", 0) > 100: risk_factors.append("long_url") if features.get("has_at_symbol"): risk_factors.append("at_symbol_in_url") if features.get("subdomain_count", 0) > 2: risk_factors.append("excessive_subdomains") if features.get("special_char_count", 0) > 10: risk_factors.append("high_special_chars") if features.get("has_double_slash"): risk_factors.append("double_slash_in_path") # Overall risk level using the max score (conservative) if max_score > 0.8: overall_risk = "critical" elif max_score > 0.6: overall_risk = "high" elif max_score > 0.4: overall_risk = "medium" else: overall_risk = "low" # risk_score 0-100 for the Heroku backend (used in riskScore field) risk_score_100 = round(max_score * 100, 1) return { "url": url, "aggregate": { "average_threat_score": round(avg_score, 3), "max_threat_score": round(max_score, 3), "risk_score": risk_score_100, "overall_risk_level": overall_risk, "risk_factors": risk_factors, "models_flagged": sum(1 for p in predictions.values() if p.get("prediction", 0) == 1), "models_total": len(predictions), }, "model_predictions": predictions, "features_analyzed": features, "timestamp": datetime.utcnow().isoformat(), } @api.post("/scan-threats") async def api_scan_threats(request: Request): """Threat scanning – called by backend mlService.scanForThreats()""" body = await request.json() query = json.dumps(body.get("data", body), indent=1)[:3000] result = gemini_service.analyze(f"Scan these indicators for threats:\n{query}") return result @api.post("/api/insights/generate") async def api_generate_insights(request: Request): """AI insights – called by backend mlService.getAIInsights()""" body = await request.json() query = body.get("query", "") context = body.get("context", {}) result = gemini_service.analyze(query, context) return { "insights": result.get("response", ""), "confidence": result.get("confidence", 0), "timestamp": datetime.utcnow().isoformat(), } @api.post("/api/models/predict") async def api_model_predict(request: Request): """Model prediction – called by backend mlService.getModelPrediction()""" body = await request.json() model_type = body.get("model_type", "phishing_detection") input_data = body.get("input_data", {}) result = ml_loader.predict(model_type, input_data) return result @api.get("/api/models/list") @api.get("/models") async def api_list_models(): result = [] for name in ml_loader.MODEL_NAMES: result.append({ "name": name, "loaded": name in ml_loader.models, "source": "ml_model" if name in ml_loader.models else "heuristic", }) return {"models": result, "total": len(ml_loader.MODEL_NAMES), "loaded": len(ml_loader.models)} # ============================================================================ # PHASE-3 ENDPOINTS — Real HF Transformer models # ============================================================================ @api.post("/api/v2/url-classify") async def api_v2_url_classify(request: Request): """URL phishing/malware classification using elftsdmr/malware-url-detect (BERT). Body: { "url": "https://..." }""" body = await request.json() url = body.get("url", "").strip() if not url: return JSONResponse(status_code=400, content={"detail": "url required"}) return transformer_loader.predict_url_phishing(url) @api.post("/api/v2/dga-detect") async def api_v2_dga_detect(request: Request): """DGA-generated domain detection using YangYang-Research/dga-detection. Body: { "domain": "abc123xyz.com" }""" body = await request.json() domain = body.get("domain", "").strip() if not domain: return JSONResponse(status_code=400, content={"detail": "domain required"}) return transformer_loader.predict_dga(domain) @api.post("/api/v2/security-chat") async def api_v2_security_chat(request: Request): """Cybersecurity Q&A via ZySec-AI/SecurityLLM (HF Inference API). Body: { "query": "...", "max_tokens": 512 }""" body = await request.json() query = body.get("query", "").strip() if not query: return JSONResponse(status_code=400, content={"detail": "query required"}) max_tokens = int(body.get("max_tokens", 512)) result = transformer_loader.security_chat(query, max_tokens=max_tokens) # Auto-fallback to Gemini when LLM unavailable if result.get("error") and gemini_service.ready: gemini_result = gemini_service.analyze(query) gemini_result["source"] = "gemini-fallback" gemini_result["llm_error"] = result["error"] return gemini_result return result @api.get("/api/v2/status") async def api_v2_status(): """Status of phase-3 transformer models — what's loaded, what failed, what's available.""" return transformer_loader.status() @api.post("/api/v2/batch-analyze") async def api_v2_batch_analyze(request: Request): """Batch URL analysis for the distributed agent system. Accepts up to 20 URLs in one call — avoids the per-URL HTTP overhead when multiple agents run concurrent scans. Body: { "urls": ["https://...", "https://...", ...] } Response: { "results": [ { ...same shape as /analyze-url... }, ... ], "total": N, "elapsed_ms": N } """ import time body = await request.json() urls = body.get("urls", []) if not urls: return JSONResponse(status_code=400, content={"detail": "urls array required"}) if len(urls) > 20: return JSONResponse(status_code=400, content={"detail": "Maximum 20 URLs per batch"}) t0 = time.time() results = [] for url in urls: url = str(url).strip() if not url: results.append({"url": url, "error": "empty url"}) continue features = extract_url_features(url) predictions = {} for model_name in ml_loader.MODEL_NAMES: predictions[model_name] = ml_loader.predict(model_name, features) scores = [p.get("threat_score", 0.2) for p in predictions.values()] avg_score = sum(scores) / len(scores) if scores else 0 max_score = max(scores) if scores else 0 risk_factors = [] if not features.get("is_https"): risk_factors.append("no_https") if features.get("has_ip_address"): risk_factors.append("ip_in_url") if features.get("has_suspicious_tld"): risk_factors.append("suspicious_tld") if features.get("url_length", 0) > 100: risk_factors.append("long_url") if features.get("has_at_symbol"): risk_factors.append("at_symbol_in_url") if features.get("subdomain_count", 0) > 2: risk_factors.append("excessive_subdomains") if max_score > 0.8: overall_risk = "critical" elif max_score > 0.6: overall_risk = "high" elif max_score > 0.4: overall_risk = "medium" else: overall_risk = "low" results.append({ "url": url, "aggregate": { "average_threat_score": round(avg_score, 3), "max_threat_score": round(max_score, 3), "risk_score": round(max_score * 100, 1), "overall_risk_level": overall_risk, "risk_factors": risk_factors, "models_flagged": sum(1 for p in predictions.values() if p.get("prediction", 0) == 1), }, "model_predictions": predictions, "features_analyzed": features, }) elapsed_ms = round((time.time() - t0) * 1000, 1) return { "results": results, "total": len(results), "elapsed_ms": elapsed_ms, "timestamp": datetime.utcnow().isoformat(), } @api.get("/api/v2/explain/{model_name}") async def api_v2_explain(model_name: str): """Feature importance / explainability for a loaded ML model. Supports RandomForest, GradientBoosting (gini importance) and LogisticRegression (coefficients). Path param: model_name — one of phishing_detection | malware_detection | anomaly_detection | web_attack_detection """ return ml_loader.feature_importance(model_name) @api.post("/api/v2/ioc-scan") async def api_v2_ioc_scan(request: Request): """Multi-field Indicator of Compromise scanner. Runs all available ML models + DGA detector + BERT classifier against multiple IOC types in a single call. Body: { "url": "https://...", # optional "domain": "example.com", # optional "ip": "1.2.3.4", # optional "hash": "abc123...", # optional (future use) "context": { ... } # optional — passed to Gemini if available } Response: per-indicator results with unified risk summary. """ body = await request.json() url = body.get("url", "").strip() domain = body.get("domain", "").strip() context = body.get("context", {}) ioc_results: Dict[str, Any] = {} all_scores: List[float] = [] # --- URL analysis --- if url: features = extract_url_features(url) url_preds = {name: ml_loader.predict(name, features) for name in ml_loader.MODEL_NAMES} url_scores = [p.get("threat_score", 0.2) for p in url_preds.values()] url_max = max(url_scores) if url_scores else 0 all_scores.append(url_max) # BERT-based URL classifier (lazy-loaded) bert_result = transformer_loader.predict_url_phishing(url) ioc_results["url"] = { "value": url, "ml_predictions": url_preds, "bert_classification": bert_result, "max_threat_score": round(url_max, 3), } # --- Domain analysis --- effective_domain = domain if not effective_domain and url: from urllib.parse import urlparse as _urlparse effective_domain = _urlparse(url).hostname or "" if effective_domain: dga_result = transformer_loader.predict_dga(effective_domain) dga_score = dga_result.get("threat_score", 0.0) all_scores.append(dga_score) ioc_results["domain"] = { "value": effective_domain, "dga_detection": dga_result, } # --- Overall risk summary --- overall_score = max(all_scores) if all_scores else 0.0 if overall_score > 0.8: risk_level = "critical" elif overall_score > 0.6: risk_level = "high" elif overall_score > 0.4: risk_level = "medium" elif overall_score > 0.0: risk_level = "low" else: risk_level = "unknown" # Optional Gemini enrichment when available ai_analysis = None if gemini_service.ready and (url or domain): target = url or domain ai_analysis = gemini_service.analyze( f"Scan these IOC indicators for threats: url={target}, domain={effective_domain}", context=context, ) return { "ioc_results": ioc_results, "summary": { "overall_risk_level": risk_level, "overall_threat_score": round(overall_score, 3), "risk_score": round(overall_score * 100, 1), "indicators_analyzed": len(ioc_results), }, "ai_analysis": ai_analysis, "timestamp": datetime.utcnow().isoformat(), } @api.post("/api/v2/url-enrich") async def api_v2_url_enrich(request: Request): """Enriched URL analysis: combines heuristic + ML models + BERT transformer into a single scored response with explanations. This is the richer alternative to /analyze-url for the AI Deep Scan mode. Body: { "url": "https://..." } """ body = await request.json() url = body.get("url", "").strip() if not url: return JSONResponse(status_code=400, content={"detail": "url required"}) features = extract_url_features(url) # Core ML predictions ml_predictions = {name: ml_loader.predict(name, features) for name in ml_loader.MODEL_NAMES} ml_scores = [p.get("threat_score", 0.2) for p in ml_predictions.values()] ml_max = max(ml_scores) if ml_scores else 0 # BERT transformer bert_result = transformer_loader.predict_url_phishing(url) bert_score = bert_result.get("threat_score", 0.0) if not bert_result.get("error") else None # DGA check on hostname from urllib.parse import urlparse as _urlparse hostname = _urlparse(url).hostname or "" dga_result = transformer_loader.predict_dga(hostname) if hostname else None dga_score = dga_result.get("threat_score", 0.0) if dga_result and not dga_result.get("error") else 0.0 # Feature importances for top model top_model = ml_loader.MODEL_NAMES[0] top_score = ml_scores[0] for i, score in enumerate(ml_scores): if score > top_score: top_score = score top_model = ml_loader.MODEL_NAMES[i] explainability = ml_loader.feature_importance(top_model) # Fuse scores: weighted average of ML (60%), BERT (30% if available), DGA (10%) if bert_score is not None: fused_score = 0.60 * ml_max + 0.30 * bert_score + 0.10 * dga_score else: fused_score = 0.80 * ml_max + 0.20 * dga_score if fused_score > 0.8: risk_level = "critical" elif fused_score > 0.6: risk_level = "high" elif fused_score > 0.4: risk_level = "medium" else: risk_level = "low" risk_factors = [] if not features.get("is_https"): risk_factors.append("no_https") if features.get("has_ip_address"): risk_factors.append("ip_in_url") if features.get("has_suspicious_tld"): risk_factors.append("suspicious_tld") if features.get("url_length", 0) > 100: risk_factors.append("long_url") if features.get("has_at_symbol"): risk_factors.append("at_symbol_in_url") if features.get("subdomain_count", 0) > 2: risk_factors.append("excessive_subdomains") if dga_result and dga_result.get("is_threat"): risk_factors.append("dga_domain") if bert_result and bert_result.get("is_threat"): risk_factors.append("bert_flagged_malicious") return { "url": url, "risk_level": risk_level, "risk_score": round(fused_score * 100, 1), "fused_threat_score": round(fused_score, 4), "components": { "ml_max_threat_score": round(ml_max, 4), "bert_threat_score": round(bert_score, 4) if bert_score is not None else None, "dga_threat_score": round(dga_score, 4), }, "ml_predictions": ml_predictions, "bert_classification": bert_result, "dga_detection": dga_result, "explainability": explainability, "risk_factors": risk_factors, "features_analyzed": features, "timestamp": datetime.utcnow().isoformat(), } @api.post("/api/v2/chat") async def api_v2_chat(request: Request): """Security Chatbot — Wave 2 implementation. Context-grounded security assistant that: 1. Accepts a live context blob (telemetry, scan history, IOCs) and grounds every answer in it (RAG-lite over the current request — no persistent vector store, which requires paid GPU on free tier). 2. Translates raw ML signals / IOCs into plain human language when the user includes raw data and asks for an explanation. 3. Falls back to the Security LLM (Mistral-7B via HF Inference API) or the ML-powered heuristic analyzer when Gemini is offline. NOTE ON REAL-TIME "LEARNING": HuggingFace Spaces cpu-basic does NOT support fine-tuning or persistent model updates at runtime. What this endpoint does instead is context injection — every turn receives the latest system telemetry + scan history in the prompt, so the chatbot's answers are grounded in live data without weight updates. This is the correct RAG-lite pattern for free-tier inference. Request body (stable contract): { "message": "What is the risk of this domain: evil.xyz?", "session_id": "optional-uuid-for-conversation-continuity", "conversation_history": [ {"role": "user", "content": "..."}, ... ], "context": { "telemetry": { "cpu": 42, "ram": 67, "net_in_kbps": 120, ... }, "recent_scans": [ { "url": "...", "risk_score": 78, "category": "phishing" }, ... ], "active_threats": [ { "type": "phishing", "severity": "high", "source": "..." } ], "behavioral_alerts": [ { "pattern": "...", "score": 0.9 } ], "translate": true # if true → force plain-language IOC translation mode } } Response contract (preserved from Wave 1): { "response": "...", # natural language answer grounded in context "confidence": 0.0-1.0, "risk_level": "low|medium|high|critical", "risk_score": 0-100, "model_used": "...", "session_id": "...", "sources": [ { "type": "...", "label": "...", "value": "..." } ], "timestamp": "..." } """ body = await request.json() message = body.get("message", "").strip() session_id = body.get("session_id", "") history = body.get("conversation_history", []) context = body.get("context", {}) if not message: return JSONResponse(status_code=400, content={"detail": "message required"}) # ── Build grounding context string and source citations ────────────── grounding_lines: List[str] = [] sources: List[Dict] = [] # System telemetry telemetry = context.get("telemetry", {}) if telemetry: cpu = telemetry.get("cpu", telemetry.get("cpu_percent")) ram = telemetry.get("ram", telemetry.get("ram_percent")) net_in = telemetry.get("net_in_kbps", telemetry.get("net_in")) telem_parts = [] if cpu is not None: telem_parts.append(f"CPU {cpu}%") if ram is not None: telem_parts.append(f"RAM {ram}%") if net_in is not None: telem_parts.append(f"Net↓ {net_in} kbps") if telem_parts: grounding_lines.append(f"SYSTEM TELEMETRY (live, last 5s): {', '.join(telem_parts)}") sources.append({"type": "telemetry", "label": "System Telemetry", "value": ", ".join(telem_parts)}) # Recent scan history recent_scans = context.get("recent_scans", []) if recent_scans: scan_summary = "; ".join( f"{s.get('url', 'unknown')} → {s.get('category', '?')} (score {s.get('risk_score', '?')}/100)" for s in recent_scans[:5] ) grounding_lines.append(f"RECENT URL SCANS (last 5): {scan_summary}") sources.append({"type": "scan_history", "label": "Recent Scans", "value": f"{len(recent_scans)} scans"}) # Active threats active_threats = context.get("active_threats", []) if active_threats: threat_summary = "; ".join( f"{t.get('type', '?')} [{t.get('severity', '?')}] from {t.get('source', '?')}" for t in active_threats[:5] ) grounding_lines.append(f"ACTIVE THREATS: {threat_summary}") sources.append({"type": "threats", "label": "Active Threats", "value": f"{len(active_threats)} active"}) # Behavioral alerts behavioral_alerts = context.get("behavioral_alerts", []) if behavioral_alerts: alert_summary = "; ".join( f"'{a.get('pattern', a.get('type', '?'))}' (score {a.get('score', '?')})" for a in behavioral_alerts[:3] ) grounding_lines.append(f"BEHAVIORAL ALERTS: {alert_summary}") sources.append({"type": "behavioral", "label": "Behavioral Alerts", "value": f"{len(behavioral_alerts)} alerts"}) # Machine→Human translation mode translate_mode = context.get("translate", False) raw_data = context.get("raw_data", "") if translate_mode and raw_data: grounding_lines.append( f"RAW DATA TO TRANSLATE INTO PLAIN LANGUAGE:\n{str(raw_data)[:1500]}" ) sources.append({"type": "raw_translation", "label": "Raw Signal", "value": str(raw_data)[:100]}) # Construct enriched prompt grounding_block = "\n".join(grounding_lines) translate_instruction = ( "\n\nIMPORTANT: The user has requested plain-language translation of raw machine data. " "Explain all technical signals, scores, and IOCs in simple, non-technical language first, " "then give actionable security advice." if translate_mode else "" ) conversation_instruction = "" if history: conversation_instruction = "\n\nCONVERSATION CONTEXT (recent turns):\n" + "\n".join( f"{h.get('role','?').upper()}: {str(h.get('content',''))[:300]}" for h in history[-6:] ) enriched_message = message if grounding_block: enriched_message = ( f"[LIVE SYSTEM CONTEXT — use this to ground your answer]\n" f"{grounding_block}\n" f"{translate_instruction}" f"{conversation_instruction}\n\n" f"[USER MESSAGE]\n{message}" ) # ── Run ML analysis on any URLs in the message (always available) ── import re as _re url_matches = _re.findall(r'https?://[^\s"\'<>]+', message) ml_url_context: List[str] = [] detected_risk_score = 0.0 detected_risk_level = "unknown" if url_matches and ml_loader.ready: for url in url_matches[:3]: features = extract_url_features(url) preds = {name: ml_loader.predict(name, features) for name in ml_loader.MODEL_NAMES} scores = [p.get("threat_score", 0.2) for p in preds.values()] mx = max(scores) if scores else 0 if mx > detected_risk_score: detected_risk_score = mx flagged = [n for n, p in preds.items() if p.get("prediction", 0) == 1] ml_url_context.append( f"ML scan of {url[:60]}: max_threat={mx:.2%}, flagged_by={flagged or 'none'}" ) sources.append({"type": "ml_scan", "label": f"ML scan: {url[:50]}", "value": f"risk {mx:.0%}"}) if ml_url_context: enriched_message += "\n\n[ML PRE-SCAN RESULTS]\n" + "\n".join(ml_url_context) if detected_risk_score > 0.8: detected_risk_level = "critical" elif detected_risk_score > 0.6: detected_risk_level = "high" elif detected_risk_score > 0.4: detected_risk_level = "medium" elif detected_risk_score > 0.0: detected_risk_level = "low" # ── Grounding: add ML model state to context ────────────────────── if ml_loader.ready: n_models = len(ml_loader.models) sources.append({"type": "ml_models", "label": "ML Models", "value": f"{n_models}/4 loaded"}) # ── LLM cascade: Gemini → Security LLM → ML heuristic ──────────── # Path 1: Gemini (premium, requires GEMINI_API_KEY env var) if gemini_service.ready: result = gemini_service.analyze(enriched_message, context=context, history=history) # Merge detected risk if ML gave a stronger signal if detected_risk_score > 0 and detected_risk_level != "unknown": ml_score_100 = round(detected_risk_score * 100, 1) if ml_score_100 > result.get("risk_score", 0): result["risk_score"] = ml_score_100 result["risk_level"] = detected_risk_level result["session_id"] = session_id result["sources"] = sources return result # Path 2: Security LLM via HF Inference API (Mistral-7B) llm_result = transformer_loader.security_chat(enriched_message, max_tokens=600) if llm_result.get("response") and not llm_result.get("error"): text = llm_result["response"] text_lower = text.lower() if "critical" in text_lower: risk_level, risk_score = "critical", 85.0 elif "high" in text_lower: risk_level, risk_score = "high", 70.0 elif "medium" in text_lower: risk_level, risk_score = "medium", 45.0 else: risk_level, risk_score = "low", 20.0 # Override with ML scan if stronger if detected_risk_score > 0: ml_score_100 = round(detected_risk_score * 100, 1) if ml_score_100 > risk_score: risk_score = ml_score_100 risk_level = detected_risk_level return { "response": text, "confidence": 0.72, "risk_level": risk_level, "risk_score": risk_score, "model_used": llm_result.get("model_id", transformer_loader.SECURITY_LLM_REPO), "session_id": session_id, "sources": sources, "timestamp": datetime.utcnow().isoformat(), } # Path 3: ML heuristic fallback (always available) fallback = gemini_service._fallback(enriched_message) # Override with ML scan risk if available if detected_risk_score > 0: ml_score_100 = round(detected_risk_score * 100, 1) if ml_score_100 > fallback.get("risk_score", 0): fallback["risk_score"] = ml_score_100 fallback["risk_level"] = detected_risk_level fallback["session_id"] = session_id fallback["sources"] = sources return fallback @api.post("/api/analysis/network") async def api_analyze_network(request: Request): """Network traffic analysis – called by backend mlService.analyzeNetworkTraffic()""" body = await request.json() traffic_data = body.get("traffic_data", {}) result = gemini_service.analyze( f"Analyze this network traffic for security threats:\n{json.dumps(traffic_data, indent=1)[:3000]}" ) return result @api.post("/api/ai/execute-task") async def api_execute_task(request: Request): """AI task execution – called by backend mlService.executeAITask()""" body = await request.json() task_type = body.get("task_type", "") task_data = body.get("task_data", {}) result = gemini_service.analyze( f"Execute cybersecurity task '{task_type}':\n{json.dumps(task_data, indent=1)[:3000]}" ) return result @api.post("/api/browser/analyze") async def api_analyze_browser(request: Request): """Browser session analysis""" body = await request.json() session = body.get("session_data", body) result = gemini_service.analyze( f"Analyze this browser session for security threats:\n{json.dumps(session, indent=1)[:3000]}" ) return result @api.post("/api/threat-feeds/analyze") async def api_analyze_threat_feed(request: Request): """Threat feed analysis""" body = await request.json() result = gemini_service.analyze( f"Analyze this threat feed:\n{json.dumps(body, indent=1)[:3000]}" ) return result @api.post("/api/datasets/process") async def api_process_dataset(request: Request): body = await request.json() return { "status": "processed", "dataset_name": body.get("dataset_name", ""), "timestamp": datetime.utcnow().isoformat(), } # ============================================================================ # GRADIO INTERFACE (UI tabs – notebooks, training, inference, models, API) # ============================================================================ def create_interface(): with gr.Blocks(title="CyberForge AI") as demo: gr.Markdown(""" # 🔐 CyberForge AI - ML Training Platform Train cybersecurity ML models and run Jupyter notebooks on Hugging Face. """) with gr.Tabs(): # ============ NOTEBOOKS TAB ============ with gr.TabItem("📓 Notebooks"): gr.Markdown("### Run ML Pipeline Notebooks\nExecute the CyberForge ML notebooks directly in the cloud.") with gr.Row(): with gr.Column(scale=1): notebook_dropdown = gr.Dropdown( choices=get_available_notebooks(), label="Select Notebook", value=get_available_notebooks()[0] if get_available_notebooks() else None, ) refresh_btn = gr.Button("🔄 Refresh List") view_btn = gr.Button("👁 View Content", variant="secondary") execute_btn = gr.Button("▶ Execute Notebook", variant="primary") gr.Markdown("### Run Single Cell") cell_number = gr.Number(label="Cell Number", value=1, minimum=1) run_cell_btn = gr.Button("Run Cell") with gr.Column(scale=2): notebook_status = gr.Markdown("Select a notebook to view or execute.") notebook_output = gr.Markdown("", label="Output") def refresh_notebooks(): notebooks = get_available_notebooks() return gr.update(choices=notebooks, value=notebooks[0] if notebooks else None) refresh_btn.click(refresh_notebooks, outputs=notebook_dropdown) view_btn.click(read_notebook_content, inputs=notebook_dropdown, outputs=notebook_output) execute_btn.click(execute_notebook, inputs=notebook_dropdown, outputs=[notebook_status, notebook_output]) run_cell_btn.click(run_notebook_cell, inputs=[notebook_dropdown, cell_number], outputs=notebook_output) # ============ TRAIN MODEL TAB ============ with gr.TabItem("🎯 Train Model"): gr.Markdown("### Train a Security ML Model\nUpload your dataset and train a model for threat detection.") with gr.Row(): with gr.Column(): task_dropdown = gr.Dropdown(choices=SECURITY_TASKS, label="Security Task", value="Phishing Detection") model_dropdown = gr.Dropdown(choices=list(MODEL_TYPES.keys()), label="Model Type", value="Random Forest") data_upload = gr.File(label="Upload Training Data (CSV)", file_types=[".csv"]) train_btn = gr.Button("🚀 Train Model", variant="primary") with gr.Column(): train_output = gr.Markdown("Upload data and click Train to begin.") model_path_output = gr.Textbox(label="Model Path", visible=False) metrics_output = gr.Textbox(label="Metrics JSON", visible=False) train_btn.click(train_model_from_data, inputs=[data_upload, model_dropdown, task_dropdown], outputs=[train_output, model_path_output, metrics_output]) # ============ INFERENCE TAB ============ with gr.TabItem("🔍 Inference"): gr.Markdown("### Run Model Inference\nLoad a trained model and make predictions.") with gr.Row(): with gr.Column(): model_upload = gr.File(label="Upload Model (.pkl)") features_input = gr.Textbox(label="Features (JSON)", value='{"url_length": 50, "has_https": 1, "digit_count": 5}', lines=5) predict_btn = gr.Button("🎯 Predict", variant="primary") with gr.Column(): prediction_output = gr.Textbox(label="Prediction Result", lines=10) predict_btn.click(run_inference, inputs=[model_upload, features_input], outputs=prediction_output) # ============ MODELS TAB ============ with gr.TabItem("📦 Models"): gr.Markdown("### Trained Models") models_list = gr.Markdown(list_trained_models()) refresh_models_btn = gr.Button("🔄 Refresh") refresh_models_btn.click(list_trained_models, outputs=models_list) # ============ API TAB ============ with gr.TabItem("🔌 API"): gr.Markdown(""" ## API Integration — v2.0 ### Core Endpoints (stable, used by Heroku backend) | Method | Endpoint | Description | |--------|----------|-------------| | GET | `/health` | Health check — version, models, services | | POST | `/analyze` | AI chat / general analysis (Gemini or ML fallback) | | POST | `/analyze-url` | URL threat analysis — 4 ML models + risk factors | | POST | `/scan-threats` | Deep threat scan (Gemini or ML fallback) | | POST | `/api/insights/generate` | AI-generated security insights | | POST | `/api/models/predict` | Direct ML model prediction | | GET | `/models` | List loaded models with source info | ### v2 Endpoints (new — for distributed agents + AI Deep Scan) | Method | Endpoint | Description | |--------|----------|-------------| | POST | `/api/v2/batch-analyze` | Batch analyze up to 20 URLs (for multi-agent) | | GET | `/api/v2/explain/{model_name}` | Feature importances / explainability | | POST | `/api/v2/url-classify` | BERT-based URL phishing classifier | | POST | `/api/v2/dga-detect` | DGA domain detection (entropy heuristic) | | POST | `/api/v2/ioc-scan` | Multi-IOC scanner: URL + domain + DGA + BERT | | POST | `/api/v2/url-enrich` | Enriched analysis: ML + BERT + DGA + explainability | | POST | `/api/v2/chat` | Security chatbot — context-grounded, RAG-lite, machine→human translation | | POST | `/api/v2/security-chat` | Security LLM Q&A (Mistral-7B via HF Inference API) | | GET | `/api/v2/status` | Transformer model load status | ### Quick Examples ```bash # Batch analyze URLs (for distributed agents) curl -X POST https://che237-cyberforge.hf.space/api/v2/batch-analyze \\ -H "Content-Type: application/json" \\ -d '{"urls": ["https://example.com", "http://evil.xyz/login?id=1"]}' # Enriched URL analysis (Deep Scan mode) curl -X POST https://che237-cyberforge.hf.space/api/v2/url-enrich \\ -H "Content-Type: application/json" \\ -d '{"url": "https://example.com"}' # Feature explainability curl https://che237-cyberforge.hf.space/api/v2/explain/phishing_detection # IOC multi-field scan curl -X POST https://che237-cyberforge.hf.space/api/v2/ioc-scan \\ -H "Content-Type: application/json" \\ -d '{"url": "http://192.168.1.1/login", "domain": "192.168.1.1"}' # Chatbot (Wave 2 extension point) curl -X POST https://che237-cyberforge.hf.space/api/v2/chat \\ -H "Content-Type: application/json" \\ -d '{"message": "Is this URL suspicious: http://login.g00gle.com.xyz/verify"}' ``` """) gr.Markdown("---\n**CyberForge AI** | [GitHub](https://github.com/Che237/cyberforge) | [Datasets](https://huggingface.co/datasets/Che237/cyberforge-datasets)") return demo # ============================================================================ # MAIN – Mount Gradio on FastAPI # ============================================================================ # Initialize services at startup logger.info("🚀 Initializing CyberForge AI services...") gemini_service.initialize() ml_loader.initialize() logger.info("✅ Services initialized") # Create Gradio app and mount it on FastAPI demo = create_interface() app = gr.mount_gradio_app(api, demo, path="/") if __name__ == "__main__": import uvicorn port = int(os.environ.get("PORT", 7860)) uvicorn.run(app, host="0.0.0.0", port=port)