Hass/ha-text-ai
1
0
Fork 0
mirror of https://github.com/smkrv/ha-text-ai.git synced 2026-07-21 22:54:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: v2.5.1 security/ML/reliability patch

Browse Source 
Security:
- H1 DNS rebinding TOCTOU: validate_endpoint returns (endpoint, resolved_ips);
  create_pinned_session builds an aiohttp session with a custom resolver that
  returns only the pre-validated IPs, closing the re-resolve gap.
- H2/M3 Shared session cookie pollution: integration no longer uses HA shared
  clientsession; isolated session with DummyCookieJar prevents cross-integration
  cookie leaks.
- Cloud metadata / link-local block added to allow_local_network mode to
  prevent IMDS exfiltration on cloud VMs.
- L3 hard cap extended to history.async_get_history (not only service schema).

ML/LLM correctness:
- _is_openai_reasoning_model uses regex with gpt-5-chat* blacklist and handles
  OpenRouter-style openai/ prefix. Future o5/gpt-6 auto-recognized.
- reasoning_effort=minimal for gpt-5 family, low for o-series.
- Gemini 2.5 Pro gets thinking_budget=128 (Pro rejects 0, flash accepts 0).
- Anthropic extracts first type=text block instead of hardcoded content[0].
- /no_think dedup uses word-boundary regex instead of substring.
- DeepSeek-reasoner detected: skips /no_think; preserves reasoning_content.

Reliability:
- Exception chaining added to Gemini-block re-raises.
- Top-level imports for json/re/hashlib (no more lazy stdlib imports).
- allow_local_network logged at INFO, not WARNING on every setup.

Code style:
- PEP 604 type hints across all .py files (dict/list/| None).

