Extend safe mode (#31927)

* Extend safe mode

* Add safe mode boolean to config JSON output and default Lovelace

* Add safe mode to frontend

* Update accent color

homeassistant/bootstrap.py

@@ -1,5 +1,6 @@
 """Provide methods to bootstrap a Home Assistant instance."""
 import asyncio
+import contextlib
 import logging
 import logging.handlers
 import os
@@ -7,12 +8,14 @@ import sys
 from time import monotonic
 from typing import Any, Dict, Optional, Set
 
+from async_timeout import timeout
 import voluptuous as vol
 
 from homeassistant import config as conf_util, config_entries, core, loader
 from homeassistant.components import http
 from homeassistant.const import (
     EVENT_HOMEASSISTANT_CLOSE,
+    EVENT_HOMEASSISTANT_STOP,
     REQUIRED_NEXT_PYTHON_DATE,
     REQUIRED_NEXT_PYTHON_VER,
 )
@@ -80,8 +83,7 @@ async def async_setup_hass(
             config_dict = await conf_util.async_hass_config_yaml(hass)
         except HomeAssistantError as err:
             _LOGGER.error(
-                "Failed to parse configuration.yaml: %s. Falling back to safe mode",
+                "Failed to parse configuration.yaml: %s. Activating safe mode", err,
-                err,
             )
         else:
             if not is_virtual_env():
@@ -93,8 +95,30 @@ async def async_setup_hass(
         finally:
             clear_secret_cache()
 
-    if safe_mode or config_dict is None or not basic_setup_success:
+    if config_dict is None:
+        safe_mode = True
+
+    elif not basic_setup_success:
+        _LOGGER.warning("Unable to set up core integrations. Activating safe mode")
+        safe_mode = True
+
+    elif "frontend" not in hass.config.components:
+        _LOGGER.warning("Detected that frontend did not load. Activating safe mode")
+        # Ask integrations to shut down. It's messy but we can't
+        # do a clean stop without knowing what is broken
+        hass.async_track_tasks()
+        hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP, {})
+        with contextlib.suppress(asyncio.TimeoutError):
+            async with timeout(10):
+                await hass.async_block_till_done()
+
+        safe_mode = True
+        hass = core.HomeAssistant()
+        hass.config.config_dir = config_dir
+
+    if safe_mode:
         _LOGGER.info("Starting in safe mode")
+        hass.config.safe_mode = True
 
         http_conf = (await http.async_get_last_config(hass)) or {}
 
@@ -283,7 +307,7 @@ def _get_domains(hass: core.HomeAssistant, config: Dict[str, Any]) -> Set[str]:
     domains = set(key.split(" ")[0] for key in config.keys() if key != core.DOMAIN)
 
     # Add config entry domains
-    if "safe_mode" not in config:
+    if not hass.config.safe_mode:
         domains.update(hass.config_entries.async_domains())
 
     # Make sure the Hass.io component is loaded

homeassistant/components/frontend/__init__.py

@@ -508,6 +508,23 @@ def websocket_get_themes(hass, connection, msg):
 
     Async friendly.
     """
+    if hass.config.safe_mode:
+        connection.send_message(
+            websocket_api.result_message(
+                msg["id"],
+                {
+                    "themes": {
+                        "safe_mode": {
+                            "primary-color": "#db4437",
+                            "accent-color": "#eeee02",
+                        }
+                    },
+                    "default_theme": "safe_mode",
+                },
+            )
+        )
+        return
+
     connection.send_message(
         websocket_api.result_message(
             msg["id"],

homeassistant/components/lovelace/__init__.py

@@ -89,6 +89,9 @@ class LovelaceStorage:
 
     async def async_load(self, force):
         """Load config."""
+        if self._hass.config.safe_mode:
+            raise ConfigNotFound
+
         if self._data is None:
             await self._load()
 

homeassistant/components/zeroconf/__init__.py

@@ -108,7 +108,6 @@ def setup(hass, config):
 
     def stop_zeroconf(_):
         """Stop Zeroconf."""
-        zeroconf.unregister_service(info)
         zeroconf.close()
 
     hass.bus.listen_once(EVENT_HOMEASSISTANT_STOP, stop_zeroconf)

homeassistant/core.py

@@ -1288,6 +1288,9 @@ class Config:
         # List of allowed external dirs to access
         self.whitelist_external_dirs: Set[str] = set()
 
+        # If Home Assistant is running in safe mode
+        self.safe_mode: bool = False
+
     def distance(self, lat: float, lon: float) -> Optional[float]:
         """Calculate distance from Home Assistant.
 
@@ -1350,6 +1353,7 @@ class Config:
             "whitelist_external_dirs": self.whitelist_external_dirs,
             "version": __version__,
             "config_source": self.config_source,
+            "safe_mode": self.safe_mode,
         }
 
     def set_time_zone(self, time_zone_str: str) -> None:

homeassistant/loader.py

@@ -41,7 +41,6 @@ DATA_INTEGRATIONS = "integrations"
 DATA_CUSTOM_COMPONENTS = "custom_components"
 PACKAGE_CUSTOM_COMPONENTS = "custom_components"
 PACKAGE_BUILTIN = "homeassistant.components"
-LOOKUP_PATHS = [PACKAGE_CUSTOM_COMPONENTS, PACKAGE_BUILTIN]
 CUSTOM_WARNING = (
     "You are using a custom integration for %s which has not "
     "been tested by Home Assistant. This component might "
@@ -67,6 +66,9 @@ async def _async_get_custom_components(
     hass: "HomeAssistant",
 ) -> Dict[str, "Integration"]:
     """Return list of custom integrations."""
+    if hass.config.safe_mode:
+        return {}
+
     try:
         import custom_components
     except ImportError:
@@ -178,7 +180,7 @@ class Integration:
 
         Will create a stub manifest.
         """
-        comp = _load_file(hass, domain, LOOKUP_PATHS)
+        comp = _load_file(hass, domain, _lookup_path(hass))
 
         if comp is None:
             return None
@@ -464,7 +466,7 @@ class Components:
             component: Optional[ModuleType] = integration.get_component()
         else:
             # Fallback to importing old-school
-            component = _load_file(self._hass, comp_name, LOOKUP_PATHS)
+            component = _load_file(self._hass, comp_name, _lookup_path(self._hass))
 
         if component is None:
             raise ImportError(f"Unable to load {comp_name}")
@@ -546,3 +548,10 @@ def _async_mount_config_dir(hass: "HomeAssistant") -> bool:
     if hass.config.config_dir not in sys.path:
         sys.path.insert(0, hass.config.config_dir)
     return True
+
+
+def _lookup_path(hass: "HomeAssistant") -> List[str]:
+    """Return the lookup paths for legacy lookups."""
+    if hass.config.safe_mode:
+        return [PACKAGE_BUILTIN]
+    return [PACKAGE_CUSTOM_COMPONENTS, PACKAGE_BUILTIN]

homeassistant/util/logging.py

@@ -80,6 +80,7 @@ class AsyncHandler:
 
     def _process(self) -> None:
         """Process log in a thread."""
+        try:
             while True:
                 record = asyncio.run_coroutine_threadsafe(
                     self._queue.get(), self.loop
@@ -90,6 +91,8 @@ class AsyncHandler:
                     return
 
                 self.handler.emit(record)
+        except asyncio.CancelledError:
+            self.handler.close()
 
     def createLock(self) -> None:
         """Ignore lock stuff."""

tests/components/frontend/test_init.py

@@ -126,6 +126,16 @@ async def test_themes_api(hass, hass_ws_client):
     assert msg["result"]["default_theme"] == "default"
     assert msg["result"]["themes"] == {"happy": {"primary-color": "red"}}
 
+    # safe mode
+    hass.config.safe_mode = True
+    await client.send_json({"id": 6, "type": "frontend/get_themes"})
+    msg = await client.receive_json()
+
+    assert msg["result"]["default_theme"] == "safe_mode"
+    assert msg["result"]["themes"] == {
+        "safe_mode": {"primary-color": "#db4437", "accent-color": "#eeee02"}
+    }
+
 
 async def test_themes_set_theme(hass, hass_ws_client):
     """Test frontend.set_theme service."""

tests/components/lovelace/test_init.py

@@ -38,6 +38,13 @@ async def test_lovelace_from_storage(hass, hass_ws_client, hass_storage):
 
     assert response["result"] == {"yo": "hello"}
 
+    # Test with safe mode
+    hass.config.safe_mode = True
+    await client.send_json({"id": 8, "type": "lovelace/config"})
+    response = await client.receive_json()
+    assert not response["success"]
+    assert response["error"]["code"] == "config_not_found"
+
 
 async def test_lovelace_from_storage_save_before_load(
     hass, hass_ws_client, hass_storage

tests/test_bootstrap.py

@@ -250,7 +250,8 @@ async def test_setup_hass(
     log_no_color = Mock()
 
     with patch(
-        "homeassistant.config.async_hass_config_yaml", return_value={"browser": {}}
+        "homeassistant.config.async_hass_config_yaml",
+        return_value={"browser": {}, "frontend": {}},
     ):
         hass = await bootstrap.async_setup_hass(
             config_dir=get_test_config_dir(),
@@ -263,6 +264,7 @@ async def test_setup_hass(
         )
 
     assert "browser" in hass.config.components
+    assert "safe_mode" not in hass.config.components
 
     assert len(mock_enable_logging.mock_calls) == 1
     assert mock_enable_logging.mock_calls[0][1] == (
@@ -382,3 +384,32 @@ async def test_setup_hass_invalid_core_config(
         )
 
     assert "safe_mode" in hass.config.components
+
+
+async def test_setup_safe_mode_if_no_frontend(
+    mock_enable_logging,
+    mock_is_virtual_env,
+    mock_mount_local_lib_path,
+    mock_ensure_config_exists,
+    mock_process_ha_config_upgrade,
+):
+    """Test we setup safe mode if frontend didn't load."""
+    verbose = Mock()
+    log_rotate_days = Mock()
+    log_file = Mock()
+    log_no_color = Mock()
+
+    with patch(
+        "homeassistant.config.async_hass_config_yaml", return_value={"browser": {}}
+    ):
+        hass = await bootstrap.async_setup_hass(
+            config_dir=get_test_config_dir(),
+            verbose=verbose,
+            log_rotate_days=log_rotate_days,
+            log_file=log_file,
+            log_no_color=log_no_color,
+            skip_pip=True,
+            safe_mode=False,
+        )
+
+    assert "safe_mode" in hass.config.components

tests/test_core.py

@@ -904,6 +904,7 @@ class TestConfig(unittest.TestCase):
             "whitelist_external_dirs": set(),
             "version": __version__,
             "config_source": "default",
+            "safe_mode": False,
         }
 
         assert expected == self.config.as_dict()

tests/test_loader.py

@@ -236,3 +236,9 @@ async def test_get_config_flows(hass):
         flows = await loader.async_get_config_flows(hass)
         assert "test_2" in flows
         assert "test_1" not in flows
+
+
+async def test_get_custom_components_safe_mode(hass):
+    """Test that we get empty custom components in safe mode."""
+    hass.config.safe_mode = True
+    assert await loader.async_get_custom_components(hass) == {}