esphome/esphome/platformio/extra_script.py

"""Run a PlatformIO library ``extraScript`` against a fake SCons env.

The shim execs the script with a stand-in ``env``, captures ``env.Append``
calls (everything else is a logged no-op), and folds the result into the
library's build flags. No sandboxing: the script runs with full process
access, so it carries the same trust as the library's own source.
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass, field
import logging
import os
from pathlib import Path
import shlex
from typing import TYPE_CHECKING

from esphome.core import EsphomeError

if TYPE_CHECKING:
    from esphome.platformio.library import ConvertedLibrary

_LOGGER = logging.getLogger(__name__)


def apply_extra_script(
    component: ConvertedLibrary,
    board_mcu: Callable[[], str],
    pio_platform: str,
) -> None:
    """Run a library's ``extraScript`` and fold its captured env vars into
    ``component.data["build"]["flags"]``.

    ``board_mcu`` is a callable so its lookup runs only when a script will.
    """
    extra_script = component.data.get("build", {}).get("extraScript")
    if not extra_script:
        return
    if not isinstance(extra_script, str):
        # A list/dict value would raise an opaque TypeError on the join below
        raise EsphomeError(
            f"extraScript of library {component.name} must be a string, "
            f"got {type(extra_script).__name__}"
        )
    # Resolve and confine to the library's source dir so a malicious
    # library.json can't escape (e.g. ``"extraScript": "../../etc/passwd"``).
    source_path = component.source_dir
    library_root = source_path.resolve()
    script_path = (source_path / extra_script).resolve()
    if not script_path.is_relative_to(library_root):
        # More hostile than a missing script; must not be quieter than it
        raise EsphomeError(
            f"extraScript {extra_script} of library {component.name} escapes "
            "the library directory"
        )
    if not script_path.is_file():
        # A declared-but-absent script is a broken or half-downloaded
        # package, not an unsupported script; PlatformIO fails on it too
        raise EsphomeError(
            f"extraScript {extra_script} of library {component.name} not found"
        )
    result = run_extra_script(
        script_path,
        library_dir=source_path,
        board_mcu=board_mcu(),
        pio_platform=pio_platform,
    )
    extra_flags = captured_as_build_flags(result, library_dir=source_path)
    if not extra_flags:
        return
    flags = component.data.setdefault("build", {}).setdefault("flags", [])
    if isinstance(flags, str):
        flags = [flags]
    elif not isinstance(flags, list):
        # A null/dict value coerced through a list wrapper would inject a
        # non-string into the compiler command line; fail naming the library
        raise EsphomeError(
            f"Library {component.name} has a malformed build.flags "
            f"({type(flags).__name__}); expected a string or list"
        )
    component.data["build"]["flags"] = [*flags, *extra_flags]


# Keys we know how to translate back into ESPHome's build-flag pipeline.
# Other env.Append kwargs are recorded but ignored downstream.
_CAPTURED_KEYS = frozenset(
    {"CPPPATH", "LIBPATH", "LIBS", "CPPDEFINES", "LINKFLAGS", "CPPFLAGS"}
)


@dataclass
class ExtraScriptResult:
    """Build-var deltas captured from a PIO extra-script ``env.Append`` call."""

    cpppath: list[str] = field(default_factory=list)
    libpath: list[str] = field(default_factory=list)
    libs: list[str] = field(default_factory=list)
    cppdefines: list[str | tuple[str, str]] = field(default_factory=list)
    linkflags: list[str] = field(default_factory=list)
    cppflags: list[str] = field(default_factory=list)


class _FakeSConsEnv:
    """Minimal SCons ``Environment`` stand-in: ``get`` and ``Append`` work;
    every other method is a swallowed no-op so scripts don't abort."""

    def __init__(self, *, board_mcu: str, pio_env: str, pio_platform: str) -> None:
        self._vars: dict[str, str] = {
            "BOARD_MCU": board_mcu,
            "PIOPLATFORM": pio_platform,
            "PIOENV": pio_env,
        }
        self.result = ExtraScriptResult()
        self._warned_methods: set[str] = set()
        self._warned_keys: set[str] = set()
        self._warned_gets: set[str] = set()

    # ----- SCons env API the common scripts use -----

    def get(self, key: str, default: str | None = None) -> str | None:
        if key not in self._vars and key not in self._warned_gets:
            # A script branching on an unmodelled var silently takes the
            # default branch; make that diagnosable from a normal build log
            self._warned_gets.add(key)
            _LOGGER.warning(
                "PIO extra-script env.get(%r) is not modelled; returning the default",
                key,
            )
        return self._vars.get(key, default)

    def __getitem__(self, key: str) -> str:
        # Scripts also read env["BOARD_MCU"]; without this the broad
        # handler would discard every flag the script captured
        return self._vars[key]

    def Append(self, **kwargs) -> None:  # noqa: N802 (SCons API name)
        for key, value in kwargs.items():
            if key not in _CAPTURED_KEYS:
                # Warn once per key so a loop of Appends cannot spam
                if key not in self._warned_keys:
                    self._warned_keys.add(key)
                    _LOGGER.warning(
                        "PIO extra-script env.Append(%s=...) is not captured; ignoring",
                        key,
                    )
                continue
            items = list(value) if isinstance(value, (list, tuple)) else [value]
            bucket = getattr(self.result, key.lower())
            bucket.extend(items)

    # Same keys, same flattened capture; ordering/dedup don't matter since
    # the consumer re-orders anyway
    Prepend = Append
    AppendUnique = Append
    PrependUnique = Append

    # ----- Everything else is a no-op so unsupported scripts don't crash -----

    def __getattr__(self, name: str):
        def _noop(*args, **kwargs):
            # Once per method: a script whose whole effect is env.Replace()
            # must be diagnosable from a normal build log
            if name not in self._warned_methods:
                self._warned_methods.add(name)
                _LOGGER.warning(
                    "PIO extra-script env.%s(...) is not supported; ignoring", name
                )

        return _noop


