Skip to content

Built-in esptool flash, install,info,erase_flash (via Pythonx) #46

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
241 changes: 241 additions & 0 deletions lib/esptool_helper.ex
Original file line number Diff line number Diff line change
@@ -0,0 +1,241 @@
defmodule ExAtomVM.EsptoolHelper do
@moduledoc """
Module for setting up and using esptool through Pythonx.
"""

@doc """
Initializes Python environment with project configuration.
We use locked main branch esptool version, pending a stable 5.x release,
as we need the read to memory (instead of only to file) features.
"""
def setup do
case Code.ensure_loaded(Pythonx) do
{:module, Pythonx} ->
Application.ensure_all_started(:pythonx)

Pythonx.uv_init("""
[project]
name = "project"
version = "0.0.0"
requires-python = "==3.13.*"
dependencies = [
"esptool @ git+https://github.com/espressif/esptool.git@6f0d779"
]
""")

_ ->
{:error, :pythonx_not_available,
"The :pythonx dependency is not available. Please add it to your mix.exs dependencies.\n{:pythonx, \"~> 0.4.0\"}"}
end
end

def flash_pythonx(tool_args) do
# https://github.com/espressif/esptool/blob/master/docs/en/esptool/scripting.rst

tool_args =
if not Enum.member?(tool_args, "--port") do
selected_device = select_device()

["--port", selected_device["port"]] ++ tool_args
else
tool_args
end

{_result, globals} =
try do
Pythonx.eval(
"""
import esptool
import sys

command = [arg.decode('utf-8') for arg in tool_args]

def flash_esp():
esptool.main(command)

if __name__ == "__main__":
try:
result = flash_esp()
result = True
except SystemExit as e:
exit_code = int(str(e))
result = exit_code == 0
except Exception as e:
print(f"Warning: {e}")
result = True

""",
%{"tool_args" => tool_args}
)
rescue
e in Pythonx.Error ->
IO.inspect("Pythonx error occurred: #{inspect(e)}")
exit({:shutdown, 1})
end

Pythonx.decode(globals["result"])
end

@doc """
Erases flash of an ESP32 device.
--after "no-reset" is needed for keeping USB-OTG devices like esp32-S2 in a good state.
"""
def erase_flash(tool_args \\ ["--chip", "auto", "--after", "no-reset"]) do
tool_args =
if not Enum.member?(tool_args, "--port") do
selected_device = select_device()

confirmation =
IO.gets(
"\nAre you sure you want to erase the flash of\n#{selected_device["chip_family_name"]} - Port: #{selected_device["port"]} MAC: #{selected_device["mac_address"]} ? [N/y]: "
)

case String.trim(confirmation) do
input when input in ["Y", "y"] ->
IO.puts("Erasing..")

_ ->
IO.puts("Flash erase cancelled.")
exit({:shutdown, 0})
end

["--port", selected_device["port"]] ++ tool_args ++ ["erase-flash"]
else
tool_args ++ ["erase-flash"]
end

{_result, globals} =
try do
Pythonx.eval(
"""
import esptool

command = [arg.decode('utf-8') for arg in tool_args]

def flash_esp():
esptool.main(command)

if __name__ == "__main__":
try:
result = flash_esp()
result = True
except SystemExit as e:
exit_code = int(str(e))
result = exit_code == 0
except Exception as e:
print(f"Warning: {e}")
result = False
""",
%{"tool_args" => tool_args}
)
rescue
e in Pythonx.Error ->
IO.inspect("Pythonx error occurred: #{inspect(e)}")
exit({:shutdown, 1})
end

Pythonx.decode(globals["result"])
end

def connected_devices do
{_result, globals} =
try do
Pythonx.eval(
"""
from esptool.cmds import (detect_chip, read_flash, attach_flash)
import serial.tools.list_ports as list_ports
import re

ports = []
for port in list_ports.comports():
if port.vid is None:
continue
ports.append(port.device)

result = []
for port in ports:
try:
with detect_chip(port) as esp:
description = esp.get_chip_description()
features = esp.get_chip_features()
mac_addr = ':'.join(['%02X' % b for b in esp.read_mac()])

# chips like esp32-s2 can have more specific names, so we call this chip family
# https://github.com/espressif/esptool/blob/807d02b0c5eb07ba46f871a492c84395fb9f37be/esptool/targets/esp32s2.py#L167
chip_family_name = esp.CHIP_NAME

# read 128 bytes at 0x10030
attach_flash(esp)
app_header = read_flash(esp, 0x10030, 128, None)
app_header_strings = [s for s in re.split('\\x00', app_header.decode('utf-8', errors='replace')) if s]

usb_mode = esp.get_usb_mode()

# this is needed to keep USB-OTG boards like esp32-S2 in a good state
esp.run_stub()

result.append({"port": port, "chip_family_name": chip_family_name,
"features": features, "build_info": app_header_strings,
"mac_address": mac_addr, "usb_mode": usb_mode
})
except Exception as e:
print(f"Error: {e}")
result = []
""",
%{}
)
rescue
e in Pythonx.Error ->
{:error, "Pythonx error occurred: #{inspect(e)}"}
end

