397 lines
13 KiB
Python
397 lines
13 KiB
Python
"""Custom loader."""
|
|
from __future__ import annotations
|
|
|
|
from collections import OrderedDict
|
|
from collections.abc import Iterator
|
|
import fnmatch
|
|
from io import StringIO, TextIOWrapper
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any, TextIO, TypeVar, Union, overload
|
|
|
|
import yaml
|
|
|
|
try:
|
|
from yaml import CSafeLoader as FastestAvailableSafeLoader
|
|
|
|
HAS_C_LOADER = True
|
|
except ImportError:
|
|
HAS_C_LOADER = False
|
|
from yaml import SafeLoader as FastestAvailableSafeLoader # type: ignore[misc]
|
|
|
|
from homeassistant.exceptions import HomeAssistantError
|
|
|
|
from .const import SECRET_YAML
|
|
from .objects import Input, NodeListClass, NodeStrClass
|
|
|
|
# mypy: allow-untyped-calls, no-warn-return-any
|
|
|
|
JSON_TYPE = Union[list, dict, str] # pylint: disable=invalid-name
|
|
_DictT = TypeVar("_DictT", bound=dict)
|
|
|
|
_LOGGER = logging.getLogger(__name__)
|
|
|
|
|
|
class Secrets:
|
|
"""Store secrets while loading YAML."""
|
|
|
|
def __init__(self, config_dir: Path) -> None:
|
|
"""Initialize secrets."""
|
|
self.config_dir = config_dir
|
|
self._cache: dict[Path, dict[str, str]] = {}
|
|
|
|
def get(self, requester_path: str, secret: str) -> str:
|
|
"""Return the value of a secret."""
|
|
current_path = Path(requester_path)
|
|
|
|
secret_dir = current_path
|
|
while True:
|
|
secret_dir = secret_dir.parent
|
|
|
|
try:
|
|
secret_dir.relative_to(self.config_dir)
|
|
except ValueError:
|
|
# We went above the config dir
|
|
break
|
|
|
|
secrets = self._load_secret_yaml(secret_dir)
|
|
|
|
if secret in secrets:
|
|
_LOGGER.debug(
|
|
"Secret %s retrieved from secrets.yaml in folder %s",
|
|
secret,
|
|
secret_dir,
|
|
)
|
|
return secrets[secret]
|
|
|
|
raise HomeAssistantError(f"Secret {secret} not defined")
|
|
|
|
def _load_secret_yaml(self, secret_dir: Path) -> dict[str, str]:
|
|
"""Load the secrets yaml from path."""
|
|
if (secret_path := secret_dir / SECRET_YAML) in self._cache:
|
|
return self._cache[secret_path]
|
|
|
|
_LOGGER.debug("Loading %s", secret_path)
|
|
try:
|
|
secrets = load_yaml(str(secret_path))
|
|
|
|
if not isinstance(secrets, dict):
|
|
raise HomeAssistantError("Secrets is not a dictionary")
|
|
|
|
if "logger" in secrets:
|
|
logger = str(secrets["logger"]).lower()
|
|
if logger == "debug":
|
|
_LOGGER.setLevel(logging.DEBUG)
|
|
else:
|
|
_LOGGER.error(
|
|
"Error in secrets.yaml: 'logger: debug' expected, but 'logger: %s' found",
|
|
logger,
|
|
)
|
|
del secrets["logger"]
|
|
except FileNotFoundError:
|
|
secrets = {}
|
|
|
|
self._cache[secret_path] = secrets
|
|
|
|
return secrets
|
|
|
|
|
|
class SafeLoader(FastestAvailableSafeLoader):
|
|
"""The fastest available safe loader."""
|
|
|
|
def __init__(self, stream: Any, secrets: Secrets | None = None) -> None:
|
|
"""Initialize a safe line loader."""
|
|
self.stream = stream
|
|
if isinstance(stream, str):
|
|
self.name = "<unicode string>"
|
|
elif isinstance(stream, bytes):
|
|
self.name = "<byte string>"
|
|
else:
|
|
self.name = getattr(stream, "name", "<file>")
|
|
super().__init__(stream)
|
|
self.secrets = secrets
|
|
|
|
def get_name(self) -> str:
|
|
"""Get the name of the loader."""
|
|
return self.name
|
|
|
|
def get_stream_name(self) -> str:
|
|
"""Get the name of the stream."""
|
|
return self.stream.name or ""
|
|
|
|
|
|
class SafeLineLoader(yaml.SafeLoader):
|
|
"""Loader class that keeps track of line numbers."""
|
|
|
|
def __init__(self, stream: Any, secrets: Secrets | None = None) -> None:
|
|
"""Initialize a safe line loader."""
|
|
super().__init__(stream)
|
|
self.secrets = secrets
|
|
|
|
def compose_node(self, parent: yaml.nodes.Node, index: int) -> yaml.nodes.Node: # type: ignore[override]
|
|
"""Annotate a node with the first line it was seen."""
|
|
last_line: int = self.line
|
|
node: yaml.nodes.Node = super().compose_node(parent, index) # type: ignore[assignment]
|
|
node.__line__ = last_line + 1 # type: ignore[attr-defined]
|
|
return node
|
|
|
|
def get_name(self) -> str:
|
|
"""Get the name of the loader."""
|
|
return self.name
|
|
|
|
def get_stream_name(self) -> str:
|
|
"""Get the name of the stream."""
|
|
return self.stream.name or ""
|
|
|
|
|
|
LoaderType = Union[SafeLineLoader, SafeLoader]
|
|
|
|
|
|
def load_yaml(fname: str, secrets: Secrets | None = None) -> JSON_TYPE:
|
|
"""Load a YAML file."""
|
|
try:
|
|
with open(fname, encoding="utf-8") as conf_file:
|
|
return parse_yaml(conf_file, secrets)
|
|
except UnicodeDecodeError as exc:
|
|
_LOGGER.error("Unable to read file %s: %s", fname, exc)
|
|
raise HomeAssistantError(exc) from exc
|
|
|
|
|
|
def parse_yaml(
|
|
content: str | TextIO | StringIO, secrets: Secrets | None = None
|
|
) -> JSON_TYPE:
|
|
"""Parse YAML with the fastest available loader."""
|
|
if not HAS_C_LOADER:
|
|
return _parse_yaml_pure_python(content, secrets)
|
|
try:
|
|
return _parse_yaml(SafeLoader, content, secrets)
|
|
except yaml.YAMLError:
|
|
# Loading failed, so we now load with the slow line loader
|
|
# since the C one will not give us line numbers
|
|
if isinstance(content, (StringIO, TextIO, TextIOWrapper)):
|
|
# Rewind the stream so we can try again
|
|
content.seek(0, 0)
|
|
return _parse_yaml_pure_python(content, secrets)
|
|
|
|
|
|
def _parse_yaml_pure_python(
|
|
content: str | TextIO | StringIO, secrets: Secrets | None = None
|
|
) -> JSON_TYPE:
|
|
"""Parse YAML with the pure python loader (this is very slow)."""
|
|
try:
|
|
return _parse_yaml(SafeLineLoader, content, secrets)
|
|
except yaml.YAMLError as exc:
|
|
_LOGGER.error(str(exc))
|
|
raise HomeAssistantError(exc) from exc
|
|
|
|
|
|
def _parse_yaml(
|
|
loader: type[SafeLoader] | type[SafeLineLoader],
|
|
content: str | TextIO,
|
|
secrets: Secrets | None = None,
|
|
) -> JSON_TYPE:
|
|
"""Load a YAML file."""
|
|
# If configuration file is empty YAML returns None
|
|
# We convert that to an empty dict
|
|
return (
|
|
yaml.load(content, Loader=lambda stream: loader(stream, secrets))
|
|
or OrderedDict()
|
|
)
|
|
|
|
|
|
@overload
|
|
def _add_reference(
|
|
obj: list | NodeListClass,
|
|
loader: LoaderType,
|
|
node: yaml.nodes.Node,
|
|
) -> NodeListClass:
|
|
...
|
|
|
|
|
|
@overload
|
|
def _add_reference(
|
|
obj: str | NodeStrClass,
|
|
loader: LoaderType,
|
|
node: yaml.nodes.Node,
|
|
) -> NodeStrClass:
|
|
...
|
|
|
|
|
|
@overload
|
|
def _add_reference(obj: _DictT, loader: LoaderType, node: yaml.nodes.Node) -> _DictT:
|
|
...
|
|
|
|
|
|
def _add_reference(obj, loader: LoaderType, node: yaml.nodes.Node): # type: ignore[no-untyped-def]
|
|
"""Add file reference information to an object."""
|
|
if isinstance(obj, list):
|
|
obj = NodeListClass(obj)
|
|
if isinstance(obj, str):
|
|
obj = NodeStrClass(obj)
|
|
setattr(obj, "__config_file__", loader.get_name())
|
|
setattr(obj, "__line__", node.start_mark.line)
|
|
return obj
|
|
|
|
|
|
def _include_yaml(loader: LoaderType, node: yaml.nodes.Node) -> JSON_TYPE:
|
|
"""Load another YAML file and embeds it using the !include tag.
|
|
|
|
Example:
|
|
device_tracker: !include device_tracker.yaml
|
|
|
|
"""
|
|
fname = os.path.join(os.path.dirname(loader.get_name()), node.value)
|
|
try:
|
|
return _add_reference(load_yaml(fname, loader.secrets), loader, node)
|
|
except FileNotFoundError as exc:
|
|
raise HomeAssistantError(
|
|
f"{node.start_mark}: Unable to read file {fname}."
|
|
) from exc
|
|
|
|
|
|
def _is_file_valid(name: str) -> bool:
|
|
"""Decide if a file is valid."""
|
|
return not name.startswith(".")
|
|
|
|
|
|
def _find_files(directory: str, pattern: str) -> Iterator[str]:
|
|
"""Recursively load files in a directory."""
|
|
for root, dirs, files in os.walk(directory, topdown=True):
|
|
dirs[:] = [d for d in dirs if _is_file_valid(d)]
|
|
for basename in sorted(files):
|
|
if _is_file_valid(basename) and fnmatch.fnmatch(basename, pattern):
|
|
filename = os.path.join(root, basename)
|
|
yield filename
|
|
|
|
|
|
def _include_dir_named_yaml(loader: LoaderType, node: yaml.nodes.Node) -> OrderedDict:
|
|
"""Load multiple files from directory as a dictionary."""
|
|
mapping: OrderedDict = OrderedDict()
|
|
loc = os.path.join(os.path.dirname(loader.get_name()), node.value)
|
|
for fname in _find_files(loc, "*.yaml"):
|
|
filename = os.path.splitext(os.path.basename(fname))[0]
|
|
if os.path.basename(fname) == SECRET_YAML:
|
|
continue
|
|
mapping[filename] = load_yaml(fname, loader.secrets)
|
|
return _add_reference(mapping, loader, node)
|
|
|
|
|
|
def _include_dir_merge_named_yaml(
|
|
loader: LoaderType, node: yaml.nodes.Node
|
|
) -> OrderedDict:
|
|
"""Load multiple files from directory as a merged dictionary."""
|
|
mapping: OrderedDict = OrderedDict()
|
|
loc = os.path.join(os.path.dirname(loader.get_name()), node.value)
|
|
for fname in _find_files(loc, "*.yaml"):
|
|
if os.path.basename(fname) == SECRET_YAML:
|
|
continue
|
|
loaded_yaml = load_yaml(fname, loader.secrets)
|
|
if isinstance(loaded_yaml, dict):
|
|
mapping.update(loaded_yaml)
|
|
return _add_reference(mapping, loader, node)
|
|
|
|
|
|
def _include_dir_list_yaml(
|
|
loader: LoaderType, node: yaml.nodes.Node
|
|
) -> list[JSON_TYPE]:
|
|
"""Load multiple files from directory as a list."""
|
|
loc = os.path.join(os.path.dirname(loader.get_name()), node.value)
|
|
return [
|
|
load_yaml(f, loader.secrets)
|
|
for f in _find_files(loc, "*.yaml")
|
|
if os.path.basename(f) != SECRET_YAML
|
|
]
|
|
|
|
|
|
def _include_dir_merge_list_yaml(
|
|
loader: LoaderType, node: yaml.nodes.Node
|
|
) -> JSON_TYPE:
|
|
"""Load multiple files from directory as a merged list."""
|
|
loc: str = os.path.join(os.path.dirname(loader.get_name()), node.value)
|
|
merged_list: list[JSON_TYPE] = []
|
|
for fname in _find_files(loc, "*.yaml"):
|
|
if os.path.basename(fname) == SECRET_YAML:
|
|
continue
|
|
loaded_yaml = load_yaml(fname, loader.secrets)
|
|
if isinstance(loaded_yaml, list):
|
|
merged_list.extend(loaded_yaml)
|
|
return _add_reference(merged_list, loader, node)
|
|
|
|
|
|
def _ordered_dict(loader: LoaderType, node: yaml.nodes.MappingNode) -> OrderedDict:
|
|
"""Load YAML mappings into an ordered dictionary to preserve key order."""
|
|
loader.flatten_mapping(node)
|
|
nodes = loader.construct_pairs(node)
|
|
|
|
seen: dict = {}
|
|
for (key, _), (child_node, _) in zip(nodes, node.value):
|
|
line = child_node.start_mark.line
|
|
|
|
try:
|
|
hash(key)
|
|
except TypeError as exc:
|
|
fname = loader.get_stream_name()
|
|
raise yaml.MarkedYAMLError(
|
|
context=f'invalid key: "{key}"',
|
|
context_mark=yaml.Mark(fname, 0, line, -1, None, None), # type: ignore[arg-type]
|
|
) from exc
|
|
|
|
if key in seen:
|
|
fname = loader.get_stream_name()
|
|
_LOGGER.warning(
|
|
'YAML file %s contains duplicate key "%s". Check lines %d and %d',
|
|
fname,
|
|
key,
|
|
seen[key],
|
|
line,
|
|
)
|
|
seen[key] = line
|
|
|
|
return _add_reference(OrderedDict(nodes), loader, node)
|
|
|
|
|
|
def _construct_seq(loader: LoaderType, node: yaml.nodes.Node) -> JSON_TYPE:
|
|
"""Add line number and file name to Load YAML sequence."""
|
|
(obj,) = loader.construct_yaml_seq(node)
|
|
return _add_reference(obj, loader, node)
|
|
|
|
|
|
def _env_var_yaml(loader: LoaderType, node: yaml.nodes.Node) -> str:
|
|
"""Load environment variables and embed it into the configuration YAML."""
|
|
args = node.value.split()
|
|
|
|
# Check for a default value
|
|
if len(args) > 1:
|
|
return os.getenv(args[0], " ".join(args[1:]))
|
|
if args[0] in os.environ:
|
|
return os.environ[args[0]]
|
|
_LOGGER.error("Environment variable %s not defined", node.value)
|
|
raise HomeAssistantError(node.value)
|
|
|
|
|
|
def secret_yaml(loader: LoaderType, node: yaml.nodes.Node) -> JSON_TYPE:
|
|
"""Load secrets and embed it into the configuration YAML."""
|
|
if loader.secrets is None:
|
|
raise HomeAssistantError("Secrets not supported in this YAML file")
|
|
|
|
return loader.secrets.get(loader.get_name(), node.value)
|
|
|
|
|
|
def add_constructor(tag: Any, constructor: Any) -> None:
|
|
"""Add to constructor to all loaders."""
|
|
for yaml_loader in (SafeLoader, SafeLineLoader):
|
|
yaml_loader.add_constructor(tag, constructor)
|
|
|
|
|
|
add_constructor("!include", _include_yaml)
|
|
add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, _ordered_dict)
|
|
add_constructor(yaml.resolver.BaseResolver.DEFAULT_SEQUENCE_TAG, _construct_seq)
|
|
add_constructor("!env_var", _env_var_yaml)
|
|
add_constructor("!secret", secret_yaml)
|
|
add_constructor("!include_dir_list", _include_dir_list_yaml)
|
|
add_constructor("!include_dir_merge_list", _include_dir_merge_list_yaml)
|
|
add_constructor("!include_dir_named", _include_dir_named_yaml)
|
|
add_constructor("!include_dir_merge_named", _include_dir_merge_named_yaml)
|
|
add_constructor("!input", Input.from_node)
|