esphome-flasher/esphomeflasher/__main__.py
2021-10-29 12:48:46 +02:00

222 lines
6.8 KiB
Python

from __future__ import print_function
import argparse
import sys
import time
from datetime import datetime
import esptool
import serial
from esphomeflasher import const
from esphomeflasher.common import (
ESP32ChipInfo,
EsphomeflasherError,
chip_run_stub,
configure_write_flash_args,
detect_chip,
detect_flash_size,
read_chip_info,
)
from esphomeflasher.const import (
ESP32_DEFAULT_BOOTLOADER_FORMAT,
ESP32_DEFAULT_OTA_DATA,
ESP32_DEFAULT_PARTITIONS,
)
from esphomeflasher.helpers import list_serial_ports
def parse_args(argv):
parser = argparse.ArgumentParser(prog=f"esphomeflasher {const.__version__}")
parser.add_argument("-p", "--port", help="Select the USB/COM port for uploading.")
group = parser.add_mutually_exclusive_group(required=False)
group.add_argument("--esp8266", action="store_true")
group.add_argument("--esp32", action="store_true")
group.add_argument(
"--upload-baud-rate",
type=int,
default=460800,
help="Baud rate to upload with (not for logging)",
)
parser.add_argument(
"--bootloader",
help="(ESP32-only) The bootloader to flash.",
default=ESP32_DEFAULT_BOOTLOADER_FORMAT,
)
parser.add_argument(
"--partitions",
help="(ESP32-only) The partitions to flash.",
default=ESP32_DEFAULT_PARTITIONS,
)
parser.add_argument(
"--otadata",
help="(ESP32-only) The otadata file to flash.",
default=ESP32_DEFAULT_OTA_DATA,
)
parser.add_argument(
"--no-erase", help="Do not erase flash before flashing", action="store_true"
)
parser.add_argument("--show-logs", help="Only show logs", action="store_true")
parser.add_argument("binary", help="The binary image to flash.")
return parser.parse_args(argv[1:])
def select_port(args):
if args.port is not None:
print(f"Using '{args.port}' as serial port.")
return args.port
ports = list_serial_ports()
if not ports:
raise EsphomeflasherError("No serial port found!")
if len(ports) != 1:
print("Found more than one serial port:")
for port, desc in ports:
print(f" * {port} ({desc})")
print("Please choose one with the --port argument.")
raise EsphomeflasherError
print(f"Auto-detected serial port: {ports[0][0]}")
return ports[0][0]
def show_logs(serial_port):
print("Showing logs:")
with serial_port:
while True:
try:
raw = serial_port.readline()
except serial.SerialException:
print("Serial port closed!")
return
text = raw.decode(errors="ignore")
line = text.replace("\r", "").replace("\n", "")
time_ = datetime.now().time().strftime("[%H:%M:%S]")
message = time_ + line
try:
print(message)
except UnicodeEncodeError:
print(message.encode("ascii", "backslashreplace"))
def run_esphomeflasher(argv):
args = parse_args(argv)
port = select_port(args)
if args.show_logs:
serial_port = serial.Serial(port, baudrate=115200)
show_logs(serial_port)
return
try:
# pylint: disable=consider-using-with
firmware = open(args.binary, "rb")
except IOError as err:
raise EsphomeflasherError(f"Error opening binary: {err}") from err
chip = detect_chip(port, args.esp8266, args.esp32)
info = read_chip_info(chip)
print()
print("Chip Info:")
print(f" - Chip Family: {info.family}")
print(f" - Chip Model: {info.model}")
if isinstance(info, ESP32ChipInfo):
print(f" - Number of Cores: {info.num_cores}")
print(f" - Max CPU Frequency: {info.cpu_frequency}")
print(f" - Has Bluetooth: {'YES' if info.has_bluetooth else 'NO'}")
print(f" - Has Embedded Flash: {'YES' if info.has_embedded_flash else 'NO'}")
print(
f" - Has Factory-Calibrated ADC: {'YES' if info.has_factory_calibrated_adc else 'NO'}"
)
else:
print(f" - Chip ID: {info.chip_id:08X}")
print(f" - MAC Address: {info.mac}")
stub_chip = chip_run_stub(chip)
flash_size = None
if args.upload_baud_rate != 115200:
try:
stub_chip.change_baud(args.upload_baud_rate)
except esptool.FatalError as err:
raise EsphomeflasherError(
f"Error changing ESP upload baud rate: {err}"
) from err
# Check if the higher baud rate works
try:
flash_size = detect_flash_size(stub_chip)
except EsphomeflasherError:
# Go back to old baud rate by recreating chip instance
print(
f"Chip does not support baud rate {args.upload_baud_rate}, changing to 115200"
)
# pylint: disable=protected-access
stub_chip._port.close()
chip = detect_chip(port, args.esp8266, args.esp32)
stub_chip = chip_run_stub(chip)
if flash_size is None:
flash_size = detect_flash_size(stub_chip)
print(f" - Flash Size: {flash_size}")
mock_args = configure_write_flash_args(
info, firmware, flash_size, args.bootloader, args.partitions, args.otadata
)
print(f" - Flash Mode: {mock_args.flash_mode}")
print(f" - Flash Frequency: {mock_args.flash_freq.upper()}Hz")
try:
stub_chip.flash_set_parameters(esptool.flash_size_bytes(flash_size))
except esptool.FatalError as err:
raise EsphomeflasherError(f"Error setting flash parameters: {err}") from err
if not args.no_erase:
try:
esptool.erase_flash(stub_chip, mock_args)
except esptool.FatalError as err:
raise EsphomeflasherError(f"Error while erasing flash: {err}") from err
try:
esptool.write_flash(stub_chip, mock_args)
except esptool.FatalError as err:
raise EsphomeflasherError(f"Error while writing flash: {err}") from err
print("Hard Resetting...")
stub_chip.hard_reset()
print("Done! Flashing is complete!")
print()
if args.upload_baud_rate != 115200:
# pylint: disable=protected-access
stub_chip._port.baudrate = 115200
time.sleep(0.05) # get rid of crap sent during baud rate change
# pylint: disable=protected-access
stub_chip._port.flushInput()
# pylint: disable=protected-access
show_logs(stub_chip._port)
def main():
try:
if len(sys.argv) <= 1:
from esphomeflasher import gui
return gui.main() or 0
return run_esphomeflasher(sys.argv) or 0
except EsphomeflasherError as err:
msg = str(err)
if msg:
print(msg)
return 1
except KeyboardInterrupt:
return 1
if __name__ == "__main__":
sys.exit(main())