Pythonx.decode(globals["result"])
|> Enum.map(fn device ->
Map.put(device, "atomvm_installed", Enum.member?(device["build_info"], "atomvm-esp32"))
end)
end

def select_device do
devices = connected_devices()

selected_device =
case length(devices) do
0 ->
IO.puts(
"Found no esp32 devices..\nYou may have to hold BOOT button down while plugging in the device"
)

exit({:shutdown, 1})

1 ->
hd(devices)

_ ->
IO.puts("\nMultiple ESP32 devices found:")

devices
|> Enum.with_index(1)
|> Enum.each(fn {device, index} ->
IO.puts(
"#{index}. #{device["chip_family_name"]} - Port: #{device["port"]} MAC: #{device["mac_address"]}"
)
end)

selected =
IO.gets("\nSelect device (1-#{length(devices)}): ")
|> String.trim()
|> Integer.parse()

case selected do
{num, _} when num > 0 and num <= length(devices) ->
Enum.at(devices, num - 1)

_ ->
IO.puts("Invalid selection.")
exit({:shutdown, 1})
end
end

selected_device
end
end
27 changes: 27 additions & 0 deletions lib/mix/tasks/esp32.erase_flash.ex
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
defmodule Mix.Tasks.Atomvm.Esp32.EraseFlash do
@moduledoc """
Mix task to get erase the flash of an connected ESP32 devices.
"""
use Mix.Task

@shortdoc "Erase flash of ESP32"

@impl Mix.Task
def run(_args) do
with :ok <- ExAtomVM.EsptoolHelper.setup(),
result <- ExAtomVM.EsptoolHelper.erase_flash() do
case result do
true -> exit({:shutdown, 0})
false -> exit({:shutdown, 1})
end
else
{:error, :pythonx_not_available, message} ->
IO.puts("\nError: #{message}")
exit({:shutdown, 1})

{:error, reason} ->
IO.puts("Error: #{reason}")
exit({:shutdown, 1})
end
end
end
106 changes: 106 additions & 0 deletions lib/mix/tasks/esp32.info.ex
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
defmodule Mix.Tasks.Atomvm.Esp32.Info do
@moduledoc """
Mix task to get information about connected ESP32 devices.
"""
use Mix.Task

@shortdoc "Get information about connected ESP32 devices"

@impl Mix.Task
def run(_args) do
with :ok <- ExAtomVM.EsptoolHelper.setup(),
devices <- ExAtomVM.EsptoolHelper.connected_devices() do
case length(devices) do
0 ->
IO.puts(
"Found no esp32 devices..\nYou may have to hold BOOT button down while plugging in the device"
)

count ->
IO.puts("Found #{count} connected esp32:")
end

if length(devices) > 1 do
Enum.each(devices, fn device ->
IO.puts(
"#{format_atomvm_status(device["atomvm_installed"])}#{device["chip_family_name"]} - Port: #{device["port"]}"
)
end)
end

Enum.each(devices, fn device ->
IO.puts("\n━━━━━━━━━━━━━━━━━━━━━━")

IO.puts(
"#{format_atomvm_status(device["atomvm_installed"])}#{device["chip_family_name"]} - Port: #{device["port"]}"
)

IO.puts("USB_MODE: #{device["usb_mode"]}")
IO.puts("MAC: #{device["mac_address"]}")
IO.puts("AtomVM installed: #{device["atomvm_installed"]}")

IO.puts("\nBuild Information:")

Enum.each(format_build_info(device["build_info"]), fn build_info ->
IO.puts(build_info)
end)

IO.puts("\nFeatures:")

Enum.each(device["features"], fn feature ->
IO.puts(" · #{feature}")
end)
end)

IO.puts("\n")
else
{:error, :pythonx_not_available, message} ->
IO.puts("\nError: #{message}")
exit({:shutdown, 1})

{:error, reason} ->
IO.puts("\nError: Failed to get ESP32 device information")
IO.puts("Reason: #{reason}")
exit({:shutdown, 1})
end
end

defp format_build_info(build_info) when is_list(build_info) and length(build_info) == 5 do
[version, target, time, date, sdk] =
build_info
|> Enum.map(&sanitize_string/1)

[
" Version: #{version}",
" Target: #{target}",
" Built: #{time} #{date}",
" SDK: #{sdk}"
]
end

defp format_build_info(build_info) when is_list(build_info) do
build_info
|> Enum.map(&sanitize_string/1)
|> Enum.with_index(1)
|> Enum.map(fn {info, index} -> " Info #{index}: #{info}" end)
end

defp format_build_info(_) do
[" Build info not available or corrupted"]
end

defp sanitize_string(str) when is_binary(str) do
str
# Remove non-printable characters while preserving spaces
|> String.replace(~r/[^\x20-\x7E\s]/u, "")
|> case do
"" -> "<unreadable>"
sanitized -> sanitized
end
end

defp sanitize_string(_), do: "<invalid>"

defp format_atomvm_status(true), do: "✅"
defp format_atomvm_status(_), do: "❌"
end
Loading