def run_extra_script(
    script_path: Path,
    *,
    library_dir: Path,
    board_mcu: str,
    pio_platform: str,
) -> ExtraScriptResult:
    """Execute ``script_path`` with a fake SCons env and return captured vars.

    Runs with ``library_dir`` as CWD so relative lookups resolve against
    the library tree. A crashed script warns and returns an empty result,
    never a partial capture.
    """
    env = _FakeSConsEnv(
        board_mcu=board_mcu,
        pio_env=f"esphome_{board_mcu}",
        pio_platform=pio_platform,
    )
    try:
        source = script_path.read_text(encoding="utf-8")
    except OSError as err:
        # An unreadable declared script is a broken package, exactly like a
        # missing one; must not be quieter than that case
        raise EsphomeError(f"extraScript {script_path} is unreadable: {err}") from err
    except UnicodeDecodeError as e:
        # A content problem, best-effort like a SyntaxError below
        _LOGGER.warning(
            "PIO extra-script %s (in %s) is not UTF-8 (%r); ignoring its output",
            script_path,
            library_dir.name,
            e,
        )
        return ExtraScriptResult()
    old_cwd = Path.cwd()
    try:
        # Inside the try: a SyntaxError in a vendored script is just as
        # best-effort as a runtime failure
        code = compile(source, str(script_path), "exec")
        os.chdir(library_dir)
        exec(  # noqa: S102 pylint: disable=exec-used
            code,
            {
                "Import": lambda *_args: None,  # SCons-side import; harmless here
                "env": env,
                "__file__": str(script_path),
                "__name__": "__pio_extra_script__",
            },
        )
    except SystemExit as e:
        if not e.code:
            # sys.exit() / sys.exit(0) is a normal PlatformIO script ending;
            # the capture is complete
            return env.result
        _LOGGER.warning(
            "PIO extra-script %s (in %s) exited with status %r; ignoring its output",
            script_path,
            library_dir.name,
            e.code,
        )
        return ExtraScriptResult()
    except Exception as e:  # noqa: BLE001  # pylint: disable=broad-exception-caught
        # Discard any partial capture: half-applied flags could build wrong
        # firmware that links cleanly.
        _LOGGER.warning(
            "PIO extra-script %s (in %s) raised %r; ignoring its output",
            script_path,
            library_dir.name,
            e,
        )
        return ExtraScriptResult()
    finally:
        os.chdir(old_cwd)
    return env.result


def captured_as_build_flags(
    result: ExtraScriptResult, *, library_dir: Path
) -> list[str]:
    """Translate captured env vars into -L/-l/-D/raw build flags.

    ``LIBPATH`` entries are made relative to ``library_dir`` so the
    generated build files stay portable.
    """
    flags: list[str] = []

    def _strs(bucket: list, kind: str) -> list[str]:
        # Third-party scripts legally append SCons nodes, ints, or dicts;
        # stringifying those into flags would hand the compiler garbage
        good = [entry for entry in bucket if isinstance(entry, str)]
        for entry in bucket:
            if not isinstance(entry, str):
                _LOGGER.warning("Ignoring unsupported %s entry %r", kind, entry)
        return good

    library_root = library_dir.resolve()

    def _anchored(path: str) -> str:
        # Anchor relative paths to library_dir; the script's CWD has been
        # restored by now
        resolved = (library_dir / path).resolve()
        try:
            return str(resolved.relative_to(library_root))
        except ValueError:
            return str(resolved)

    # shlex.quote so a spaced path survives lex_build_flags as one token
    flags.extend(
        f"-I{shlex.quote(_anchored(path))}" for path in _strs(result.cpppath, "CPPPATH")
    )
    flags.extend(
        f"-L{shlex.quote(_anchored(path))}" for path in _strs(result.libpath, "LIBPATH")
    )
    flags.extend(f"-l{shlex.quote(lib)}" for lib in _strs(result.libs, "LIBS"))
    for define in result.cppdefines:
        # SCons also accepts dict/list CPPDEFINES; formatting those blind
        # would hand the compiler garbage like -D{'FOO': '1'}
        if isinstance(define, (tuple, list)) and len(define) == 2:
            flags.append(f"-D{define[0]}={define[1]}")
        elif isinstance(define, str):
            flags.append(f"-D{define}")
        else:
            _LOGGER.warning("Ignoring unsupported CPPDEFINES entry %r", define)
    # Each captured entry is one argv token in SCons; quote so the
    # lex_build_flags round-trip cannot split a spaced value into two
    flags.extend(shlex.quote(f) for f in _strs(result.linkflags, "LINKFLAGS"))
    flags.extend(shlex.quote(f) for f in _strs(result.cppflags, "CPPFLAGS"))
    return flags