No breaking changes for users. Internal API: validate_endpoint return type
changed from str to tuple[str, list[str]].
This commit is contained in:
SMKRV
2026-04-17 01:58:16 +03:00
parent e8b9b911ef
commit 6e635f7e2f
11 changed files with 367 additions and 204 deletions
Download Patch File Download Diff File Expand all files Collapse all files
custom_components/ha_text_ai/__init__.py
+12 -10
View File
@@ -9,7 +9,7 @@ The HA Text AI integration.
from __future__ import annotations
import logging
from typing import Any, Dict
from typing import Any
import asyncio
@@ -20,12 +20,11 @@ from homeassistant.const import CONF_API_KEY, CONF_NAME
from homeassistant.core import HomeAssistant, ServiceCall, SupportsResponse
from homeassistant.exceptions import ConfigEntryNotReady, HomeAssistantError
from homeassistant.helpers import config_validation as cv
from homeassistant.helpers import aiohttp_client
from homeassistant.util import dt as dt_util
from .coordinator import HATextAICoordinator
from .api_client import APIClient
from .utils import normalize_name, safe_log_data, validate_endpoint
from .utils import create_pinned_session, normalize_name, safe_log_data, validate_endpoint
from .providers import get_default_endpoint, get_default_model, build_auth_headers
from .const import (
DOMAIN,
@@ -110,7 +109,7 @@ def get_coordinator_by_instance(hass: HomeAssistant, instance: str) -> HATextAIC
raise HomeAssistantError(f"Instance {instance} not found")
async def async_setup(hass: HomeAssistant, config: Dict[str, Any]) -> bool:
async def async_setup(hass: HomeAssistant, config: dict[str, Any]) -> bool:
"""Set up the Home Assistant Text AI component."""
# Initialize domain data storage
hass.data.setdefault(DOMAIN, {})
@@ -274,22 +273,26 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
_LOGGER.error("API provider not specified")
raise ConfigEntryNotReady("API provider is required")
session = aiohttp_client.async_get_clientsession(hass)
model = config.get(CONF_MODEL, get_default_model(api_provider))
raw_endpoint = config.get(CONF_API_ENDPOINT, get_default_endpoint(api_provider))
allow_local = config.get(CONF_ALLOW_LOCAL_NETWORK, DEFAULT_ALLOW_LOCAL_NETWORK)
if allow_local:
_LOGGER.warning(
_LOGGER.info(
"Local network mode enabled for endpoint %s — "
"SSRF protection disabled, API credentials may be sent without TLS",
"SSRF protection relaxed for self-hosted proxies",
raw_endpoint,
)
try:
endpoint = await validate_endpoint(hass, raw_endpoint, allow_local=allow_local)
endpoint, resolved_ips = await validate_endpoint(
hass, raw_endpoint, allow_local=allow_local
)
except ValueError as err:
_LOGGER.error("Invalid API endpoint: %s", err)
raise ConfigEntryNotReady(f"Invalid API endpoint: {err}") from err
# Pinned session closes DNS-rebinding TOCTOU and isolates cookies
# from other integrations sharing the same endpoint hostname.
session = create_pinned_session(hass, endpoint, resolved_ips)
# API key can now be updated via options
api_key = config.get(CONF_API_KEY, entry.data.get(CONF_API_KEY))
instance_name = entry.data.get(CONF_NAME, entry.entry_id)
@@ -363,7 +366,6 @@ async def async_update_options(hass: HomeAssistant, entry: ConfigEntry) -> None:
_LOGGER.info("Options updated for %s, reloading integration", entry.title)
await hass.config_entries.async_reload(entry.entry_id)
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Unload a config entry."""
try:
custom_components/ha_text_ai/api_client.py
+108 -64
View File
@@ -8,9 +8,11 @@ API Client for HA Text AI.
"""
from __future__ import annotations
import logging
import asyncio
from typing import Any, Dict, List, Optional
import json
import logging
import re
from typing import Any
from aiohttp import ClientSession, ClientTimeout
from homeassistant.exceptions import HomeAssistantError
@@ -29,7 +31,6 @@ from .const import (
_LOGGER = logging.getLogger(__name__)
class APIClient:
"""API Client for OpenAI and Anthropic."""
@@ -37,11 +38,11 @@ class APIClient:
self,
session: ClientSession,
endpoint: str,
headers: Dict[str, str],
headers: dict[str, str],
api_provider: str,
model: str,
api_timeout: int = DEFAULT_API_TIMEOUT,
api_key: Optional[str] = None,
api_key: str | None = None,
) -> None:
"""Initialize API client."""
self.session = session
@@ -89,8 +90,8 @@ class APIClient:
async def _make_request(
self,
url: str,
payload: Dict[str, Any],
) -> Dict[str, Any]:
payload: dict[str, Any],
) -> dict[str, Any]:
"""Make API request with retry logic for transient errors only.
Retries on:
@@ -174,7 +175,7 @@ class APIClient:
raise HomeAssistantError("API request failed after all retries")
@staticmethod
def _parse_retry_after(value: Optional[str]) -> Optional[float]:
def _parse_retry_after(value: str | None) -> float | None:
"""Parse Retry-After header (seconds). Caps at 60s to avoid long stalls."""
if not value:
return None
@@ -189,13 +190,13 @@ class APIClient:
async def create(
self,
model: str,
messages: List[Dict[str, str]],
messages: list[dict[str, str]],
temperature: float,
max_tokens: int,
structured_output: bool = False,
json_schema: Optional[str] = None,
json_schema: str | None = None,
disable_thinking: bool = False,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""Create completion using appropriate API."""
try:
self._validate_parameters(temperature, max_tokens)
@@ -224,43 +225,60 @@ class APIClient:
_LOGGER.error("API request failed: %s", str(e))
raise HomeAssistantError(f"API request failed: {str(e)}") from e
@staticmethod
def _is_openai_reasoning_model(model: str) -> bool:
"""Detect OpenAI reasoning models (o-series and GPT-5 family).
# Non-reasoning variants whose names otherwise overlap with the
# reasoning prefix set (e.g. "gpt-5-chat-latest" is classic chat).
_OPENAI_NON_REASONING_PATTERNS: tuple[str, ...] = ("gpt-5-chat",)
_OPENAI_REASONING_REGEX = re.compile(
r"^(?:o\d+|gpt-[5-9](?:\.\d+)?)(?:[-_].*)?$"
)
@classmethod
def _is_openai_reasoning_model(cls, model: str) -> bool:
"""Detect OpenAI reasoning models (o-series and GPT-5+ family).
Reasoning models require max_completion_tokens (not max_tokens),
do not accept custom temperature, and use "developer" role instead
of "system". Cutoff: models released 2025-09 and later are all
reasoning-by-default (o3, o4-mini, gpt-5, gpt-5-mini, gpt-5-nano).
of "system". Uses a regex so future o5/gpt-6 releases are caught
without code change. Explicitly excludes chat-variants
(e.g. gpt-5-chat-latest) which are classic chat models.
"""
if not model:
return False
m = model.lower().lstrip()
# Match bare model names and dated variants (o3-2025-04-16 etc.)
return (
m.startswith(("o1", "o3", "o4-mini", "o4"))
or m.startswith(("gpt-5", "gpt5"))
)
m = model.strip().lower()
# OpenRouter-style "openai/o3" prefix — strip provider namespace.
if "/" in m:
m = m.rsplit("/", 1)[-1]
for non_reasoning in cls._OPENAI_NON_REASONING_PATTERNS:
if m.startswith(non_reasoning):
return False
return bool(cls._OPENAI_REASONING_REGEX.match(m))
@staticmethod
def _convert_system_to_developer(
messages: List[Dict[str, str]],
) -> List[Dict[str, str]]:
messages: list[dict[str, str]],
) -> list[dict[str, str]]:
"""Rename role "system" to "developer" for OpenAI reasoning models."""
return [
{**m, "role": "developer"} if m.get("role") == "system" else m
for m in messages
]
@staticmethod
# Matches /no_think only as a standalone soft-switch token (word-bounded),
# not when users discuss the concept ("discuss /no_think semantics").
_NO_THINK_TOKEN_RE = re.compile(r"(?:^|\s)/no_think(?:\s|$)")
@classmethod
def _apply_no_think_tag(
messages: List[Dict[str, str]],
) -> List[Dict[str, str]]:
cls,
messages: list[dict[str, str]],
) -> list[dict[str, str]]:
"""Append Qwen-style /no_think soft switch to the last user message.
Why: Qwen3 reasoning models treat "/no_think" in the last user turn as a
request to skip thinking. Non-Qwen models ignore the trailing token
harmlessly, so this is safe to apply to all OpenAI-compatible backends.
Uses word-boundary regex for dedup so that user content mentioning
"/no_think" mid-sentence isn't mistaken for an existing soft switch.
"""
if not messages:
return messages
@@ -268,8 +286,8 @@ class APIClient:
for i in range(len(patched) - 1, -1, -1):
if patched[i].get("role") == "user":
content = patched[i].get("content", "")
if "/no_think" not in content:
patched[i]["content"] = f"{content} /no_think".strip()
if not cls._NO_THINK_TOKEN_RE.search(content):
patched[i]["content"] = f"{content.rstrip()} /no_think".lstrip()
break
return patched
@@ -285,7 +303,6 @@ class APIClient:
"""
if not text or "<think>" not in text:
return text
import re
pattern = re.compile(r"<think>.*?</think>", flags=re.DOTALL)
cleaned = text
# Iterative pass: each iteration peels one layer of nested tags.
@@ -303,14 +320,13 @@ class APIClient:
@staticmethod
def _apply_structured_output(
payload: Dict[str, Any],
payload: dict[str, Any],
structured_output: bool,
json_schema: Optional[str],
json_schema: str | None,
) -> None:
"""Apply OpenAI-compatible structured output to payload in-place."""
if not (structured_output and json_schema):
return
import json
try:
schema = json.loads(json_schema)
payload["response_format"] = {
@@ -328,16 +344,28 @@ class APIClient:
async def _create_deepseek_completion(
self,
model: str,
messages: List[Dict[str, str]],
messages: list[dict[str, str]],
temperature: float,
max_tokens: int,
structured_output: bool = False,
json_schema: Optional[str] = None,
json_schema: str | None = None,
disable_thinking: bool = False,
) -> Dict[str, Any]:
"""Create completion using DeepSeek API."""
) -> dict[str, Any]:
"""Create completion using DeepSeek API.
DeepSeek-reasoner (R1) is a reasoning model: it ignores /no_think
(thinking is always on by design) and emits reasoning_content as a
separate field alongside content. We skip the no_think append for
this model and preserve reasoning_content in the response payload
so it's available for logging/debug.
"""
url = f"{self.endpoint}/chat/completions"
final_messages = self._apply_no_think_tag(messages) if disable_thinking else messages
is_reasoner = "reasoner" in model.lower()
final_messages = (
self._apply_no_think_tag(messages)
if (disable_thinking and not is_reasoner)
else messages
)
payload = {
"model": model,
"messages": final_messages,
@@ -348,13 +376,18 @@ class APIClient:
self._apply_structured_output(payload, structured_output, json_schema)
data = await self._make_request(url, payload)
content = data["choices"][0]["message"]["content"]
if disable_thinking:
message = data["choices"][0]["message"]
content = message.get("content", "")
reasoning = message.get("reasoning_content")
if disable_thinking and not is_reasoner:
content = self._strip_think_blocks(content)
return {
"choices": [
{
"message": {"content": content},
"message": {
"content": content,
**({"reasoning_content": reasoning} if reasoning else {}),
},
}
],
"usage": {
@@ -367,13 +400,13 @@ class APIClient:
async def _create_openai_completion(
self,
model: str,
messages: List[Dict[str, str]],
messages: list[dict[str, str]],
temperature: float,
max_tokens: int,
structured_output: bool = False,
json_schema: Optional[str] = None,
json_schema: str | None = None,
disable_thinking: bool = False,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""Create completion using OpenAI API.
Reasoning models (o-series, gpt-5 family) require a different payload
@@ -388,13 +421,16 @@ class APIClient:
if is_reasoning:
prepared_messages = self._convert_system_to_developer(messages)
payload: Dict[str, Any] = {
payload: dict[str, Any] = {
"model": model,
"messages": prepared_messages,
"max_completion_tokens": max_tokens,
}
if disable_thinking:
payload["reasoning_effort"] = "low"
# gpt-5+ supports "minimal" (cheapest, lowest-CoT). o-series
# rejects "minimal" and accepts low/medium/high — fall back to "low".
effort = "minimal" if model.lower().startswith(("gpt-5", "gpt5")) else "low"
payload["reasoning_effort"] = effort
else:
prepared_messages = (
self._apply_no_think_tag(messages) if disable_thinking else messages
@@ -429,13 +465,13 @@ class APIClient:
async def _create_anthropic_completion(
self,
model: str,
messages: List[Dict[str, str]],
messages: list[dict[str, str]],
temperature: float,
max_tokens: int,
structured_output: bool = False,
json_schema: Optional[str] = None,
json_schema: str | None = None,
disable_thinking: bool = False,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""Create completion using Anthropic API."""
url = f"{self.endpoint}/v1/messages"
@@ -455,10 +491,9 @@ class APIClient:
# schema strings (built from templates/webhook data) could otherwise
# break out of the JSON fence and rewrite the system instruction.
if structured_output and json_schema:
import json as _json
try:
_json.loads(json_schema)
except _json.JSONDecodeError as err:
json.loads(json_schema)
except json.JSONDecodeError as err:
_LOGGER.warning(
"Anthropic: invalid JSON schema, ignoring structured_output: %s", err
)
@@ -489,10 +524,14 @@ class APIClient:
payload["system"] = system_prompt
data = await self._make_request(url, payload)
# Anthropic returns text in "content" array; extended thinking arrives
# as a separate thinking-type block (not <think> tags), so no strip
# is required. Extended thinking is opt-in — we simply never request it.
content = data["content"][0]["text"]
# Anthropic returns an array of content blocks; if extended thinking
# is ever enabled the first block may be type="thinking". Find the
# first text-type block instead of hardcoding index [0].
content = ""
for block in data.get("content", []):
if block.get("type") == "text":
content = block.get("text", "")
break
return {
"choices": [
{
@@ -509,13 +548,13 @@ class APIClient:
async def _create_gemini_completion(
self,
model: str,
messages: List[Dict[str, str]],
messages: list[dict[str, str]],
temperature: float,
max_tokens: int,
structured_output: bool = False,
json_schema: Optional[str] = None,
json_schema: str | None = None,
disable_thinking: bool = False,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""Create completion using Gemini API with google-genai library.
Args:
@@ -566,7 +605,6 @@ class APIClient:
parsed_schema = None
if structured_output and json_schema:
try:
import json
parsed_schema = json.loads(json_schema)
_LOGGER.debug("Gemini structured output enabled with schema")
except json.JSONDecodeError as e:
@@ -589,10 +627,14 @@ class APIClient:
config.response_mime_type = "application/json"
config.response_schema = parsed_schema
# Disable thinking for Gemini 2.5+ models (ignored by 2.0 and earlier)
# Disable thinking for Gemini 2.5+ models. Flash variants accept
# thinking_budget=0 (fully off); Pro variants reject 0 and require
# at least 128 tokens of thinking. 2.0 and earlier ignore the field.
if disable_thinking:
m_lower = model.lower()
budget = 128 if "2.5-pro" in m_lower else 0
try:
config.thinking_config = types.ThinkingConfig(thinking_budget=0)
config.thinking_config = types.ThinkingConfig(thinking_budget=budget)
except (AttributeError, TypeError) as err:
_LOGGER.debug(
"ThinkingConfig not supported by this google-genai version: %s", err
@@ -685,10 +727,12 @@ class APIClient:
except ImportError as e:
_LOGGER.error("Google Gemini library not installed: %s", e)
raise HomeAssistantError("Missing dependency: google-genai. Please install it.")
raise HomeAssistantError(
"Missing dependency: google-genai. Please install it."
) from e
except Exception as e:
_LOGGER.error("Gemini API error: %s", e)
raise HomeAssistantError("Gemini API request failed")
raise HomeAssistantError(f"Gemini API request failed: {e}") from e
async def shutdown(self) -> None:
"""Shutdown API client."""
custom_components/ha_text_ai/config_flow.py
+26 -21
View File
@@ -9,14 +9,13 @@ Config flow for HA text AI integration.
from __future__ import annotations
import logging
from typing import Any, Dict, Optional
from typing import Any
import voluptuous as vol
from homeassistant import config_entries
from homeassistant.const import CONF_API_KEY, CONF_NAME
from homeassistant.core import callback
from homeassistant.config_entries import ConfigFlowResult
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.helpers import selector
from .const import (
@@ -61,13 +60,17 @@ from .const import (
)
from homeassistant.util import dt as dt_util
from .utils import normalize_name, safe_log_data, validate_endpoint
from .utils import (
create_pinned_session,
normalize_name,
safe_log_data,
validate_endpoint,
)
from .providers import get_default_endpoint, get_default_model, build_auth_headers
_LOGGER = logging.getLogger(__name__)
def _build_parameter_schema(data: Dict[str, Any]) -> dict:
def _build_parameter_schema(data: dict[str, Any]) -> dict:
"""Build shared parameter schema fields used by both ConfigFlow and OptionsFlow."""
return {
vol.Optional(
@@ -100,7 +103,6 @@ def _build_parameter_schema(data: Dict[str, Any]) -> dict:
): bool,
}
class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
"""Handle a config flow for HA text AI."""
@@ -112,7 +114,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
self._data = {}
self._provider = None
async def async_step_user(self, user_input: Optional[Dict[str, Any]] = None) -> ConfigFlowResult:
async def async_step_user(self, user_input: dict[str, Any] | None = None) -> ConfigFlowResult:
"""Handle the initial step."""
if user_input is None:
return self.async_show_form(
@@ -131,7 +133,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
return await self.async_step_provider()
def _build_provider_schema(
self, data: Optional[Dict[str, Any]] = None
self, data: dict[str, Any] | None = None
) -> vol.Schema:
"""Build provider configuration schema with optional defaults from data."""
defaults = data or {}
@@ -150,7 +152,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
schema_dict.update(_build_parameter_schema(defaults))
return vol.Schema(schema_dict)
async def async_step_provider(self, user_input: Optional[Dict[str, Any]] = None) -> ConfigFlowResult:
async def async_step_provider(self, user_input: dict[str, Any] | None = None) -> ConfigFlowResult:
"""Handle provider configuration step."""
self._errors = {}
@@ -229,7 +231,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
return normalized
async def _async_validate_api(self, user_input: Dict[str, Any]) -> bool:
async def _async_validate_api(self, user_input: dict[str, Any]) -> bool:
"""Validate API connection using provider registry."""
try:
if CONF_API_KEY not in user_input:
@@ -239,7 +241,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
try:
allow_local = user_input.get(CONF_ALLOW_LOCAL_NETWORK, DEFAULT_ALLOW_LOCAL_NETWORK)
endpoint = await validate_endpoint(
endpoint, resolved_ips = await validate_endpoint(
self.hass, user_input[CONF_API_ENDPOINT], allow_local=allow_local
)
except ValueError as err:
@@ -253,7 +255,9 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
return False
return True
session = async_get_clientsession(self.hass)
# Pinned session ensures the reachability check goes to the same
# IP that will later be used by api_client (no DNS rebinding).
session = create_pinned_session(self.hass, endpoint, resolved_ips)
headers = build_auth_headers(self._provider, user_input[CONF_API_KEY])
from .providers import get_provider_config
@@ -274,7 +278,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
self._errors["base"] = "cannot_connect"
return False
async def _create_entry(self, user_input: Dict[str, Any]) -> ConfigFlowResult:
async def _create_entry(self, user_input: dict[str, Any]) -> ConfigFlowResult:
"""Create the config entry with unique_id deduplication."""
instance_name = user_input[CONF_NAME]
normalized_name = normalize_name(instance_name)
@@ -314,7 +318,6 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
"""Get the options flow for this handler."""
return OptionsFlowHandler()
class OptionsFlowHandler(config_entries.OptionsFlow):
"""Handle options flow."""
@@ -326,7 +329,9 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
return False
try:
endpoint = await validate_endpoint(self.hass, endpoint, allow_local=allow_local)
endpoint, resolved_ips = await validate_endpoint(
self.hass, endpoint, allow_local=allow_local
)
except ValueError as err:
_LOGGER.error("Endpoint validation failed: %s", err)
self._errors["base"] = "cannot_connect"
@@ -335,7 +340,7 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
if provider == API_PROVIDER_GEMINI:
return True
session = async_get_clientsession(self.hass)
session = create_pinned_session(self.hass, endpoint, resolved_ips)
headers = build_auth_headers(provider, api_key)
from .providers import get_provider_config
@@ -356,11 +361,11 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
self._errors["base"] = "cannot_connect"
return False
async def async_step_init(self, user_input: Optional[Dict[str, Any]] = None) -> ConfigFlowResult:
async def async_step_init(self, user_input: dict[str, Any] | None = None) -> ConfigFlowResult:
"""Handle provider selection step."""
if not hasattr(self, "_errors"):
self._errors: dict[str, str] = {}
self._selected_provider: Optional[str] = None
self._selected_provider: str | None = None
current_data = {**self.config_entry.data, **self.config_entry.options}
current_provider = current_data.get(CONF_API_PROVIDER, API_PROVIDER_OPENAI)
@@ -386,7 +391,7 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
}
)
async def async_step_settings(self, user_input: Optional[Dict[str, Any]] = None) -> ConfigFlowResult:
async def async_step_settings(self, user_input: dict[str, Any] | None = None) -> ConfigFlowResult:
"""Handle settings configuration step."""
self._errors = {}
current_data = {**self.config_entry.data, **self.config_entry.options}
@@ -477,8 +482,8 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
def _get_settings_schema(
self,
provider: str,
current_data: Dict[str, Any],
user_input: Optional[Dict[str, Any]],
current_data: dict[str, Any],
user_input: dict[str, Any] | None,
default_endpoint: str,
default_model: str,
) -> vol.Schema:
custom_components/ha_text_ai/const.py
+1 -1
View File
@@ -29,7 +29,7 @@ API_PROVIDERS: Final = [
API_PROVIDER_GEMINI
]
VERSION: Final = "2.5.0"
VERSION: Final = "2.5.1"
# Default endpoints
DEFAULT_OPENAI_ENDPOINT: Final = "https://api.openai.com/v1"
custom_components/ha_text_ai/coordinator.py
+24 -25
View File
@@ -12,7 +12,7 @@ import asyncio
import logging
import os
from datetime import timedelta
from typing import Any, Dict, List, Optional
from typing import Any
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
@@ -40,7 +40,6 @@ from .utils import normalize_name
_LOGGER = logging.getLogger(__name__)
class HATextAICoordinator(DataUpdateCoordinator):
"""Home Assistant Text AI Conversation Coordinator."""
@@ -101,9 +100,9 @@ class HATextAICoordinator(DataUpdateCoordinator):
self._is_rate_limited = False
self._is_maintenance = False
self.endpoint_status = "ready"
self._system_prompt: Optional[str] = None
self._system_prompt: str | None = None
self._last_response: Dict[str, Any] = {
self._last_response: dict[str, Any] = {
"timestamp": dt_util.utcnow().isoformat(),
"question": "",
"response": "",
@@ -132,7 +131,7 @@ class HATextAICoordinator(DataUpdateCoordinator):
# Convenience accessors for backward compatibility
# ------------------------------------------------------------------
@property
def _conversation_history(self) -> List[Dict[str, Any]]:
def _conversation_history(self) -> list[dict[str, Any]]:
return self._history.conversation_history
@property
@@ -155,12 +154,12 @@ class HATextAICoordinator(DataUpdateCoordinator):
# Last response
# ------------------------------------------------------------------
@property
def last_response(self) -> Dict[str, Any]:
def last_response(self) -> dict[str, Any]:
"""Get the last response."""
return self._last_response
@last_response.setter
def last_response(self, value: Dict[str, Any]) -> None:
def last_response(self, value: dict[str, Any]) -> None:
self._last_response = value
# ------------------------------------------------------------------
@@ -173,7 +172,7 @@ class HATextAICoordinator(DataUpdateCoordinator):
except Exception as err:
_LOGGER.error("Error updating HA state for %s: %s", self.instance_name, err)
async def _async_update_data(self) -> Dict[str, Any]:
async def _async_update_data(self) -> dict[str, Any]:
"""Update coordinator data."""
try:
current_state = self._get_current_state()
@@ -209,14 +208,14 @@ class HATextAICoordinator(DataUpdateCoordinator):
async def async_ask_question(
self,
question: str,
model: Optional[str] = None,
temperature: Optional[float] = None,
max_tokens: Optional[int] = None,
system_prompt: Optional[str] = None,
context_messages: Optional[int] = None,
model: str | None = None,
temperature: float | None = None,
max_tokens: int | None = None,
system_prompt: str | None = None,
context_messages: int | None = None,
structured_output: bool = False,
json_schema: Optional[str] = None,
disable_thinking: Optional[bool] = None,
json_schema: str | None = None,
disable_thinking: bool | None = None,
) -> dict:
"""Process question with context management."""
if self.client is None:
@@ -279,11 +278,11 @@ class HATextAICoordinator(DataUpdateCoordinator):
self,
question: str,
model: str,
messages: List[Dict[str, str]],
messages: list[dict[str, str]],
temperature: float,
max_tokens: int,
structured_output: bool = False,
json_schema: Optional[str] = None,
json_schema: str | None = None,
disable_thinking: bool = False,
) -> dict:
"""Send request to AI provider and return structured response.
@@ -348,12 +347,12 @@ class HATextAICoordinator(DataUpdateCoordinator):
async def async_get_history(
self,
limit: Optional[int] = None,
filter_model: Optional[str] = None,
start_date: Optional[str] = None,
limit: int | None = None,
filter_model: str | None = None,
start_date: str | None = None,
include_metadata: bool = False,
sort_order: str = "newest",
) -> List[Dict[str, Any]]:
) -> list[dict[str, Any]]:
"""Get conversation history with optional filtering."""
return await self._history.async_get_history(
limit=limit,
@@ -383,7 +382,7 @@ class HATextAICoordinator(DataUpdateCoordinator):
return STATE_ERROR
return STATE_READY
def _get_safe_initial_state(self) -> Dict[str, Any]:
def _get_safe_initial_state(self) -> dict[str, Any]:
return {
"state": STATE_ERROR,
"metrics": {},
@@ -403,7 +402,7 @@ class HATextAICoordinator(DataUpdateCoordinator):
"normalized_name": self.normalized_name,
}
def _get_sanitized_last_response(self) -> Dict[str, Any]:
def _get_sanitized_last_response(self) -> dict[str, Any]:
"""Get sanitized version of last response with truncation."""
response = self.last_response.copy()
@@ -422,7 +421,7 @@ class HATextAICoordinator(DataUpdateCoordinator):
def _calculate_uptime(self) -> float:
return (dt_util.utcnow() - self._start_time).total_seconds()
def _get_truncated_system_prompt(self) -> Optional[str]:
def _get_truncated_system_prompt(self) -> str | None:
if not self._system_prompt:
return None
if len(self._system_prompt) <= 4096:
@@ -430,7 +429,7 @@ class HATextAICoordinator(DataUpdateCoordinator):
return self._system_prompt[:4096] + TRUNCATION_INDICATOR
@staticmethod
def _validate_update_data(data: Dict[str, Any]) -> None:
def _validate_update_data(data: dict[str, Any]) -> None:
for key in ("state", "metrics", "last_response"):
if key not in data:
raise ValueError(f"Missing required key: {key}")
custom_components/ha_text_ai/history.py
+12 -12
View File
@@ -14,7 +14,7 @@ import os
import shutil
import traceback
from datetime import datetime
from typing import Any, Dict, List, Optional
from typing import Any
import aiofiles
@@ -34,7 +34,6 @@ MAX_ARCHIVE_FILES = 3
_LOGGER = logging.getLogger(__name__)
class AsyncFileHandler:
"""Async context manager for file operations."""
@@ -49,7 +48,6 @@ class AsyncFileHandler:
async def __aexit__(self, exc_type, exc_val, exc_tb):
await self.file.close()
def _assert_not_symlink(path: str) -> None:
"""Refuse to operate on a path that resolves to a symlink.
@@ -61,7 +59,6 @@ def _assert_not_symlink(path: str) -> None:
if os.path.islink(path):
raise OSError(f"Refusing to operate on symlink: {path}")
class HistoryManager:
"""Manages conversation history for an instance."""
@@ -84,10 +81,10 @@ class HistoryManager:
history_dir, f"{normalized_name}_history.json"
)
self._max_history_file_size = MAX_HISTORY_FILE_SIZE
self._conversation_history: List[Dict[str, Any]] = []
self._conversation_history: list[dict[str, Any]] = []
@property
def conversation_history(self) -> List[Dict[str, Any]]:
def conversation_history(self) -> list[dict[str, Any]]:
return self._conversation_history
@property
@@ -388,13 +385,13 @@ class HistoryManager:
async def async_get_history(
self,
limit: Optional[int] = None,
filter_model: Optional[str] = None,
start_date: Optional[str] = None,
limit: int | None = None,
filter_model: str | None = None,
start_date: str | None = None,
include_metadata: bool = False,
sort_order: str = "newest",
default_model: str = "",
) -> List[Dict[str, Any]]:
) -> list[dict[str, Any]]:
"""Get conversation history with optional filtering and sorting."""
try:
history = self._conversation_history.copy()
@@ -425,8 +422,11 @@ class HistoryManager:
else:
history.sort(key=lambda x: x.get("timestamp", ""), reverse=True)
# Clamp limit to ABSOLUTE_MAX_HISTORY_SIZE to prevent pathological
# caller requests from producing multi-MB service payloads.
if limit and limit > 0:
history = history[:limit]
effective_limit = min(int(limit), ABSOLUTE_MAX_HISTORY_SIZE)
history = history[:effective_limit]
if include_metadata:
enriched = []
@@ -447,7 +447,7 @@ class HistoryManager:
_LOGGER.error("Error getting history: %s", e)
return []
def get_limited_history(self, max_display: int = 5) -> Dict[str, Any]:
def get_limited_history(self, max_display: int = 5) -> dict[str, Any]:
"""Get limited conversation history for sensor attributes.
Returns last `max_display` entries with truncated text for HA state.
custom_components/ha_text_ai/manifest.json
+1 -1
View File
@@ -15,5 +15,5 @@
"google-genai>=1.16.0,<2.0.0"
],
"single_config_entry": false,
"version": "2.5.0"
"version": "2.5.1"
}
custom_components/ha_text_ai/metrics.py
+8 -9
View File
@@ -13,7 +13,7 @@ import logging
import os
import re
import traceback
from typing import Any, Dict
from typing import Any
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
@@ -21,7 +21,7 @@ from homeassistant.util import dt as dt_util
_LOGGER = logging.getLogger(__name__)
DEFAULT_METRICS: Dict[str, Any] = {
DEFAULT_METRICS: dict[str, Any] = {
"total_tokens": 0,
"prompt_tokens": 0,
"completion_tokens": 0,
@@ -33,7 +33,6 @@ DEFAULT_METRICS: Dict[str, Any] = {
"min_latency": 0,
}
class MetricsManager:
"""Manages performance metrics for an instance."""
@@ -46,10 +45,10 @@ class MetricsManager:
self.hass = hass
self.instance_name = instance_name
self._metrics_file = metrics_file
self._performance_metrics: Dict[str, Any] = DEFAULT_METRICS.copy()
self._performance_metrics: dict[str, Any] = DEFAULT_METRICS.copy()
@property
def metrics(self) -> Dict[str, Any]:
def metrics(self) -> dict[str, Any]:
return self._performance_metrics
async def async_initialize(self) -> None:
@@ -57,7 +56,7 @@ class MetricsManager:
loaded = await self._load_metrics()
self._performance_metrics = loaded or DEFAULT_METRICS.copy()
async def _load_metrics(self) -> Dict[str, Any] | None:
async def _load_metrics(self) -> dict[str, Any] | None:
try:
exists = await self.hass.async_add_executor_job(
os.path.exists, self._metrics_file
@@ -108,7 +107,7 @@ class MetricsManager:
await self._save_metrics()
async def get_current_metrics(self) -> Dict[str, Any]:
async def get_current_metrics(self) -> dict[str, Any]:
"""Get current performance metrics."""
return self._performance_metrics.copy()
@@ -116,7 +115,7 @@ class MetricsManager:
self,
error: Exception,
model: str,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""Record an error in metrics and return error details."""
self._performance_metrics["total_errors"] += 1
self._performance_metrics["failed_requests"] += 1
@@ -146,7 +145,7 @@ class MetricsManager:
if len(error_msg) > 256:
error_msg = error_msg[:256] + "..."
error_details: Dict[str, Any] = {
error_details: dict[str, Any] = {
"timestamp": dt_util.utcnow().isoformat(),
"model": model,
"instance": self.instance_name,
custom_components/ha_text_ai/providers.py
-4
View File
@@ -62,7 +62,6 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
},
}
def get_provider_config(provider: str) -> dict[str, Any]:
"""Get full provider configuration.
@@ -73,17 +72,14 @@ def get_provider_config(provider: str) -> dict[str, Any]:
raise ValueError(f"Unknown API provider: {provider}")
return PROVIDER_REGISTRY[provider]
def get_default_endpoint(provider: str) -> str:
"""Get default API endpoint for a provider."""
return get_provider_config(provider)["default_endpoint"]
def get_default_model(provider: str) -> str:
"""Get default model for a provider."""
return get_provider_config(provider)["default_model"]
def build_auth_headers(provider: str, api_key: str) -> dict[str, str]:
"""Build authentication headers for a provider."""
config = get_provider_config(provider)
custom_components/ha_text_ai/sensor.py
+3 -4
View File
@@ -10,7 +10,7 @@ from __future__ import annotations
import logging
import math
from typing import Any, Dict
from typing import Any
from homeassistant.components.sensor import (
SensorEntity,
SensorEntityDescription,
@@ -79,7 +79,6 @@ _LOGGER = logging.getLogger(__name__)
_ATTR_TEXT_LIMIT = 2048
_ATTR_PROMPT_LIMIT = 512
async def async_setup_entry(
hass: HomeAssistant,
entry: ConfigEntry,
@@ -185,7 +184,7 @@ class HATextAISensor(CoordinatorEntity, SensorEntity):
return None
return value
def _sanitize_attributes(self, attributes: Dict[str, Any]) -> Dict[str, Any]:
def _sanitize_attributes(self, attributes: dict[str, Any]) -> dict[str, Any]:
"""Sanitize all attributes for JSON serialization."""
sanitized = {
key: self._sanitize_value(value)
@@ -231,7 +230,7 @@ class HATextAISensor(CoordinatorEntity, SensorEntity):
return ENTITY_ICON
@property
def extra_state_attributes(self) -> Dict[str, Any]:
def extra_state_attributes(self) -> dict[str, Any]:
"""Return entity specific state attributes."""
if not self.coordinator.data:
return {}
custom_components/ha_text_ai/utils.py
+172 -53
View File
@@ -8,14 +8,22 @@ Utility functions for HA Text AI integration.
"""
from __future__ import annotations
import hashlib
import ipaddress
import logging
import socket
from typing import Any
from urllib.parse import urlparse
import aiohttp
from aiohttp.abc import AbstractResolver
from aiohttp.resolver import DefaultResolver
from homeassistant.const import CONF_API_KEY
from homeassistant.core import HomeAssistant
from homeassistant.helpers.aiohttp_client import async_create_clientsession
_LOGGER = logging.getLogger(__name__)
def normalize_name(name: str) -> str:
"""Normalize name to conform to HA naming convention using underscores.
@@ -27,12 +35,10 @@ def normalize_name(name: str) -> str:
normalized = ''.join(c if c.isalnum() or c == '_' else '_' for c in name)
normalized = '_'.join(filter(None, normalized.split('_'))).lower()
if not normalized:
import hashlib
digest = hashlib.sha256(name.encode("utf-8", errors="replace")).hexdigest()[:8]
normalized = f"instance_{digest}"
return normalized
def safe_log_data(
data: dict[str, Any],
sensitive_keys: tuple[str, ...] = (CONF_API_KEY,),
@@ -40,11 +46,9 @@ def safe_log_data(
"""Filter sensitive keys from data for safe logging."""
return {k: "***" if k in sensitive_keys else v for k, v in data.items()}
class _RestrictedIPError(ValueError):
"""Raised when an IP address is in a restricted range."""
def _check_ip_restricted(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
"""Check if an IP address is in a restricted range."""
return (
@@ -56,14 +60,136 @@ def _check_ip_restricted(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) ->
or addr.is_unspecified
)
def _is_cloud_metadata_or_unsafe(
addr: ipaddress.IPv4Address | ipaddress.IPv6Address,
) -> bool:
"""Block link-local and cloud instance-metadata addresses.
async def validate_endpoint(hass: HomeAssistant, endpoint: str, *, allow_local: bool = False) -> str:
"""Validate API endpoint URL for security.
These must be blocked even in allow_local_network mode:
- IPv4 link-local (169.254.0.0/16) covers AWS/GCP/Azure IMDS 169.254.169.254.
- IPv6 link-local (fe80::/10).
- Multicast and unspecified addresses.
Without this check, a self-hosted HA running on a cloud VM could be
tricked into exfiltrating cloud credentials via IMDS.
"""
return addr.is_multicast or addr.is_unspecified or addr.is_link_local
class _PinnedResolver(AbstractResolver):
"""aiohttp resolver that returns pre-validated IPs for a single hostname.
Why: prevents DNS-rebinding attacks. After validate_endpoint has
confirmed the hostname resolves to a safe IP, we pin that IP in the
resolver used by the aiohttp session. aiohttp then skips its own
DNS lookup on each request and uses the pinned IP, closing the
TOCTOU gap between validation and actual HTTP call.
"""
def __init__(
self,
pinned: dict[str, list[tuple[str, int]]],
fallback: AbstractResolver | None = None,
) -> None:
self._pinned = pinned
self._fallback = fallback or DefaultResolver()
async def resolve(
self,
host: str,
port: int = 0,
family: int = socket.AF_INET,
) -> list[dict[str, Any]]:
entries = self._pinned.get(host.lower())
if entries is None:
return await self._fallback.resolve(host, port, family)
return [
{
"hostname": host,
"host": ip,
"port": port or default_port,
"family": _family_for(ip),
"proto": 0,
"flags": 0,
}
for ip, default_port in entries
]
async def close(self) -> None:
await self._fallback.close()
def _family_for(ip: str) -> int:
"""Return AF_INET or AF_INET6 based on the IP literal."""
try:
return socket.AF_INET6 if ":" in ip else socket.AF_INET
except Exception:
return socket.AF_INET
async def resolve_hostname_ips(
hass: HomeAssistant,
hostname: str,
) -> list[str]:
"""Resolve hostname to all its IPs via the event-loop-safe executor.
Returns a list of IP strings (may contain both IPv4 and IPv6).
Raises ValueError on resolution failure.
"""
try:
addrinfos = await hass.async_add_executor_job(
socket.getaddrinfo, hostname, None
)
except socket.gaierror as err:
raise ValueError(f"Cannot resolve hostname: {hostname}") from err
ips = []
seen: set[str] = set()
for _family, _type, _proto, _canonname, sockaddr in addrinfos:
ip = sockaddr[0]
if ip not in seen:
seen.add(ip)
ips.append(ip)
if not ips:
raise ValueError(f"No IPs for hostname: {hostname}")
return ips
def create_pinned_session(
hass: HomeAssistant,
endpoint: str,
resolved_ips: list[str],
) -> aiohttp.ClientSession:
"""Create an isolated aiohttp session with pinned DNS and no cookie jar.
Addresses two issues at once:
- DNS rebinding: aiohttp will reuse the pinned IPs from validate_endpoint
rather than re-resolving the hostname on each request.
- Cookie pollution: DummyCookieJar prevents cookies from leaking between
this integration and other HA components sharing the same domain.
"""
parsed = urlparse(endpoint)
hostname = (parsed.hostname or "").lower()
port = parsed.port or (443 if parsed.scheme == "https" else 80)
pinned: dict[str, list[tuple[str, int]]] = {
hostname: [(ip, port) for ip in resolved_ips]
}
connector = aiohttp.TCPConnector(resolver=_PinnedResolver(pinned))
return async_create_clientsession(
hass,
cookie_jar=aiohttp.DummyCookieJar(),
connector=connector,
)
async def validate_endpoint(
hass: HomeAssistant,
endpoint: str,
*,
allow_local: bool = False,
) -> tuple[str, list[str]]:
"""Validate API endpoint URL for security and pin resolved IPs.
Ensures HTTPS-only and blocks private/reserved IP ranges (SSRF protection).
When allow_local is True, permits private IPs and HTTP scheme for self-hosted proxies.
Uses async DNS resolution to avoid blocking the event loop.
Returns the validated endpoint stripped of trailing slashes.
Returns: (validated_endpoint_without_trailing_slash, resolved_ips).
The resolved IPs are intended for pinning in aiohttp resolver via
create_pinned_session(), closing the DNS-rebinding TOCTOU gap.
Raises:
ValueError: If the endpoint fails validation.
@@ -81,54 +207,47 @@ async def validate_endpoint(hass: HomeAssistant, endpoint: str, *, allow_local:
if not hostname:
raise ValueError("Invalid endpoint URL: no hostname")
if allow_local:
# Even in local mode, block multicast and unspecified addresses
resolved_ips: list[str] = []
# Collect and check all resolved IPs (or IP literal directly).
def _collect(ips: list[str]) -> None:
for ip in ips:
if ip not in resolved_ips:
resolved_ips.append(ip)
try:
addr = ipaddress.ip_address(hostname)
_is_ip_literal = True
try:
addr = ipaddress.ip_address(hostname)
except ValueError:
_is_ip_literal = False
except ValueError:
addr = None
_is_ip_literal = False
if _is_ip_literal:
if addr.is_multicast or addr.is_unspecified:
raise ValueError("Multicast and unspecified addresses are not allowed")
else:
# Not an IP literal — resolve and check
try:
addrinfos = await hass.async_add_executor_job(
socket.getaddrinfo, hostname, None
)
for family, _type, _proto, _canonname, sockaddr in addrinfos:
resolved = ipaddress.ip_address(sockaddr[0])
if resolved.is_multicast or resolved.is_unspecified:
raise ValueError(
"Hostname resolves to multicast/unspecified address"
)
except socket.gaierror as err:
raise ValueError(f"Cannot resolve hostname: {hostname}") from err
if _is_ip_literal:
_collect([hostname])
else:
# Full SSRF protection — block private/reserved IPs
try:
addr = ipaddress.ip_address(hostname)
if _check_ip_restricted(addr):
raise _RestrictedIPError("Private/reserved IP addresses are not allowed")
except _RestrictedIPError:
raise
except ValueError:
# Not an IP literal — resolve hostname and check all resolved IPs
# to prevent DNS rebinding attacks
try:
addrinfos = await hass.async_add_executor_job(
socket.getaddrinfo, hostname, None
)
for family, _type, _proto, _canonname, sockaddr in addrinfos:
ip_str = sockaddr[0]
resolved_addr = ipaddress.ip_address(ip_str)
if _check_ip_restricted(resolved_addr):
raise ValueError(
"Hostname resolves to a restricted IP range"
)
except socket.gaierror as err:
raise ValueError(f"Cannot resolve hostname: {hostname}") from err
addrinfos = await hass.async_add_executor_job(
socket.getaddrinfo, hostname, None
)
except socket.gaierror as err:
raise ValueError(f"Cannot resolve hostname: {hostname}") from err
_collect([sockaddr[0] for (*_, sockaddr) in addrinfos])
return endpoint.rstrip("/")
if not resolved_ips:
raise ValueError(f"No IPs resolved for hostname: {hostname}")
# Validate each resolved IP against the selected policy.
for ip_str in resolved_ips:
ip_obj = ipaddress.ip_address(ip_str)
if allow_local:
if _is_cloud_metadata_or_unsafe(ip_obj):
raise ValueError(
"Link-local/metadata/multicast addresses are not allowed"
)
else:
if _check_ip_restricted(ip_obj):
raise _RestrictedIPError(
"Private/reserved IP addresses are not allowed"
)
return endpoint.rstrip("/"), resolved_ips