Run post scripts for factory binaries for flashing (#3003)

Co-authored-by: Oxan van Leeuwen <oxan@oxanvanleeuwen.nl>

esphome/components/esp32/__init__.py

@@ -2,8 +2,9 @@ from dataclasses import dataclass
 from typing import Union
 from pathlib import Path
 import logging
+import os
 
-from esphome.helpers import write_file_if_changed
+from esphome.helpers import copy_file_if_changed, write_file_if_changed
 from esphome.const import (
     CONF_BOARD,
     CONF_FRAMEWORK,
@@ -295,6 +296,8 @@ async def to_code(config):
     conf = config[CONF_FRAMEWORK]
     cg.add_platformio_option("platform", conf[CONF_PLATFORM_VERSION])
 
+    cg.add_platformio_option("extra_scripts", ["post:post_build.py"])
+
     if conf[CONF_TYPE] == FRAMEWORK_ESP_IDF:
         cg.add_platformio_option("framework", "espidf")
         cg.add_build_flag("-DUSE_ESP_IDF")
@@ -412,3 +415,10 @@ def copy_files():
             CORE.relative_build_path("partitions.csv"),
             IDF_PARTITIONS_CSV,
         )
+
+    dir = os.path.dirname(__file__)
+    post_build_file = os.path.join(dir, "post_build.py")
+    copy_file_if_changed(
+        post_build_file,
+        CORE.relative_build_path("post_build.py"),
+    )

esphome/components/esp32/post_build.py

@@ -0,0 +1,43 @@
+# Source https://github.com/letscontrolit/ESPEasy/pull/3845#issuecomment-1005864664
+
+import esptool
+
+# pylint: disable=E0602
+Import("env")  # noqa
+
+
+def esp32_create_combined_bin(source, target, env):
+    print("Generating combined binary for serial flashing")
+    app_offset = 0x10000
+
+    new_file_name = env.subst("$BUILD_DIR/${PROGNAME}-factory.bin")
+    sections = env.subst(env.get("FLASH_EXTRA_IMAGES"))
+    firmware_name = env.subst("$BUILD_DIR/${PROGNAME}.bin")
+    chip = env.get("BOARD_MCU")
+    flash_size = env.BoardConfig().get("upload.flash_size")
+    cmd = [
+        "--chip",
+        chip,
+        "merge_bin",
+        "-o",
+        new_file_name,
+        "--flash_size",
+        flash_size,
+    ]
+    print("    Offset | File")
+    for section in sections:
+        sect_adr, sect_file = section.split(" ", 1)
+        print(f" -  {sect_adr} | {sect_file}")
+        cmd += [sect_adr, sect_file]
+
+    print(f" - {hex(app_offset)} | {firmware_name}")
+    cmd += [hex(app_offset), firmware_name]
+
+    print()
+    print(f"Using esptool.py arguments: {' '.join(cmd)}")
+    print()
+    esptool.main(cmd)
+
+
+# pylint: disable=E0602
+env.AddPostAction("$BUILD_DIR/${PROGNAME}.bin", esp32_create_combined_bin)  # noqa

esphome/components/esp8266/__init__.py

@@ -1,4 +1,5 @@
 import logging
+import os
 
 from esphome.const import (
     CONF_BOARD,
@@ -14,6 +15,7 @@ from esphome.const import (
 from esphome.core import CORE, coroutine_with_priority
 import esphome.config_validation as cv
 import esphome.codegen as cg
+from esphome.helpers import copy_file_if_changed
 
 from .const import CONF_RESTORE_FROM_FLASH, KEY_BOARD, KEY_ESP8266, esp8266_ns
 from .boards import ESP8266_FLASH_SIZES, ESP8266_LD_SCRIPTS
@@ -158,6 +160,8 @@ async def to_code(config):
     cg.add_define("ESPHOME_BOARD", config[CONF_BOARD])
     cg.add_define("ESPHOME_VARIANT", "ESP8266")
 
+    cg.add_platformio_option("extra_scripts", ["post:post_build.py"])
+
     conf = config[CONF_FRAMEWORK]
     cg.add_platformio_option("framework", "arduino")
     cg.add_build_flag("-DUSE_ARDUINO")
@@ -210,3 +214,14 @@ async def to_code(config):
 
         if ld_script is not None:
             cg.add_platformio_option("board_build.ldscript", ld_script)
+
+
+# Called by writer.py
+def copy_files():
+
+    dir = os.path.dirname(__file__)
+    post_build_file = os.path.join(dir, "post_build.py")
+    copy_file_if_changed(
+        post_build_file,
+        CORE.relative_build_path("post_build.py"),
+    )

esphome/components/esp8266/post_build.py

@@ -0,0 +1,15 @@
+import shutil
+
+# pylint: disable=E0602
+Import("env")  # noqa
+
+
+def esp8266_copy_factory_bin(source, target, env):
+    firmware_name = env.subst("$BUILD_DIR/${PROGNAME}.bin")
+    new_file_name = env.subst("$BUILD_DIR/${PROGNAME}-factory.bin")
+
+    shutil.copyfile(firmware_name, new_file_name)
+
+
+# pylint: disable=E0602
+env.AddPostAction("$BUILD_DIR/${PROGNAME}.bin", esp8266_copy_factory_bin)  # noqa

esphome/dashboard/dashboard.py

@@ -410,6 +410,17 @@ class DownloadBinaryRequestHandler(BaseHandler):
             filename = f"{storage_json.name}.bin"
             path = storage_json.firmware_bin_path
 
+        elif type == "firmware-factory.bin":
+            storage_path = ext_storage_path(settings.config_dir, configuration)
+            storage_json = StorageJSON.load(storage_path)
+            if storage_json is None:
+                self.send_error(404)
+                return
+            filename = f"{storage_json.name}.bin"
+            path = storage_json.firmware_bin_path.replace(
+                "firmware.bin", "firmware-factory.bin"
+            )
+
         else:
             args = ["esphome", "idedata", settings.rel_path(configuration)]
             rc, stdout, _ = run_system_command(*args)

esphome/writer.py

@@ -290,6 +290,11 @@ def copy_src_tree():
 
         copy_files()
 
+    elif CORE.is_esp8266:
+        from esphome.components.esp8266 import copy_files
+
+        copy_files()
+
 
 def generate_defines_h():
     define_content_l = [x.as_macro for x in CORE.defines]

platformio.ini

@@ -89,6 +89,7 @@ build_flags =
     ${common:arduino.build_flags}
     -DUSE_ESP8266
     -DUSE_ESP8266_FRAMEWORK_ARDUINO
+extra_scripts = post:esphome/components/esp8266/post_build.py
 
 ; This are common settings for the ESP32 (all variants) using Arduino.
 [common:esp32-arduino]
@@ -107,6 +108,7 @@ build_flags =
     ${common:arduino.build_flags}
     -DUSE_ESP32
     -DUSE_ESP32_FRAMEWORK_ARDUINO
+extra_scripts = post:esphome/components/esp32/post_build.py
 
 ; This are common settings for the ESP32 (all variants) using IDF.
 [common:esp32-idf]
@@ -125,6 +127,7 @@ build_flags =
     -Wno-nonnull-compare
     -DUSE_ESP32
     -DUSE_ESP32_FRAMEWORK_ESP_IDF
+extra_scripts = post:esphome/components/esp32/post_build.py
 
 ; All the actual environments are defined below.
 [env:esp8266-arduino]