Files
homeassistant/config/custom_components/heweather/config_flow.py
T
2026-08-24 21:55:33 +08:00

708 lines
27 KiB
Python

import logging
import asyncio
from typing import Optional
import voluptuous as vol
from homeassistant import config_entries
from homeassistant.core import callback
from homeassistant.helpers import entity_registry as er
from .heweather.const import (
DOMAIN,
DEFAULT_NAME,
CONF_AUTH_METHOD,
CONF_LONGITUDE,
CONF_LATITUDE,
CONF_HOST,
CONF_KEY,
CONF_STORAGE_PATH,
CONF_JWT_SUB,
CONF_JWT_KID,
CONF_DISASTERLEVEL,
CONF_DISASTERMSG,
CONF_SENSOR_LIST,
DEFAULT_HOST,
DEFAULT_AUTH_METHOD,
AUTH_METHOD,
DEFAULT_DISASTER_LEVEL_CONF,
DISASTER_LEVEL_CONF,
DEFAULT_DISASTER_MSG,
DISASTER_MSG
)
from .heweather.heweather_cert import HeWeatherCert
_LOGGER = logging.getLogger(__name__)
async def migrate_entities_for_location_change(hass, config_entry, old_longitude=None, old_latitude=None, new_longitude=None, new_latitude=None):
"""迁移实体到新的经纬度,保持实体ID和配置"""
if old_longitude is None or old_latitude is None or new_longitude is None or new_latitude is None:
return
entity_registry = er.async_get(hass)
# 获取所有与此集成相关的旧实体
entities_to_migrate = []
old_unique_id_suffix = f"_{old_longitude}_{old_latitude}"
new_unique_id_suffix = f"_{new_longitude}_{new_latitude}"
for entity_id, entity in entity_registry.entities.items():
if entity.config_entry_id == config_entry.entry_id and entity.platform == DOMAIN:
if entity.unique_id and old_unique_id_suffix in entity.unique_id:
entities_to_migrate.append((entity_id, entity))
# 迁移实体:更新unique_id而不删除实体
migrated_count = 0
for entity_id, entity in entities_to_migrate:
# 生成新的unique_id
new_unique_id = entity.unique_id.replace(old_unique_id_suffix, new_unique_id_suffix)
try:
# 更新实体的unique_id
entity_registry.async_update_entity(
entity_id,
new_unique_id=new_unique_id
)
_LOGGER.info(f"Migrated entity {entity_id}: {entity.unique_id} -> {new_unique_id}")
migrated_count += 1
except Exception as e:
_LOGGER.error(f"Failed to migrate entity {entity_id}: {e}")
# 如果迁移失败,则删除旧实体让系统重新创建
entity_registry.async_remove(entity_id)
_LOGGER.info(f"Removed entity {entity_id} due to migration failure")
if migrated_count > 0:
_LOGGER.info(f"Successfully migrated {migrated_count} entities to new location")
async def cleanup_old_entities(hass, config_entry, old_longitude=None, old_latitude=None):
"""清理旧的实体,当经纬度发生变化时(备用方法)"""
if old_longitude is None or old_latitude is None:
return
entity_registry = er.async_get(hass)
# 获取所有与此集成相关的实体
entities_to_remove = []
for entity_id, entity in entity_registry.entities.items():
if entity.config_entry_id == config_entry.entry_id and entity.platform == DOMAIN:
# 检查是否是旧的实体(基于旧的经纬度)
old_unique_id_suffix = f"_{old_longitude}_{old_latitude}"
if entity.unique_id and old_unique_id_suffix in entity.unique_id:
entities_to_remove.append(entity_id)
# 删除旧实体
for entity_id in entities_to_remove:
_LOGGER.info(f"Removing old entity: {entity_id}")
entity_registry.async_remove(entity_id)
if entities_to_remove:
_LOGGER.info(f"Cleaned up {len(entities_to_remove)} old entities")
async def cleanup_duplicate_entities(hass, config_entry):
"""清理重复的实体,保留最新的"""
entity_registry = er.async_get(hass)
# 获取所有与此集成相关的实体,按传感器类型分组
entities_by_type = {}
for entity_id, entity in entity_registry.entities.items():
if entity.config_entry_id == config_entry.entry_id and entity.platform == DOMAIN:
if entity.unique_id:
# 提取传感器类型(unique_id的第一部分)
sensor_type = entity.unique_id.split('_')[0] + '_' + entity.unique_id.split('_')[1]
if sensor_type not in entities_by_type:
entities_by_type[sensor_type] = []
entities_by_type[sensor_type].append((entity_id, entity))
# 检查每种传感器类型是否有重复
entities_to_remove = []
for sensor_type, entities in entities_by_type.items():
if len(entities) > 1:
# 有重复实体,保留最新的(通过entity_id排序,通常较新的ID会更大)
entities.sort(key=lambda x: x[0]) # 按entity_id排序
# 删除除最后一个之外的所有实体
for entity_id, entity in entities[:-1]:
entities_to_remove.append(entity_id)
_LOGGER.info(f"Found duplicate entity for {sensor_type}: {entity_id}")
# 删除重复实体
for entity_id in entities_to_remove:
_LOGGER.info(f"Removing duplicate entity: {entity_id}")
entity_registry.async_remove(entity_id)
if entities_to_remove:
_LOGGER.info(f"Cleaned up {len(entities_to_remove)} duplicate entities")
def validate_longitude(lon: str) -> bool:
try:
lon_float = float(lon)
return -180 <= lon_float <= 180
except ValueError:
return False
def validate_latitude(lat: str) -> bool:
try:
lat_float = float(lat)
return -90 <= lat_float <= 90
except ValueError:
return False
class HeWeatherConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
VERSION = 1
MINOR_VERSION = 1
SUPPORT_MULTIPLE_ENTRIES = True
@staticmethod
@callback
def async_get_options_flow(config_entry):
"""Get the options flow for this handler."""
return HeWeatherOptionsFlow(config_entry)
_main_loop: asyncio.AbstractEventLoop
_heweather_cert: HeWeatherCert
_storage_path: str
_auth_method: str
_host: str
_key: str
_jwt_pubkey: str
_jwt_sub: str
_jwt_kid: str
_longitude: str
_latitude: str
_disasterlevel: str
_disastermsg: str
def __init__(self):
self._main_loop = asyncio.get_running_loop()
self._storage_path = ''
self._auth_method = DEFAULT_AUTH_METHOD
self._host = DEFAULT_HOST
self._key = ''
self._jwt_pubkey = ''
self._jwt_sub = ''
self._jwt_kid = ''
self._longitude = ''
self._latitude = ''
self._disasterlevel = DEFAULT_DISASTER_LEVEL_CONF
self._disastermsg = DEFAULT_DISASTER_MSG
async def async_step_user(
self, user_input: Optional[dict] = None
):
#if self._async_current_entries():
# return self.async_abort(reason="single_instance_allowed")
self.hass.data.setdefault(DOMAIN, {})
if not self._storage_path:
self._storage_path = self.hass.config.path('.storage', DOMAIN)
# HeWeather Certification
self._heweather_cert = self.hass.data[DOMAIN].get('heweather_cert', None)
if not self._heweather_cert:
self._heweather_cert = HeWeatherCert(
root_path=self._storage_path,
loop=self._main_loop)
self.hass.data[DOMAIN]['heweather_cert'] = self._heweather_cert
_LOGGER.info(
'async_step_user, create heweather cert, %s', self._storage_path)
return await self.async_step_auth_method_config(user_input)
async def async_step_auth_method_config(
self, user_input: Optional[dict] = None
):
if user_input:
self._auth_method = user_input.get("auth_method", self._auth_method)
if self._auth_method == "key":
return await self.async_step_auth_apikey_config()
else:
return await self.async_step_auth_jwt_config()
return await self.__show_auth_method_config_form("")
async def __show_auth_method_config_form(self, reason: str):
return self.async_show_form(
step_id="auth_method_config",
data_schema=vol.Schema({
vol.Required(
"auth_method",
default=self._auth_method
): vol.In(AUTH_METHOD)
}),
errors={'base': reason},
last_step=False
)
async def async_step_auth_apikey_config(
self, user_input: Optional[dict] = None
):
if user_input:
if user_input.get("key", self._key) == "":
return await self.__show_auth_apikey_config_form("key is empty")
elif user_input.get("host", None) == "":
return await self.__show_auth_apikey_config_form("host is empty")
else:
self._key = user_input.get("key", self._key)
self._host = user_input.get("host", self._host)
return await self.async_step_location_config()
return await self.__show_auth_apikey_config_form("")
async def __show_auth_apikey_config_form(self, reason: str):
return self.async_show_form(
step_id="auth_apikey_config",
data_schema=vol.Schema({
vol.Required(
"key",
default=self._key
): str,
vol.Required(
"host",
default=self._host
): str
}),
errors={'base': reason},
last_step=False
)
async def async_step_auth_jwt_config(
self, user_input: Optional[dict] = None
):
if user_input:
if user_input.get("jwt_sub", self._key) == "":
return await self.__show_auth_apikey_config_form("jwt_sub is empty")
elif user_input.get("jwt_kid", None) == "":
return await self.__show_auth_apikey_config_form("jwt_kid is empty")
elif user_input.get("host", None) == "":
return await self.__show_auth_apikey_config_form("host is empty")
else:
self._jwt_sub = user_input.get("jwt_sub", self._jwt_sub)
self._jwt_kid = user_input.get("jwt_kid", self._jwt_kid)
self._host = user_input.get("host", self._host)
return await self.async_step_location_config()
await self._heweather_cert.gen_key_async()
self._jwt_pubkey = await self._heweather_cert.get_pub_key_async()
return await self.__show_auth_jwt_config_form("")
async def __show_auth_jwt_config_form(self, reason: str):
return self.async_show_form(
step_id="auth_jwt_config",
data_schema=vol.Schema({
vol.Required(
"jwt_sub",
default=self._jwt_sub
): str,
vol.Required(
"jwt_kid",
default=self._jwt_kid
): str,
vol.Required(
"host",
default=self._host
): str
}),
description_placeholders={
"jwt_pubkey": self._jwt_pubkey,
},
errors={'base': reason},
last_step=False
)
async def async_step_location_config(
self, user_input: Optional[dict] = None
):
if user_input:
longitude = user_input.get("longitude", "")
latitude = user_input.get("latitude", "")
if not longitude or not latitude:
return await self.__show_location_config_form("empty_location")
if not validate_longitude(longitude):
return await self.__show_location_config_form("invalid_longitude")
if not validate_latitude(latitude):
return await self.__show_location_config_form("invalid_latitude")
self._longitude = f"{float(longitude):.4f}"
self._latitude = f"{float(latitude):.4f}"
return await self.async_step_disaster_config()
return await self.__show_location_config_form("")
async def __show_location_config_form(self, reason: str):
return self.async_show_form(
step_id="location_config",
data_schema=vol.Schema({
vol.Required(
"longitude",
default=self._longitude
): str,
vol.Required(
"latitude",
default=self._latitude
): str,
}),
errors={'base': reason},
last_step=False
)
async def async_step_disaster_config(
self, user_input: Optional[dict] = None
):
if user_input:
self._disasterlevel = user_input.get("disasterlevel", self._disasterlevel)
self._disastermsg = user_input.get("disastermsg", self._disastermsg)
return await self.config_flow_done()
return await self.__show_disaster_config_form("")
async def __show_disaster_config_form(self, reason: str):
return self.async_show_form(
step_id="disaster_config",
data_schema=vol.Schema({
vol.Required(
"disasterlevel",
default=self._disasterlevel
): vol.In(DISASTER_LEVEL_CONF),
vol.Required(
"disastermsg",
default=self._disastermsg
): vol.In(DISASTER_MSG),
}),
errors={'base': reason},
last_step=False
)
async def config_flow_done(self):
return self.async_create_entry(
title=DEFAULT_NAME,
data={
CONF_AUTH_METHOD: self._auth_method,
CONF_KEY: self._key,
CONF_STORAGE_PATH: self._storage_path,
CONF_JWT_SUB: self._jwt_sub,
CONF_JWT_KID: self._jwt_kid,
CONF_HOST: self._host,
CONF_LONGITUDE: self._longitude,
CONF_LATITUDE: self._latitude,
CONF_DISASTERLEVEL: self._disasterlevel,
CONF_DISASTERMSG: self._disastermsg,
})
class HeWeatherOptionsFlow(config_entries.OptionsFlow):
"""Handle options flow for HeWeather."""
def __init__(self, config_entry: config_entries.ConfigEntry) -> None:
"""Initialize options flow."""
self._config_entry = config_entry
self._auth_method = config_entry.data.get(CONF_AUTH_METHOD, DEFAULT_AUTH_METHOD)
self._jwt_pubkey = ""
self._heweather_cert = None
# Initialize variables for different auth methods
self._key = config_entry.data.get(CONF_KEY, "")
self._host = config_entry.data.get(CONF_HOST, DEFAULT_HOST)
self._jwt_sub = config_entry.data.get(CONF_JWT_SUB, "")
self._jwt_kid = config_entry.data.get(CONF_JWT_KID, "")
# Initialize location and disaster config
self._longitude = config_entry.data.get(CONF_LONGITUDE, "")
self._latitude = config_entry.data.get(CONF_LATITUDE, "")
self._disasterlevel = config_entry.data.get(CONF_DISASTERLEVEL, DEFAULT_DISASTER_LEVEL_CONF)
self._disastermsg = config_entry.data.get(CONF_DISASTERMSG, DEFAULT_DISASTER_MSG)
async def async_step_init(self, user_input: Optional[dict] = None):
"""Manage the options."""
return await self.async_step_auth_method_config()
async def async_step_auth_method_config(self, user_input: Optional[dict] = None):
"""Handle authentication method configuration."""
if user_input is not None:
self._auth_method = user_input.get("auth_method", self._auth_method)
if self._auth_method == "key":
return await self.async_step_auth_apikey_config()
else:
return await self.async_step_auth_jwt_config()
return self.async_show_form(
step_id="auth_method_config",
data_schema=vol.Schema({
vol.Required(
"auth_method",
default=self._auth_method
): vol.In(AUTH_METHOD)
}),
)
async def async_step_auth_apikey_config(self, user_input: Optional[dict] = None):
"""Handle API key authentication configuration."""
if user_input is not None:
# Validate inputs
if not user_input.get("key", "").strip():
return self.async_show_form(
step_id="auth_apikey_config",
data_schema=self._get_apikey_schema(),
errors={"base": "key is empty"}
)
if not user_input.get("host", "").strip():
return self.async_show_form(
step_id="auth_apikey_config",
data_schema=self._get_apikey_schema(),
errors={"base": "host is empty"}
)
# Store auth data and proceed to location config
self._key = user_input.get("key", "")
self._host = user_input.get("host", DEFAULT_HOST)
return await self.async_step_location_config()
return self.async_show_form(
step_id="auth_apikey_config",
data_schema=self._get_apikey_schema(),
)
def _get_apikey_schema(self):
"""Get API key configuration schema."""
current_key = self._config_entry.data.get(CONF_KEY, "")
current_host = self._config_entry.data.get(CONF_HOST, DEFAULT_HOST)
return vol.Schema({
vol.Required(
"key",
default=current_key
): str,
vol.Required(
"host",
default=current_host
): str
})
async def async_step_auth_jwt_config(self, user_input: Optional[dict] = None):
"""Handle JWT authentication configuration."""
if user_input is not None:
# Validate inputs
if not user_input.get("jwt_sub", "").strip():
return self.async_show_form(
step_id="auth_jwt_config",
data_schema=self._get_jwt_schema(),
description_placeholders={"jwt_pubkey": self._jwt_pubkey},
errors={"base": "jwt_sub is empty"}
)
if not user_input.get("jwt_kid", "").strip():
return self.async_show_form(
step_id="auth_jwt_config",
data_schema=self._get_jwt_schema(),
description_placeholders={"jwt_pubkey": self._jwt_pubkey},
errors={"base": "jwt_kid is empty"}
)
if not user_input.get("host", "").strip():
return self.async_show_form(
step_id="auth_jwt_config",
data_schema=self._get_jwt_schema(),
description_placeholders={"jwt_pubkey": self._jwt_pubkey},
errors={"base": "host is empty"}
)
# Store auth data and proceed to location config
self._jwt_sub = user_input.get("jwt_sub", "")
self._jwt_kid = user_input.get("jwt_kid", "")
self._host = user_input.get("host", DEFAULT_HOST)
return await self.async_step_location_config()
# Get HeWeather cert instance and public key
if not self._heweather_cert:
# Get the cert instance from hass data
storage_path = self._config_entry.data.get(CONF_STORAGE_PATH)
if storage_path:
self.hass.data.setdefault(DOMAIN, {})
self._heweather_cert = self.hass.data[DOMAIN].get('heweather_cert', None)
if not self._heweather_cert:
import asyncio
loop = asyncio.get_running_loop()
self._heweather_cert = HeWeatherCert(root_path=storage_path, loop=loop)
self.hass.data[DOMAIN]['heweather_cert'] = self._heweather_cert
# Get existing public key (don't regenerate)
if self._heweather_cert:
self._jwt_pubkey = await self._heweather_cert.get_pub_key_async()
# Only generate new key if no existing key found
if not self._jwt_pubkey:
await self._heweather_cert.gen_key_async()
self._jwt_pubkey = await self._heweather_cert.get_pub_key_async()
return self.async_show_form(
step_id="auth_jwt_config",
data_schema=self._get_jwt_schema(),
description_placeholders={"jwt_pubkey": self._jwt_pubkey},
)
def _get_jwt_schema(self):
"""Get JWT configuration schema."""
current_jwt_sub = self._config_entry.data.get(CONF_JWT_SUB, "")
current_jwt_kid = self._config_entry.data.get(CONF_JWT_KID, "")
current_host = self._config_entry.data.get(CONF_HOST, DEFAULT_HOST)
return vol.Schema({
vol.Required(
"jwt_sub",
default=current_jwt_sub
): str,
vol.Required(
"jwt_kid",
default=current_jwt_kid
): str,
vol.Required(
"host",
default=current_host
): str
})
async def async_step_location_config(self, user_input: Optional[dict] = None):
"""Handle location configuration."""
if user_input is not None:
longitude = user_input.get("longitude", "")
latitude = user_input.get("latitude", "")
# Validate location inputs
if not longitude or not latitude:
return self.async_show_form(
step_id="location_config",
data_schema=self._get_location_schema(),
errors={"base": "empty_location"}
)
if not validate_longitude(longitude):
return self.async_show_form(
step_id="location_config",
data_schema=self._get_location_schema(),
errors={"base": "invalid_longitude"}
)
if not validate_latitude(latitude):
return self.async_show_form(
step_id="location_config",
data_schema=self._get_location_schema(),
errors={"base": "invalid_latitude"}
)
self._longitude = f"{float(longitude):.4f}"
self._latitude = f"{float(latitude):.4f}"
return await self.async_step_disaster_config()
return self.async_show_form(
step_id="location_config",
data_schema=self._get_location_schema(),
)
def _get_location_schema(self):
"""Get location configuration schema."""
current_longitude = self._config_entry.data.get(CONF_LONGITUDE, "")
current_latitude = self._config_entry.data.get(CONF_LATITUDE, "")
return vol.Schema({
vol.Required(
"longitude",
default=current_longitude
): str,
vol.Required(
"latitude",
default=current_latitude
): str,
})
async def async_step_disaster_config(self, user_input: Optional[dict] = None):
"""Handle disaster configuration."""
if user_input is not None:
self._disasterlevel = user_input.get("disasterlevel", DEFAULT_DISASTER_LEVEL_CONF)
self._disastermsg = user_input.get("disastermsg", DEFAULT_DISASTER_MSG)
return await self.options_flow_done()
return self.async_show_form(
step_id="disaster_config",
data_schema=self._get_disaster_schema(),
)
def _get_disaster_schema(self):
"""Get disaster configuration schema."""
current_disasterlevel = self._config_entry.data.get(CONF_DISASTERLEVEL, DEFAULT_DISASTER_LEVEL_CONF)
current_disastermsg = self._config_entry.data.get(CONF_DISASTERMSG, DEFAULT_DISASTER_MSG)
return vol.Schema({
vol.Required(
"disasterlevel",
default=current_disasterlevel
): vol.In(DISASTER_LEVEL_CONF),
vol.Required(
"disastermsg",
default=current_disastermsg
): vol.In(DISASTER_MSG),
})
async def options_flow_done(self):
"""Create the updated config entry."""
# Get old location data for cleanup
old_longitude = self._config_entry.data.get(CONF_LONGITUDE)
old_latitude = self._config_entry.data.get(CONF_LATITUDE)
# Check if location has changed
location_changed = (
old_longitude != self._longitude or
old_latitude != self._latitude
)
# Build the updated data based on auth method
# Preserve existing auth data when switching methods
updated_data = {
CONF_AUTH_METHOD: self._auth_method,
CONF_STORAGE_PATH: self._config_entry.data.get(CONF_STORAGE_PATH),
CONF_LONGITUDE: self._longitude,
CONF_LATITUDE: self._latitude,
CONF_DISASTERLEVEL: self._disasterlevel,
CONF_DISASTERMSG: self._disastermsg,
}
if self._auth_method == "key":
# Update API key settings, preserve existing JWT settings
updated_data.update({
CONF_KEY: self._key,
CONF_HOST: self._host,
CONF_JWT_SUB: self._config_entry.data.get(CONF_JWT_SUB, ""),
CONF_JWT_KID: self._config_entry.data.get(CONF_JWT_KID, ""),
})
else:
# Update JWT settings, preserve existing API key settings
updated_data.update({
CONF_KEY: self._config_entry.data.get(CONF_KEY, ""),
CONF_HOST: self._host,
CONF_JWT_SUB: self._jwt_sub,
CONF_JWT_KID: self._jwt_kid,
})
# Migrate entities if location changed (更平滑的实体迁移)
if location_changed and old_longitude and old_latitude:
_LOGGER.info(f"Location changed from ({old_longitude}, {old_latitude}) to ({self._longitude}, {self._latitude})")
# 先尝试迁移实体
await migrate_entities_for_location_change(
self.hass, self._config_entry,
old_longitude, old_latitude,
self._longitude, self._latitude
)
# Always check for and clean up duplicate entities
await cleanup_duplicate_entities(self.hass, self._config_entry)
# Update the config entry data directly
self.hass.config_entries.async_update_entry(
self._config_entry, data=updated_data
)
# For options flow, we should reload the entry to apply changes
# 重新加载会创建新实体(如果迁移失败的话)
await self.hass.config_entries.async_reload(self._config_entry.entry_id)
# Return empty result to indicate successful update
return self.async_create_entry(title="", data={})