esphome-flasher/esphomeflasher/common.py

import io
import struct

import esptool

from esphomeflasher.const import HTTP_REGEX
from esphomeflasher.helpers import prevent_print


class EsphomeflasherError(Exception):
    pass


class MockEsptoolArgs:
    def __init__(self, flash_size, addr_filename, flash_mode, flash_freq):
        self.compress = True
        self.no_compress = False
        self.flash_size = flash_size
        self.addr_filename = addr_filename
        self.flash_mode = flash_mode
        self.flash_freq = flash_freq
        self.no_stub = False
        self.verify = False
        self.erase_all = False
        self.encrypt = False
        self.encrypt_files = None


class ChipInfo:
    def __init__(self, family, model, mac):
        self.family = family
        self.model = model
        self.mac = mac
        self.is_esp32 = None

    def as_dict(self):
        return {
            "family": self.family,
            "model": self.model,
            "mac": self.mac,
            "is_esp32": self.is_esp32,
        }


class ESP32ChipInfo(ChipInfo):
    def __init__(
        self,
        model,
        mac,
        num_cores,
        cpu_frequency,
        has_bluetooth,
        has_embedded_flash,
        has_factory_calibrated_adc,
    ):
        super().__init__("ESP32", model, mac)
        self.num_cores = num_cores
        self.cpu_frequency = cpu_frequency
        self.has_bluetooth = has_bluetooth
        self.has_embedded_flash = has_embedded_flash
        self.has_factory_calibrated_adc = has_factory_calibrated_adc

    def as_dict(self):
        data = ChipInfo.as_dict(self)
        data.update(
            {
                "num_cores": self.num_cores,
                "cpu_frequency": self.cpu_frequency,
                "has_bluetooth": self.has_bluetooth,
                "has_embedded_flash": self.has_embedded_flash,
                "has_factory_calibrated_adc": self.has_factory_calibrated_adc,
            }
        )
        return data


class ESP8266ChipInfo(ChipInfo):
    def __init__(self, model, mac, chip_id):
        super().__init__("ESP8266", model, mac)
        self.chip_id = chip_id

    def as_dict(self):
        data = ChipInfo.as_dict(self)
        data.update(
            {
                "chip_id": self.chip_id,
            }
        )
        return data


def read_chip_property(func, *args, **kwargs):
    try:
        return prevent_print(func, *args, **kwargs)
    except esptool.FatalError as err:
        raise EsphomeflasherError(f"Reading chip details failed: {err}") from err


def read_chip_info(chip):
    mac = ":".join(f"{x:02X}" for x in read_chip_property(chip.read_mac))
    if isinstance(chip, esptool.ESP32ROM):
        model = read_chip_property(chip.get_chip_description)
        features = read_chip_property(chip.get_chip_features)
        num_cores = 2 if "Dual Core" in features else 1
        frequency = next((x for x in ("160MHz", "240MHz") if x in features), "80MHz")
        has_bluetooth = "BT" in features
        has_embedded_flash = "Embedded Flash" in features
        has_factory_calibrated_adc = "VRef calibration in efuse" in features
        return ESP32ChipInfo(
            model,
            mac,
            num_cores,
            frequency,
            has_bluetooth,
            has_embedded_flash,
            has_factory_calibrated_adc,
        )
    if isinstance(chip, esptool.ESP8266ROM):
        model = read_chip_property(chip.get_chip_description)
        chip_id = read_chip_property(chip.chip_id)
        return ESP8266ChipInfo(model, mac, chip_id)
    raise EsphomeflasherError(f"Unknown chip type {type(chip)}")


def chip_run_stub(chip):
    try:
        return chip.run_stub()
    except esptool.FatalError as err:
        raise EsphomeflasherError(
            f"Error putting ESP in stub flash mode: {err}"
        ) from err


def detect_flash_size(stub_chip):
    flash_id = read_chip_property(stub_chip.flash_id)
    return esptool.DETECTED_FLASH_SIZES.get(flash_id >> 16, "4MB")


def read_firmware_info(firmware):
    header = firmware.read(4)
    firmware.seek(0)

    magic, _, flash_mode_raw, flash_size_freq = struct.unpack("BBBB", header)
    if magic != esptool.ESPLoader.ESP_IMAGE_MAGIC:
        raise EsphomeflasherError(
            f"The firmware binary is invalid (magic byte={magic:02X}, should be {esptool.ESPLoader.ESP_IMAGE_MAGIC:02X})"
        )
    flash_freq_raw = flash_size_freq & 0x0F
    flash_mode = {0: "qio", 1: "qout", 2: "dio", 3: "dout"}.get(flash_mode_raw)
    flash_freq = {0: "40m", 1: "26m", 2: "20m", 0xF: "80m"}.get(flash_freq_raw)
    return flash_mode, flash_freq


def open_downloadable_binary(path):
    if hasattr(path, "seek"):
        path.seek(0)
        return path

    if HTTP_REGEX.match(path) is not None:
        import requests

        try:
            response = requests.get(path)
            response.raise_for_status()
        except requests.exceptions.Timeout as err:
            raise EsphomeflasherError(
                f"Timeout while retrieving firmware file '{path}': {err}"
            ) from err
        except requests.exceptions.RequestException as err:
            raise EsphomeflasherError(
                f"Error while retrieving firmware file '{path}': {err}"
            ) from err

        binary = io.BytesIO()
        binary.write(response.content)
        binary.seek(0)
        return binary

    try:
        return open(path, "rb")
    except IOError as err:
        raise EsphomeflasherError(f"Error opening binary '{path}': {err}") from err


def format_bootloader_path(path, flash_mode, flash_freq):
    return path.replace("$FLASH_MODE$", flash_mode).replace("$FLASH_FREQ$", flash_freq)


def configure_write_flash_args(
    info, firmware_path, flash_size, bootloader_path, partitions_path, otadata_path
):
    addr_filename = []
    firmware = open_downloadable_binary(firmware_path)
    flash_mode, flash_freq = read_firmware_info(firmware)
    if isinstance(info, ESP32ChipInfo):
        if flash_freq in ("26m", "20m"):
            raise EsphomeflasherError(
                f"No bootloader available for flash frequency {flash_freq}"
            )
        bootloader = open_downloadable_binary(
            format_bootloader_path(bootloader_path, flash_mode, flash_freq)
        )
        partitions = open_downloadable_binary(partitions_path)
        otadata = open_downloadable_binary(otadata_path)

        addr_filename.append((0x1000, bootloader))
        addr_filename.append((0x8000, partitions))
        addr_filename.append((0xE000, otadata))
        addr_filename.append((0x10000, firmware))
    else:
        addr_filename.append((0x0, firmware))
    return MockEsptoolArgs(flash_size, addr_filename, flash_mode, flash_freq)


def detect_chip(port, force_esp8266=False, force_esp32=False):
    if force_esp8266 or force_esp32:
        klass = esptool.ESP32ROM if force_esp32 else esptool.ESP8266ROM
        chip = klass(port)
    else:
        try:
            chip = esptool.ESPLoader.detect_chip(port)
        except esptool.FatalError as err:
            if "Wrong boot mode detected" in str(err):
                msg = "ESP is not in flash boot mode. If your board has a flashing pin, try again while keeping it pressed."
            else:
                msg = f"ESP Chip Auto-Detection failed: {err}"
            raise EsphomeflasherError(msg) from err

    try:
        chip.connect()
    except esptool.FatalError as err:
        raise EsphomeflasherError(f"Error connecting to ESP: {err}") from err

    return chip