diff --git a/.env.example b/.env.example index b523af9..b062066 100755 --- a/.env.example +++ b/.env.example @@ -78,6 +78,12 @@ ENABLE_CVE_HUNT=true # NVD API key for higher rate limits: https://nvd.nist.gov/developers/request-an-api-key #NVD_API_KEY= +# NVIDIA NIM API key for free 40 RPM endpoint +NIM_API_KEY= + +# NVIDIA NIM Model (optional - defaults to openai/gpt-oss-120b) +#NIM_MODEL= + # GitHub token for exploit search (optional, increases rate limit) #GITHUB_TOKEN= diff --git a/backend/api/v1/providers.py b/backend/api/v1/providers.py index 995a303..9cfa00a 100644 --- a/backend/api/v1/providers.py +++ b/backend/api/v1/providers.py @@ -335,6 +335,7 @@ async def toggle_provider(provider_id: str, req: ToggleRequest): # Whitelist of env keys that can be modified via UI ALLOWED_ENV_KEYS = { "ANTHROPIC_API_KEY", "OPENAI_API_KEY", "GEMINI_API_KEY", "GOOGLE_API_KEY", + "NIM_API_KEY", "NIM_MODEL", "NIM_BASE_URL", "OPENROUTER_API_KEY", "TOGETHER_API_KEY", "FIREWORKS_API_KEY", "OLLAMA_HOST", "LMSTUDIO_HOST", "ENABLE_SMART_ROUTER", "ENABLE_REASONING", "ENABLE_CVE_HUNT", diff --git a/backend/api/v1/settings.py b/backend/api/v1/settings.py index d21d876..57bcdd9 100755 --- a/backend/api/v1/settings.py +++ b/backend/api/v1/settings.py @@ -71,6 +71,7 @@ class SettingsUpdate(BaseModel): llm_model: Optional[str] = None anthropic_api_key: Optional[str] = None openai_api_key: Optional[str] = None + nim_api_key: Optional[str] = None openrouter_api_key: Optional[str] = None gemini_api_key: Optional[str] = None together_api_key: Optional[str] = None @@ -103,6 +104,7 @@ class SettingsResponse(BaseModel): llm_model: str = "" has_anthropic_key: bool = False has_openai_key: bool = False + has_nim_key: bool = False has_openrouter_key: bool = False has_gemini_key: bool = False has_together_key: bool = False @@ -172,7 +174,9 @@ def _load_settings_from_env() -> dict: # Detect provider from which keys are set provider = "claude" - if os.getenv("ANTHROPIC_API_KEY"): + if os.getenv("NIM_API_KEY"): + provider = "nim" + elif os.getenv("ANTHROPIC_API_KEY"): provider = "claude" elif os.getenv("OPENAI_API_KEY"): provider = "openai" @@ -184,6 +188,7 @@ def _load_settings_from_env() -> dict: "llm_model": os.getenv("DEFAULT_LLM_MODEL", ""), "anthropic_api_key": os.getenv("ANTHROPIC_API_KEY", ""), "openai_api_key": os.getenv("OPENAI_API_KEY", ""), + "nim_api_key": os.getenv("NIM_API_KEY", ""), "openrouter_api_key": os.getenv("OPENROUTER_API_KEY", ""), "gemini_api_key": os.getenv("GEMINI_API_KEY", ""), "together_api_key": os.getenv("TOGETHER_API_KEY", ""), @@ -224,6 +229,7 @@ async def get_settings(): llm_model=_settings.get("llm_model", ""), has_anthropic_key=bool(_settings["anthropic_api_key"] or os.getenv("ANTHROPIC_API_KEY")), has_openai_key=bool(_settings["openai_api_key"] or os.getenv("OPENAI_API_KEY")), + has_nim_key=bool(_settings.get("nim_api_key") or os.getenv("NIM_API_KEY")), has_openrouter_key=bool(_settings["openrouter_api_key"] or os.getenv("OPENROUTER_API_KEY")), has_gemini_key=bool(_settings.get("gemini_api_key") or os.getenv("GEMINI_API_KEY")), has_together_key=bool(_settings.get("together_api_key") or os.getenv("TOGETHER_API_KEY")), @@ -275,6 +281,12 @@ async def update_settings(settings_data: SettingsUpdate): os.environ["OPENAI_API_KEY"] = settings_data.openai_api_key env_updates["OPENAI_API_KEY"] = settings_data.openai_api_key + if settings_data.nim_api_key is not None: + _settings["nim_api_key"] = settings_data.nim_api_key + if settings_data.nim_api_key: + os.environ["NIM_API_KEY"] = settings_data.nim_api_key + env_updates["NIM_API_KEY"] = settings_data.nim_api_key + if settings_data.openrouter_api_key is not None: _settings["openrouter_api_key"] = settings_data.openrouter_api_key if settings_data.openrouter_api_key: @@ -564,6 +576,11 @@ CLOUD_MODELS = { "codex": [ {"model_id": "codex-mini-latest", "display_name": "Codex Mini", "context_length": 192000}, ], + "nim": [ + {"model_id": "openai/gpt-oss-120b", "display_name": "NVIDIA GPT-OSS 120B", "context_length": 32768}, + {"model_id": "meta/llama-3.1-70b-instruct", "display_name": "Llama 3.1 70B (NIM)", "context_length": 32768}, + {"model_id": "meta/llama-3.1-405b-instruct", "display_name": "Llama 3.1 405B (NIM)", "context_length": 32768}, + ], } diff --git a/backend/config.py b/backend/config.py index bf7efd9..a6d49ab 100755 --- a/backend/config.py +++ b/backend/config.py @@ -32,6 +32,8 @@ class Settings(BaseSettings): # LLM Settings ANTHROPIC_API_KEY: Optional[str] = os.getenv("ANTHROPIC_API_KEY") OPENAI_API_KEY: Optional[str] = os.getenv("OPENAI_API_KEY") + NIM_API_KEY: Optional[str] = os.getenv("NIM_API_KEY") + NIM_BASE_URL: str = os.getenv("NIM_BASE_URL", "https://integrate.api.nvidia.com/v1/chat/completions") OPENROUTER_API_KEY: Optional[str] = os.getenv("OPENROUTER_API_KEY") GEMINI_API_KEY: Optional[str] = os.getenv("GEMINI_API_KEY") AZURE_OPENAI_API_KEY: Optional[str] = os.getenv("AZURE_OPENAI_API_KEY") diff --git a/backend/core/autonomous_agent.py b/backend/core/autonomous_agent.py index 51a121d..a0e40d1 100755 --- a/backend/core/autonomous_agent.py +++ b/backend/core/autonomous_agent.py @@ -217,6 +217,8 @@ except ImportError: ANTHROPIC_AVAILABLE = False anthropic = None + + # Try to import openai try: import openai @@ -367,6 +369,7 @@ class LLMClient: self.azure_openai_api_version = os.getenv("AZURE_OPENAI_API_VERSION", "2024-02-01") self.azure_openai_deployment = os.getenv("AZURE_OPENAI_DEPLOYMENT", "") self.codex_key = os.getenv("CODEX_API_KEY", "") + self.nim_key = os.getenv("NIM_API_KEY", "") self.ollama_model = os.getenv("OLLAMA_MODEL", "llama3.2") self.configured_model = os.getenv("DEFAULT_LLM_MODEL", "") # User-configured model name self.client = None @@ -419,6 +422,14 @@ class LLMClient: def _initialize_provider(self): """Initialize the first available LLM provider""" + # 0. Try NVIDIA NIM (NVIDIA's OpenAI-compatible endpoint) + if self.nim_key: + self.client = "nim" # Placeholder - uses HTTP requests + self.provider = "nim" + self.model_name = self.configured_model or os.getenv("NIM_MODEL", "openai/gpt-oss-120b") + print(f"[LLM] NVIDIA NIM initialized (model: {self.model_name})") + return + # 1. Try Claude (Anthropic) if ANTHROPIC_AVAILABLE and self.anthropic_key: try: @@ -573,6 +584,7 @@ class LLMClient: "openai_lib": OPENAI_AVAILABLE, "ollama_available": self._check_ollama(), "lmstudio_available": self._check_lmstudio(), + "has_nim_key": bool(self.nim_key), "has_google_key": bool(self.google_key), "has_together_key": bool(self.together_key), "has_fireworks_key": bool(self.fireworks_key), @@ -639,6 +651,14 @@ class LLMClient: ) return message.content[0].text + elif self.provider == "nim": + return await self._generate_openai_compatible( + prompt, system or default_system, max_tokens, + url=os.getenv("NIM_BASE_URL", "https://integrate.api.nvidia.com/v1/chat/completions"), + api_key=self.nim_key, + model=self.model_name or "openai/gpt-oss-120b", + ) + elif self.provider == "openai": response = self.client.chat.completions.create( model=self.model_name or "gpt-4o", diff --git a/backend/core/knowledge_processor.py b/backend/core/knowledge_processor.py index d8f9cbe..0eb1f4a 100644 --- a/backend/core/knowledge_processor.py +++ b/backend/core/knowledge_processor.py @@ -177,7 +177,7 @@ class KnowledgeProcessor: def _extract_text_pdf(self, file_path: Path) -> str: """Extract text from PDF.""" - if not HAS_PYPDKF2: + if not HAS_PYPDF2: logger.warning("PyPDF2 not installed - PDF extraction unavailable. Install: pip install PyPDF2") # Try reading as text fallback try: diff --git a/backend/core/smart_router/provider_registry.py b/backend/core/smart_router/provider_registry.py index 2bb883a..7b9ca96 100644 --- a/backend/core/smart_router/provider_registry.py +++ b/backend/core/smart_router/provider_registry.py @@ -77,6 +77,13 @@ class Provider: # Default provider definitions DEFAULT_PROVIDERS: List[Dict] = [ + # === NVIDIA NIM Provider (Tier 2 - Free) === + { + "id": "nim", "name": "NVIDIA NIM", "auth_type": "api_key", + "api_format": "openai_compat", "base_url": "https://integrate.api.nvidia.com/v1", + "tier": 2, "default_model": os.getenv("NIM_MODEL", "openai/gpt-oss-120b"), + "env_key": "NIM_API_KEY", + }, # === OAuth Providers (Tier 1 - Subscription) === { "id": "claude_code", "name": "Claude Code", "auth_type": "oauth", diff --git a/backend/core/smart_router/token_extractor.py b/backend/core/smart_router/token_extractor.py index 2555c2e..99e1ff2 100644 --- a/backend/core/smart_router/token_extractor.py +++ b/backend/core/smart_router/token_extractor.py @@ -32,7 +32,7 @@ class TokenExtractor: """Extracts tokens from CLI tools installed on the system.""" EXTRACTORS = [ - "claude_code", "codex_cli", "gemini_cli", "cursor", + "claude_code", "codex_cli", "cursor", "copilot", "iflow", "qwen_code", "kiro", ] @@ -129,46 +129,6 @@ class TokenExtractor: pass return None - def _extract_gemini_cli(self) -> Optional[ExtractedToken]: - """Gemini CLI: ~/.gemini/oauth_creds.json""" - creds_path = Path.home() / ".gemini" / "oauth_creds.json" - if not creds_path.exists(): - return None - try: - data = json.loads(creds_path.read_text()) - access_token = data.get("access_token") - refresh = data.get("refresh_token") - # Gemini CLI uses multiple field names for expiry - expires = ( - data.get("expiry_date") # Gemini CLI uses this (milliseconds) - or data.get("expires_at") - or data.get("expiry") - ) - if access_token and access_token.startswith("ya29."): - # Parse expiry: may be ms timestamp, unix timestamp, or ISO string - if isinstance(expires, (int, float)): - # If > 1e12, it's milliseconds — convert to seconds - if expires > 1e12: - expires = expires / 1000.0 - elif isinstance(expires, str): - try: - from datetime import datetime - dt = datetime.fromisoformat(expires.replace("Z", "+00:00")) - expires = dt.timestamp() - except Exception: - expires = None - return ExtractedToken( - provider_id="gemini_cli", - credential_type="oauth", - token=access_token, - refresh_token=refresh, - expires_at=expires, - label="Gemini CLI", - ) - except Exception: - pass - return None - def _extract_cursor(self) -> Optional[ExtractedToken]: """Cursor: SQLite state.vscdb database""" is_mac = platform.system() == "Darwin" diff --git a/backend/core/smart_router/token_refresher.py b/backend/core/smart_router/token_refresher.py index 46d29a7..e8420a1 100644 --- a/backend/core/smart_router/token_refresher.py +++ b/backend/core/smart_router/token_refresher.py @@ -104,8 +104,6 @@ class TokenRefresher: return await self._refresh_anthropic(account_id, refresh_token) elif provider_id == "codex_cli": return await self._refresh_openai(account_id, refresh_token) - elif provider_id == "gemini_cli": - return await self._refresh_google(account_id, refresh_token) else: logger.debug( f"TokenRefresher: No refresh method for {provider_id}" @@ -115,74 +113,6 @@ class TokenRefresher: logger.error(f"TokenRefresher: Refresh failed for {provider_id}: {e}") return False - async def _refresh_google(self, account_id: str, refresh_token: str) -> bool: - """Refresh a Google OAuth token.""" - try: - import aiohttp - - # Gemini CLI client_id/secret - extracted at runtime from CLI binary or set via env - gemini_client_id = os.environ.get("GEMINI_CLI_CLIENT_ID", "") - gemini_client_secret = os.environ.get("GEMINI_CLI_CLIENT_SECRET", "") - if not gemini_client_id or not gemini_client_secret: - # Try extracting from installed Gemini CLI - try: - from backend.core.smart_router.token_extractor import TokenExtractor - extractor = TokenExtractor() - creds = extractor._extract_gemini_oauth_creds() - if creds: - gemini_client_id = creds.get("client_id", "") - gemini_client_secret = creds.get("client_secret", "") - except Exception: - pass - payload = { - "grant_type": "refresh_token", - "refresh_token": refresh_token, - "client_id": gemini_client_id, - "client_secret": gemini_client_secret, - } - async with aiohttp.ClientSession() as session: - async with session.post(GOOGLE_TOKEN_URL, data=payload, timeout=aiohttp.ClientTimeout(total=10)) as resp: - if resp.status == 200: - data = await resp.json() - new_token = data.get("access_token") - expires_in = data.get("expires_in", 3600) - if new_token: - self._registry.update_credential( - account_id, - new_token, - time.time() + expires_in, - ) - # Save refreshed token to Gemini CLI creds file - self._save_gemini_cli_token(new_token, expires_in) - logger.info(f"TokenRefresher: Google token refreshed (expires in {expires_in}s)") - return True - else: - body = await resp.text() - logger.warning(f"TokenRefresher: Google refresh failed: {resp.status} {body[:200]}") - except ImportError: - logger.warning("TokenRefresher: aiohttp not available for token refresh") - except Exception as e: - logger.error(f"TokenRefresher: Google refresh error: {e}") - return False - - @staticmethod - def _save_gemini_cli_token(new_token: str, expires_in: int): - """Write refreshed token back to Gemini CLI credentials file on disk.""" - import json - from pathlib import Path - - creds_path = Path.home() / ".gemini" / "oauth_creds.json" - if not creds_path.exists(): - return - try: - data = json.loads(creds_path.read_text()) - data["access_token"] = new_token - data["expiry_date"] = int((time.time() + expires_in) * 1000) - creds_path.write_text(json.dumps(data, indent=2)) - logger.debug("TokenRefresher: Saved refreshed token to Gemini CLI creds file") - except Exception as e: - logger.debug(f"TokenRefresher: Failed to save Gemini CLI token: {e}") - async def _refresh_anthropic(self, account_id: str, refresh_token: str) -> bool: """Refresh an Anthropic/Claude OAuth token.""" try: diff --git a/backend/main.py b/backend/main.py index 6fa398b..57b2acd 100755 --- a/backend/main.py +++ b/backend/main.py @@ -128,11 +128,15 @@ async def health_check(): # Check LLM availability anthropic_key = os.getenv("ANTHROPIC_API_KEY", "") openai_key = os.getenv("OPENAI_API_KEY", "") + nim_key = os.getenv("NIM_API_KEY", "") llm_status = "not_configured" llm_provider = None - if anthropic_key and anthropic_key not in ["", "your-anthropic-api-key"]: + if nim_key and nim_key not in ["", "your-nim-api-key"]: + llm_status = "configured" + llm_provider = "nim" + elif anthropic_key and anthropic_key not in ["", "your-anthropic-api-key"]: llm_status = "configured" llm_provider = "claude" elif openai_key and openai_key not in ["", "your-openai-api-key"]: @@ -146,7 +150,7 @@ async def health_check(): "llm": { "status": llm_status, "provider": llm_provider, - "message": "AI agent ready" if llm_status == "configured" else "Set ANTHROPIC_API_KEY or OPENAI_API_KEY to enable AI features" + "message": "AI agent ready" if llm_status == "configured" else "Set ANTHROPIC_API_KEY, OPENAI_API_KEY or NIM_API_KEY to enable AI features" } } diff --git a/backend/requirements.txt b/backend/requirements.txt index 0f40b4d..2c10f87 100755 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -38,3 +38,6 @@ apscheduler>=3.10.0 httpx>=0.26.0 pytest>=7.4.0 pytest-asyncio>=0.23.0 + +# Document Processing +PyPDF2>=3.0.0 diff --git a/core/llm_manager.py b/core/llm_manager.py index d28627c..1e89837 100755 --- a/core/llm_manager.py +++ b/core/llm_manager.py @@ -25,8 +25,19 @@ logger = logging.getLogger(__name__) class LLMManager: """Manage multiple LLM providers""" - def __init__(self, config: Dict): + def __init__(self, config: Optional[Dict] = None): """Initialize LLM manager""" + if config is None: + # Try to load from default config file or environment + config = {} + try: + config_path = Path("config/config.json") + if config_path.exists(): + with open(config_path, 'r') as f: + config = json.load(f) + except Exception as e: + logger.warning(f"Could not load default config: {e}") + self.config = config.get('llm', {}) self.default_profile_name = self.config.get('default_profile', 'gemini_pro_default') self.profiles = self.config.get('profiles', {}) @@ -34,8 +45,37 @@ class LLMManager: self.active_profile = self.profiles.get(self.default_profile_name, {}) # Load active profile settings - self.provider = self.active_profile.get('provider', 'gemini').lower() - self.model = self.active_profile.get('model', 'gemini-pro') + self.provider = self.active_profile.get('provider', '').lower() + self.model = self.active_profile.get('model', '') + + # Overriding priority: If NIM_API_KEY is in env and we are using a default/empty provider, use nim + if (not self.provider or self.provider == 'gemini') and os.getenv("NIM_API_KEY"): + self.provider = "nim" + # If the model was specifically gemini-pro (default) or empty, change to a NIM model + if not self.model or self.model == 'gemini-pro': + self.model = os.getenv("DEFAULT_LLM_MODEL", "meta/llama-3.1-70b-instruct") + elif not self.provider: + # Detect from environment if not in config + if os.getenv("ANTHROPIC_API_KEY"): + self.provider = "claude" + elif os.getenv("OPENAI_API_KEY"): + self.provider = "gpt" + elif os.getenv("GEMINI_API_KEY"): + self.provider = "gemini" + else: + self.provider = "gemini" # Final fallback + + # Default models per provider if still empty + default_models = { + "nim": "meta/llama-3.1-70b-instruct", + "claude": "claude-3-5-sonnet-20240620", + "gpt": "gpt-4o", + "gemini": "gemini-1.5-pro" + } + + if not self.model: + self.model = os.getenv("DEFAULT_LLM_MODEL", default_models.get(self.provider, "gemini-pro")) + self.api_key = self._get_api_key(self.active_profile.get('api_key', '')) self.temperature = self.active_profile.get('temperature', 0.7) self.max_tokens = self.active_profile.get('max_tokens', 4096) @@ -151,6 +191,8 @@ class LLMManager: try: if self.provider == 'claude': raw_response = self._generate_claude(prompt, system_prompt) + elif self.provider == 'nim': + raw_response = self._generate_nim(prompt, system_prompt) elif self.provider == 'gpt': raw_response = self._generate_gpt(prompt, system_prompt) elif self.provider == 'gemini': @@ -282,6 +324,42 @@ Identify any potential hallucinations, inconsistencies, or areas where the respo finally: self.hallucination_mitigation_strategy = original_mitigation_state # Restore original state + def _generate_nim(self, prompt: str, system_prompt: Optional[str] = None) -> str: + """Generate using NVIDIA NIM API (OpenAI-compatible)""" + api_key = os.getenv("NIM_API_KEY", self.api_key) + if not api_key: + raise ValueError("NIM_API_KEY not set.") + + base_url = os.getenv("NIM_BASE_URL", "https://integrate.api.nvidia.com/v1/chat/completions") + url = base_url + headers = { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json" + } + + messages = [] + if system_prompt: + messages.append({"role": "system", "content": system_prompt}) + messages.append({"role": "user", "content": prompt}) + + data = { + "model": self.model or "openai/gpt-oss-120b", + "messages": messages, + "temperature": self.temperature, + "max_tokens": self.max_tokens + } + + try: + response = requests.post(url, headers=headers, json=data, timeout=180) + if response.status_code == 200: + result = response.json() + return result["choices"][0]["message"]["content"] + else: + raise ValueError(f"NIM API error {response.status_code}: {response.text}") + except Exception as e: + logger.error(f"NIM error: {e}") + raise + def _generate_claude(self, prompt: str, system_prompt: Optional[str] = None) -> str: """Generate using Claude API with requests (bypasses httpx/SSL issues on macOS)""" if not self.api_key: diff --git a/data/custom-knowledge/index.json b/data/custom-knowledge/index.json index 4e57a36..76affd9 100644 --- a/data/custom-knowledge/index.json +++ b/data/custom-knowledge/index.json @@ -13,7 +13,10 @@ "knowledge_entries": [] } ], - "vuln_type_index": {}, + "vuln_type_index": { + "information_disclosure": [], + "clickjacking": [] + }, "version": "1.0", - "updated_at": "2026-02-16T14:50:31.618046" + "updated_at": "2026-04-28T19:00:40.997968" } \ No newline at end of file diff --git a/docker-compose.yml b/docker-compose.yml index 486e37b..cff5b96 100755 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -12,6 +12,7 @@ services: # These override .env if set - ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY:-} - OPENAI_API_KEY=${OPENAI_API_KEY:-} + - NIM_API_KEY=${NIM_API_KEY:-} - DATABASE_URL=sqlite+aiosqlite:///./data/neurosploit.db volumes: - neurosploit-data:/app/data diff --git a/docker/Dockerfile.kali b/docker/Dockerfile.kali index fc10541..7864332 100755 --- a/docker/Dockerfile.kali +++ b/docker/Dockerfile.kali @@ -18,7 +18,7 @@ # - Full Kali apt repos available for on-demand apt-get install of any security tool # ---- Stage 1: Pre-compile Go security tools ---- -FROM golang:1.24-bookworm AS go-builder +FROM golang:1.26-bookworm AS go-builder RUN apt-get update && apt-get install -y --no-install-recommends \ git build-essential libpcap-dev \ diff --git a/docker/Dockerfile.sandbox b/docker/Dockerfile.sandbox index fcd90b4..dbdbba0 100755 --- a/docker/Dockerfile.sandbox +++ b/docker/Dockerfile.sandbox @@ -2,7 +2,7 @@ # Kali-based container with real penetration testing tools # Provides Nuclei, Naabu, and other ProjectDiscovery tools via isolated execution -FROM golang:1.24-bookworm AS go-builder +FROM golang:1.26-bookworm AS go-builder RUN apt-get update && apt-get install -y --no-install-recommends git build-essential && \ rm -rf /var/lib/apt/lists/* diff --git a/frontend/index.html b/frontend/index.html index 32b2743..9fda857 100755 --- a/frontend/index.html +++ b/frontend/index.html @@ -5,6 +5,9 @@
{subtitle}
} +{subtitle}
}