diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..71e7d36 --- /dev/null +++ b/.gitignore @@ -0,0 +1,38 @@ +# Local secrets +.HexSec +.env +.env.* +!.env.example + +# Python bytecode and caches +__pycache__/ +*.py[cod] +*$py.class +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ + +# Virtual environments +.venv/ +venv/ +env/ +ENV/ + +# Test and coverage artifacts +.coverage +.coverage.* +htmlcov/ + +# Build artifacts +build/ +dist/ +*.egg-info/ + +# Logs +*.log + +# Editor and OS files +.vscode/ +.idea/ +.DS_Store +Thumbs.db diff --git a/HexSecGPT.py b/HexSecGPT.py index d487a16..60fcc93 100644 --- a/HexSecGPT.py +++ b/HexSecGPT.py @@ -94,39 +94,149 @@ def check_dependencies(): tngtech/deepseek-r1t2-chimera:free """ # --- Configuration --- -class Config: - """System Configuration & Constants""" +class Config: + """System Configuration & Constants""" + + # API Provider Settings + PROVIDER_ENV_NAME = "HEXSECGPT_PROVIDER" + LEGACY_API_KEY_NAME = "HexSecGPT-API" + DEFAULT_PROVIDER = "openrouter" + ZAI_OPENAI_BASE_URL = "https://api.z.ai/api/paas/v4" + ZAI_LEGACY_OPENAI_BASE_URL = "https://api.z.ai/api/pass/v4" + + PROVIDERS = { + "openrouter": { + "DISPLAY_NAME": "OpenRouter", + "BASE_URL": "https://openrouter.ai/api/v1", + "BASE_URL_ENV": "HEXSECGPT_OPENROUTER_BASE_URL", + "MODEL_NAME": "deepseek/deepseek-r1-0528:free", # if not work change model run the SeeOpenRouterFreeModels.py see all free model + "MODEL_NAME_ENV": "HEXSECGPT_OPENROUTER_MODEL", + "API_KEY_NAME": "HEXSECGPT_OPENROUTER_API_KEY", + "API_KEY_PROMPT": "Enter your OpenRouter API Key (starts with sk-or-...):", + "DEFAULT_HEADERS": { + "HTTP-Referer": "https://github.com/hexsecteam", + "X-Title": "HexSecGPT-CLI" + }, + }, + "deepseek": { + "DISPLAY_NAME": "DeepSeek", + "BASE_URL": "https://api.deepseek.com", + "BASE_URL_ENV": "HEXSECGPT_DEEPSEEK_BASE_URL", + "MODEL_NAME": "deepseek-chat", + "MODEL_NAME_ENV": "HEXSECGPT_DEEPSEEK_MODEL", + "API_KEY_NAME": "HEXSECGPT_DEEPSEEK_API_KEY", + "API_KEY_PROMPT": "Enter your DeepSeek API Key:", + "DEFAULT_HEADERS": {}, + }, + "custom_openai": { + "DISPLAY_NAME": "Custom OpenAI-Compatible", + "BASE_URL": ZAI_OPENAI_BASE_URL, + "BASE_URL_ENV": "HEXSECGPT_CUSTOM_OPENAI_BASE_URL", + "MODEL_NAME": "glm-5", + "MODEL_NAME_ENV": "HEXSECGPT_CUSTOM_OPENAI_MODEL", + "API_KEY_NAME": "HEXSECGPT_CUSTOM_OPENAI_API_KEY", + "API_KEY_PROMPT": "Enter your API Key:", + "DEFAULT_HEADERS": {}, + }, + } + + # System Paths + ENV_FILE = ".HexSec" + API_KEY_NAME = LEGACY_API_KEY_NAME + + # Visual Theme + CODE_THEME = "monokai" - # API Provider Settings - PROVIDERS = { - "openrouter": { - "BASE_URL": "https://openrouter.ai/api/v1", - "MODEL_NAME": "deepseek/deepseek-r1-0528:free", # if not work change model run the SeeOpenRouterFreeModels.py see all free model - }, - "deepseek": { - "BASE_URL": "https://api.deepseek.com", - "MODEL_NAME": "deepseek-chat", - }, - } - - # Change this if you want to use DeepSeek direct - API_PROVIDER = "openrouter" - - # System Paths - ENV_FILE = ".HexSec" - API_KEY_NAME = "HexSecGPT-API" - - # Visual Theme - CODE_THEME = "monokai" - - class Colors: - USER_PROMPT = "bright_yellow" - - @classmethod - def get_provider_config(cls): - if cls.API_PROVIDER not in cls.PROVIDERS: - return None - return cls.PROVIDERS[cls.API_PROVIDER] + class Colors: + USER_PROMPT = "bright_yellow" + + @classmethod + def get_provider_name(cls): + provider_name = os.getenv(cls.PROVIDER_ENV_NAME, cls.DEFAULT_PROVIDER).strip().lower() + provider_name = provider_name.replace("-", "_").replace(" ", "_") + if provider_name not in cls.PROVIDERS: + return cls.DEFAULT_PROVIDER + return provider_name + + @classmethod + def normalize_provider_name(cls, value: str): + if not value: + return None + normalized = value.strip().lower().replace("-", "_").replace(" ", "_") + aliases = { + "1": "openrouter", + "2": "deepseek", + "3": "custom_openai", + "custom": "custom_openai", + "openai_compatible": "custom_openai", + } + normalized = aliases.get(normalized, normalized) + if normalized not in cls.PROVIDERS: + return None + return normalized + + @classmethod + def get_provider_config(cls, provider_name=None): + provider_name = provider_name or cls.get_provider_name() + provider = cls.PROVIDERS.get(provider_name) + if not provider: + return None + + config = dict(provider) + + base_url_env = config.get("BASE_URL_ENV") + if base_url_env: + config["BASE_URL"] = os.getenv(base_url_env, config["BASE_URL"]).strip() + if ( + provider_name == "custom_openai" + and config["BASE_URL"] == cls.ZAI_LEGACY_OPENAI_BASE_URL + ): + config["BASE_URL"] = cls.ZAI_OPENAI_BASE_URL + + model_name_env = config.get("MODEL_NAME_ENV") + if model_name_env: + config["MODEL_NAME"] = os.getenv(model_name_env, config["MODEL_NAME"]).strip() + + config["PROVIDER_NAME"] = provider_name + return config + + @classmethod + def get_provider_label(cls, provider_name=None): + config = cls.get_provider_config(provider_name) + if not config: + return "Unknown Provider" + return config["DISPLAY_NAME"] + + @classmethod + def get_api_key_name(cls, provider_name=None): + config = cls.get_provider_config(provider_name) + if not config: + return cls.LEGACY_API_KEY_NAME + return config["API_KEY_NAME"] + + @classmethod + def get_api_key(cls, provider_name=None): + api_key_name = cls.get_api_key_name(provider_name) + return os.getenv(api_key_name) or os.getenv(cls.LEGACY_API_KEY_NAME) + + @classmethod + def get_missing_provider_settings(cls, provider_name=None): + config = cls.get_provider_config(provider_name) + if not config: + return [cls.PROVIDER_ENV_NAME] + + missing_settings = [] + + if not cls.get_api_key(config["PROVIDER_NAME"]): + missing_settings.append(config["API_KEY_NAME"]) + + if config.get("BASE_URL_ENV") and not config["BASE_URL"]: + missing_settings.append(config["BASE_URL_ENV"]) + + if config.get("MODEL_NAME_ENV") and not config["MODEL_NAME"]: + missing_settings.append(config["MODEL_NAME_ENV"]) + + return missing_settings # --- UI / TUI Class --- class UI: @@ -268,30 +378,45 @@ class HexSecBrain: Hacker Mode: ENGAGED. """ - def __init__(self, api_key: str, ui: UI): - self.ui = ui - config = Config.get_provider_config() - - if not config: - ui.show_msg("System Error", "Invalid API Provider Configuration", "red") - sys.exit(1) - - self.client = openai.OpenAI( - api_key=api_key, - base_url=config["BASE_URL"], - default_headers={ - "HTTP-Referer": "https://github.com/hexsecteam", - "X-Title": "HexSecGPT-CLI" - } - ) - self.model = config["MODEL_NAME"] - self.history = [{"role": "system", "content": self.SYSTEM_PROMPT}] - - def reset(self): - self.history = [{"role": "system", "content": self.SYSTEM_PROMPT}] - - def chat(self, user_input: str) -> Generator[str, None, None]: - self.history.append({"role": "user", "content": user_input}) + def __init__(self, api_key: str, ui: UI): + self.ui = ui + config = Config.get_provider_config() + + if not config: + ui.show_msg("System Error", "Invalid API Provider Configuration", "red") + sys.exit(1) + + self.provider_name = config["PROVIDER_NAME"] + self.provider_label = config["DISPLAY_NAME"] + self.provider_config = config + self.client = openai.OpenAI( + api_key=api_key, + base_url=config["BASE_URL"], + default_headers=config["DEFAULT_HEADERS"] + ) + self.model = config["MODEL_NAME"] + self.history = [{"role": "system", "content": self.SYSTEM_PROMPT}] + + def reset(self): + self.history = [{"role": "system", "content": self.SYSTEM_PROMPT}] + + def verify_connection(self): + try: + self.client.models.list() + except Exception as model_error: + try: + self.client.chat.completions.create( + model=self.model, + messages=[{"role": "user", "content": "ping"}], + max_tokens=1 + ) + except Exception as completion_error: + raise RuntimeError( + f"Models endpoint failed: {model_error} | Chat verification failed: {completion_error}" + ) from completion_error + + def chat(self, user_input: str) -> Generator[str, None, None]: + self.history.append({"role": "user", "content": user_input}) try: stream = self.client.chat.completions.create( @@ -316,54 +441,141 @@ def chat(self, user_input: str) -> Generator[str, None, None]: yield f"Error: Connection Terminated. Reason: {str(e)}" # --- Main Application --- -class App: - def __init__(self): - self.ui = UI() - self.brain = None - - def setup(self) -> bool: - load_dotenv(dotenv_path=Config.ENV_FILE) - key = os.getenv(Config.API_KEY_NAME) - - if not key: - self.ui.banner() - self.ui.show_msg("Warning", "Encryption Key (API Key) not found.", "yellow") - if self.ui.get_input("Configure now? (y/n)").lower().startswith('y'): - return self.configure_key() - return False - - try: - with self.ui.console.status("[bold green]Verifying Neural Link...[/]"): - self.brain = HexSecBrain(key, self.ui) - self.brain.client.models.list() - time.sleep(1) - return True - except Exception as e: - self.ui.show_msg("Auth Failed", f"Key verification failed: {e}", "red") - if self.ui.get_input("Re-enter key? (y/n)").lower().startswith('y'): - return self.configure_key() - return False - - def configure_key(self) -> bool: - self.ui.banner() - self.ui.console.print("[bold yellow]Enter your API Key (starts with sk-or-...):[/]") - try: - key = pwinput(prompt=f"{colorama.Fore.CYAN}Key > {colorama.Style.RESET_ALL}", mask="*") - except: - key = input("Key > ") - - if not key.strip(): - return False - - set_key(Config.ENV_FILE, Config.API_KEY_NAME, key.strip()) - self.ui.show_msg("Success", "Key saved to encryption ring (.HexSec).", "green") - time.sleep(1) - return self.setup() - - def run_chat(self): - if not self.brain: return - self.ui.banner() - self.ui.show_msg("Connected", "HexSecGPT Uplink Established. Type '/help' for commands.", "green") +class App: + def __init__(self): + self.ui = UI() + self.brain = None + + def prompt_api_key(self): + try: + return pwinput(prompt=f"{colorama.Fore.CYAN}Key > {colorama.Style.RESET_ALL}", mask="*") + except Exception: + return input("Key > ") + + def get_api_key_prompt_text(self, provider_config): + return f"[bold yellow]{provider_config['API_KEY_PROMPT']}[/]" + + def select_provider(self): + current_provider = Config.get_provider_name() + provider_lines = [ + f"Current provider: {Config.get_provider_label(current_provider)} ({current_provider})", + "", + "[1] OpenRouter", + "[2] DeepSeek", + "[3] Custom OpenAI-Compatible", + "", + "Press Enter to keep the current provider.", + ] + self.ui.show_msg("Provider Setup", "\n".join(provider_lines), "cyan") + + while True: + choice = self.ui.get_input("Provider [1-3]").strip() + if not choice: + return current_provider + + provider_name = Config.normalize_provider_name(choice) + if provider_name: + return provider_name + + self.ui.show_msg("Invalid Provider", "Choose 1, 2, 3 or a provider name.", "red") + + def setup(self) -> bool: + load_dotenv(dotenv_path=Config.ENV_FILE, override=True) + provider_name = Config.get_provider_name() + provider_label = Config.get_provider_label(provider_name) + missing_settings = Config.get_missing_provider_settings(provider_name) + key = Config.get_api_key(provider_name) + + if missing_settings: + self.ui.banner() + self.ui.show_msg( + "Warning", + f"{provider_label} configuration is incomplete.\nMissing: {', '.join(missing_settings)}", + "yellow" + ) + if self.ui.get_input("Configure now? (y/n)").lower().startswith('y'): + return self.configure_key() + return False + + try: + with self.ui.console.status("[bold green]Verifying Neural Link...[/]"): + self.brain = HexSecBrain(key, self.ui) + self.brain.verify_connection() + time.sleep(1) + return True + except Exception as e: + self.ui.show_msg("Auth Failed", f"{provider_label} verification failed: {e}", "red") + if self.ui.get_input("Re-enter key? (y/n)").lower().startswith('y'): + return self.configure_key() + return False + + def configure_key(self) -> bool: + load_dotenv(dotenv_path=Config.ENV_FILE, override=True) + self.ui.banner() + + provider_name = self.select_provider() + provider_config = Config.get_provider_config(provider_name) + current_key = (Config.get_api_key(provider_name) or "").strip() + current_model = provider_config["MODEL_NAME"] + current_base_url = provider_config["BASE_URL"] + + set_key(Config.ENV_FILE, Config.PROVIDER_ENV_NAME, provider_name) + os.environ[Config.PROVIDER_ENV_NAME] = provider_name + + if provider_name == "custom_openai": + self.ui.console.print( + "[bold yellow]Enter the OpenAI-compatible base URL " + f"(press Enter to keep: {Config.ZAI_OPENAI_BASE_URL}):[/]" + ) + base_url = self.ui.get_input("Base URL").strip() or current_base_url + if not base_url: + self.ui.show_msg("Missing Setting", "A custom base URL is required.", "red") + return False + set_key(Config.ENV_FILE, provider_config["BASE_URL_ENV"], base_url) + os.environ[provider_config["BASE_URL_ENV"]] = base_url + + self.ui.console.print( + f"[bold yellow]Enter the model for {provider_config['DISPLAY_NAME']} " + f"(press Enter to keep: {current_model or 'none'}):[/]" + ) + model_name = self.ui.get_input("Model").strip() or current_model + if provider_name == "custom_openai" and not model_name: + self.ui.show_msg("Missing Setting", "A custom model name is required.", "red") + return False + if model_name: + set_key(Config.ENV_FILE, provider_config["MODEL_NAME_ENV"], model_name) + os.environ[provider_config["MODEL_NAME_ENV"]] = model_name + + self.ui.console.print(self.get_api_key_prompt_text(provider_config)) + key = self.prompt_api_key().strip() or current_key + + if not key: + self.ui.show_msg("Missing Setting", "An API key is required for the selected provider.", "red") + return False + + set_key(Config.ENV_FILE, provider_config["API_KEY_NAME"], key) + os.environ[provider_config["API_KEY_NAME"]] = key + self.ui.show_msg( + "Success", + f"{provider_config['DISPLAY_NAME']} configuration saved to encryption ring (.HexSec).", + "green" + ) + time.sleep(1) + return self.setup() + + def run_chat(self): + if not self.brain: return + self.ui.banner() + self.ui.show_msg( + "Connected", + ( + "HexSecGPT Uplink Established.\n" + f"Provider: {self.brain.provider_label}\n" + f"Model: {self.brain.model}\n" + "Type '/help' for commands." + ), + "green" + ) while True: try: diff --git a/README.md b/README.md index 9426ca5..e691704 100644 --- a/README.md +++ b/README.md @@ -91,9 +91,10 @@ Follow these steps to get the HexSecGPT framework running on your system. To use this framework, you **must** obtain an API key from a supported provider. These services offer free tiers that are perfect for getting started. -1. **Choose a provider:** - * **OpenRouter:** Visit [OpenRouter.ai](https://openrouter.ai/keys) to get a free API key. They provide access to a variety of models. - * **DeepSeek:** Visit the [DeepSeek Platform](https://platform.deepseek.com/api_keys) for a free API key to use their powerful models. +1. **Choose a provider:** + * **OpenRouter:** Visit [OpenRouter.ai](https://openrouter.ai/keys) to get a free API key. They provide access to a variety of models. + * **DeepSeek:** Visit the [DeepSeek Platform](https://platform.deepseek.com/api_keys) for a free API key to use their powerful models. + * **Custom OpenAI-Compatible API:** Use any provider or self-hosted endpoint that exposes an OpenAI-compatible API. The built-in default base URL is `https://api.z.ai/api/paas/v4` with `glm-5` as the default model, and you can override either value if needed. 2. **Copy your API key.** You will need to paste it into the script when prompted during the first run. @@ -133,23 +134,22 @@ If you prefer to install manually, follow these steps. --- -## :wrench: Configuration - -You can easily switch between API providers. - -1. Open the `HexSecGPT.py` file in a text editor. -2. Locate the `API_PROVIDER` variable at the top of the file. -3. Change the value to either `"openrouter"` or `"deepseek"`. - - ```python - # HexSecGPT.py - - # Change this value to "deepseek" or "openrouter" - API_PROVIDER = "openrouter" - ``` -4. Save the file. The script will now use the selected provider's API. ---- -## 📽️ Demo Setup +## :wrench: Configuration + +You can switch providers directly from the CLI without editing source code. + +1. Run `python HexSecGPT.py`. +2. Open `Configure Security Keys (API Setup)` from the main menu. +3. Select one of these providers: + * `OpenRouter` + * `DeepSeek` + * `Custom OpenAI-Compatible` +4. Enter the requested settings: + * For OpenRouter and DeepSeek: API key, plus an optional model override. + * For Custom OpenAI-Compatible: base URL, model name, and API key. If you press Enter on the base URL prompt, the app keeps `https://api.z.ai/api/paas/v4`, and if you press Enter on the model prompt, it keeps `glm-5`. The custom API key prompt does not require any `sk-` or `sk-or-` prefix. +5. The app stores the selected provider and its credentials in `.HexSec` for future sessions. +--- +## 📽️ Demo Setup ▶️ YouTube Demo: [https://www.youtube.com/watch?v=EM08JC4Mv6c](https://www.youtube.com/watch?v=EM08JC4Mv6c) @@ -161,7 +161,7 @@ Once installation and configuration are complete, run the application with this python3 HexSecGPT.py ``` -The first time you run it, you will be prompted to enter your API key. It will be saved locally for future sessions. +The first time you run it, you will be prompted to configure a provider. Your provider settings and API key are saved locally in `.HexSec` for future sessions. --- @@ -179,14 +179,15 @@ To handle this, the repository includes a model discovery script that helps you 4. The script will list all currently available FREE models 5. Choose one of the working models from the output -### 🔧 Update the provider configuration - -- Navigate to the provider file in the source code HexSecGPT.py -- Replace the existing model name with one of the working free models - -![Provider model configuration example](img/provider-model-example.jpg) - -- Save the file and restart the application +### 🔧 Update the provider configuration + +- Open `Configure Security Keys (API Setup)` from the CLI +- Select `OpenRouter` +- Replace the current model with one of the working free models + +![Provider model configuration example](img/provider-model-example.jpg) + +- Save the configuration and restart the application if needed > ⚠️ Note: Some free models may not work correctly or may be temporarily disabled by OpenRouter. > If a model fails, simply try another one from the list. diff --git a/install.bat b/install.bat index 60e97ca..a360268 100644 --- a/install.bat +++ b/install.bat @@ -49,6 +49,6 @@ echo To run HexSecGPT, run this command in this terminal: echo. echo python HexSecGPT.py echo. -echo Don't forget to get your API key from OpenRouter or DeepSeek! +echo Don't forget to prepare your API key for OpenRouter, DeepSeek, or any custom OpenAI-compatible provider! echo ====================================== pause diff --git a/install.sh b/install.sh index 0752226..a8f69c8 100644 --- a/install.sh +++ b/install.sh @@ -60,5 +60,5 @@ echo "To run HexSecGPT:" echo "1. cd HexSecGPT" echo "2. python3 HexSecGPT.py" echo "" -echo "Don't forget to get your API key from OpenRouter or DeepSeek!" +echo "Don't forget to prepare your API key for OpenRouter, DeepSeek, or any custom OpenAI-compatible provider!" echo "======================================" diff --git a/tests/test_provider_config.py b/tests/test_provider_config.py new file mode 100644 index 0000000..145d222 --- /dev/null +++ b/tests/test_provider_config.py @@ -0,0 +1,209 @@ +import importlib +import os +import sys +import types +import unittest + + +def install_dependency_stubs(): + openai = types.ModuleType("openai") + + class AuthenticationError(Exception): + pass + + class DummyOpenAI: + def __init__(self, *args, **kwargs): + self.models = types.SimpleNamespace(list=lambda: []) + self.chat = types.SimpleNamespace( + completions=types.SimpleNamespace(create=lambda **kwargs: []) + ) + + openai.AuthenticationError = AuthenticationError + openai.OpenAI = DummyOpenAI + sys.modules["openai"] = openai + + colorama = types.ModuleType("colorama") + colorama.Fore = types.SimpleNamespace(CYAN="") + colorama.Style = types.SimpleNamespace(RESET_ALL="") + colorama.init = lambda autoreset=True: None + sys.modules["colorama"] = colorama + + pwinput_module = types.ModuleType("pwinput") + pwinput_module.pwinput = lambda prompt="", mask="*": "" + sys.modules["pwinput"] = pwinput_module + + dotenv = types.ModuleType("dotenv") + dotenv.load_dotenv = lambda dotenv_path=None, override=False: None + dotenv.set_key = lambda path, key, value: None + sys.modules["dotenv"] = dotenv + + rich = types.ModuleType("rich") + sys.modules["rich"] = rich + + console_module = types.ModuleType("rich.console") + + class DummyStatus: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + class DummyConsole: + def print(self, *args, **kwargs): + return None + + def input(self, *args, **kwargs): + return "" + + def status(self, *args, **kwargs): + return DummyStatus() + + console_module.Console = DummyConsole + sys.modules["rich.console"] = console_module + + panel_module = types.ModuleType("rich.panel") + panel_module.Panel = type("Panel", (), {}) + sys.modules["rich.panel"] = panel_module + + markdown_module = types.ModuleType("rich.markdown") + markdown_module.Markdown = type("Markdown", (), {}) + sys.modules["rich.markdown"] = markdown_module + + text_module = types.ModuleType("rich.text") + text_module.Text = type("Text", (), {}) + sys.modules["rich.text"] = text_module + + live_module = types.ModuleType("rich.live") + + class DummyLive: + def __init__(self, *args, **kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def update(self, *args, **kwargs): + return None + + live_module.Live = DummyLive + sys.modules["rich.live"] = live_module + + table_module = types.ModuleType("rich.table") + table_module.Table = type("Table", (), {}) + sys.modules["rich.table"] = table_module + + spinner_module = types.ModuleType("rich.spinner") + spinner_module.Spinner = type("Spinner", (), {}) + sys.modules["rich.spinner"] = spinner_module + + align_module = types.ModuleType("rich.align") + align_module.Align = type("Align", (), {"center": staticmethod(lambda value: value)}) + sys.modules["rich.align"] = align_module + + +def load_hexsec_module(): + install_dependency_stubs() + sys.modules.pop("HexSecGPT", None) + return importlib.import_module("HexSecGPT") + + +class ProviderConfigTests(unittest.TestCase): + def setUp(self): + self.original_env = os.environ.copy() + for key in list(os.environ): + if key.startswith("HEXSECGPT_") or key == "HexSecGPT-API": + os.environ.pop(key, None) + + def tearDown(self): + os.environ.clear() + os.environ.update(self.original_env) + + def test_defaults_to_openrouter_config(self): + module = load_hexsec_module() + + self.assertEqual(module.Config.get_provider_name(), "openrouter") + + config = module.Config.get_provider_config() + self.assertEqual(config["BASE_URL"], "https://openrouter.ai/api/v1") + self.assertEqual(config["MODEL_NAME"], "deepseek/deepseek-r1-0528:free") + self.assertEqual(config["API_KEY_NAME"], "HEXSECGPT_OPENROUTER_API_KEY") + self.assertIn("HTTP-Referer", config["DEFAULT_HEADERS"]) + self.assertIn("X-Title", config["DEFAULT_HEADERS"]) + + def test_custom_openai_provider_reads_runtime_settings(self): + module = load_hexsec_module() + os.environ["HEXSECGPT_PROVIDER"] = "custom_openai" + os.environ["HEXSECGPT_CUSTOM_OPENAI_API_KEY"] = "sk-custom" + os.environ["HEXSECGPT_CUSTOM_OPENAI_BASE_URL"] = "https://example.test/v1" + os.environ["HEXSECGPT_CUSTOM_OPENAI_MODEL"] = "custom-model" + + config = module.Config.get_provider_config() + + self.assertEqual(config["BASE_URL"], "https://example.test/v1") + self.assertEqual(config["MODEL_NAME"], "custom-model") + self.assertEqual(config["API_KEY_NAME"], "HEXSECGPT_CUSTOM_OPENAI_API_KEY") + self.assertEqual(config["DEFAULT_HEADERS"], {}) + + def test_custom_openai_provider_uses_zai_defaults(self): + module = load_hexsec_module() + os.environ["HEXSECGPT_PROVIDER"] = "custom_openai" + + config = module.Config.get_provider_config() + + self.assertEqual(config["BASE_URL"], "https://api.z.ai/api/paas/v4") + self.assertEqual(config["MODEL_NAME"], "glm-5") + + def test_custom_openai_legacy_pass_url_is_normalized(self): + module = load_hexsec_module() + os.environ["HEXSECGPT_PROVIDER"] = "custom_openai" + os.environ["HEXSECGPT_CUSTOM_OPENAI_BASE_URL"] = "https://api.z.ai/api/pass/v4" + + config = module.Config.get_provider_config() + + self.assertEqual(config["BASE_URL"], "https://api.z.ai/api/paas/v4") + + def test_legacy_api_key_fallback_is_preserved(self): + module = load_hexsec_module() + os.environ["HEXSECGPT_PROVIDER"] = "deepseek" + os.environ["HexSecGPT-API"] = "legacy-key" + + self.assertEqual(module.Config.get_api_key_name(), "HEXSECGPT_DEEPSEEK_API_KEY") + self.assertEqual(module.Config.get_api_key(), "legacy-key") + + def test_custom_provider_reports_missing_required_settings(self): + module = load_hexsec_module() + os.environ["HEXSECGPT_PROVIDER"] = "custom_openai" + os.environ["HEXSECGPT_CUSTOM_OPENAI_API_KEY"] = "sk-custom" + + self.assertEqual( + module.Config.get_missing_provider_settings(), + [], + ) + + def test_custom_openai_key_prompt_does_not_require_sk_prefix(self): + module = load_hexsec_module() + app = module.App() + + prompt = app.get_api_key_prompt_text( + module.Config.get_provider_config("custom_openai") + ) + + self.assertEqual(prompt, "[bold yellow]Enter your API Key:[/]") + + def test_openrouter_key_prompt_keeps_provider_specific_hint(self): + module = load_hexsec_module() + app = module.App() + + prompt = app.get_api_key_prompt_text( + module.Config.get_provider_config("openrouter") + ) + + self.assertIn("sk-or-", prompt) + + +if __name__ == "__main__": + unittest.main()