Ryan Malloy
945939bdad
Security manager: eFuse read/burn via espefuse, security audit combining chip-id + get-security-info + eFuse summary, flash encryption status checking. Diagnostics: memory dump with hex formatting via dump-mem, performance profiling by timing esptool operations, diagnostic report combining chip-id + read-mac + flash-id. Partition manager: OTA partition table generation with auto-layout, custom partition table from config dict with validation, binary partition table parsing from flash reads at 0x8000.
447 lines
16 KiB
Python
447 lines
16 KiB
Python
"""
|
|
Partition Manager Component
|
|
|
|
Handles ESP partition table operations: generating OTA-capable tables,
|
|
custom partition layouts, and reading/analyzing partition tables from
|
|
connected devices.
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from fastmcp import Context, FastMCP
|
|
|
|
from ..config import ESPToolServerConfig
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ESP partition types and subtypes
|
|
PARTITION_TYPES = {
|
|
"app": 0x00,
|
|
"data": 0x01,
|
|
}
|
|
|
|
APP_SUBTYPES = {
|
|
"factory": 0x00,
|
|
"ota_0": 0x10,
|
|
"ota_1": 0x11,
|
|
"ota_2": 0x12,
|
|
"ota_3": 0x13,
|
|
"test": 0x20,
|
|
}
|
|
|
|
DATA_SUBTYPES = {
|
|
"ota": 0x00,
|
|
"phy": 0x01,
|
|
"nvs": 0x02,
|
|
"coredump": 0x03,
|
|
"nvs_keys": 0x04,
|
|
"efuse": 0x05,
|
|
"spiffs": 0x82,
|
|
"littlefs": 0x83,
|
|
"fat": 0x81,
|
|
}
|
|
|
|
# Size multipliers
|
|
_SIZE_MULT = {"K": 1024, "M": 1024 * 1024}
|
|
|
|
|
|
def _parse_size_spec(spec: str) -> int:
|
|
"""Parse a partition size like '1MB', '64K', '0x10000' into bytes."""
|
|
spec = spec.strip().upper()
|
|
for suffix, mult in _SIZE_MULT.items():
|
|
if spec.endswith(suffix + "B"):
|
|
return int(spec[: -len(suffix) - 1]) * mult
|
|
if spec.endswith(suffix):
|
|
return int(spec[: -len(suffix)]) * mult
|
|
return int(spec, 0)
|
|
|
|
|
|
def _format_size(size_bytes: int) -> str:
|
|
"""Format byte count as human-readable size."""
|
|
if size_bytes >= 1024 * 1024 and size_bytes % (1024 * 1024) == 0:
|
|
return f"{size_bytes // (1024 * 1024)}MB"
|
|
if size_bytes >= 1024 and size_bytes % 1024 == 0:
|
|
return f"{size_bytes // 1024}KB"
|
|
return f"{size_bytes}B"
|
|
|
|
|
|
class PartitionManager:
|
|
"""ESP partition table management"""
|
|
|
|
def __init__(self, app: FastMCP, config: ESPToolServerConfig) -> None:
|
|
self.app = app
|
|
self.config = config
|
|
self._register_tools()
|
|
|
|
async def _run_esptool(
|
|
self,
|
|
port: str,
|
|
args: list[str],
|
|
timeout: float = 30.0,
|
|
) -> dict[str, Any]:
|
|
"""Run esptool as an async subprocess."""
|
|
cmd = [self.config.esptool_path, "--port", port, *args]
|
|
proc = None
|
|
try:
|
|
proc = await asyncio.create_subprocess_exec(
|
|
*cmd,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
|
|
output = (stdout or b"").decode() + (stderr or b"").decode()
|
|
if proc.returncode != 0:
|
|
return {"success": False, "error": output.strip()[:500]}
|
|
return {"success": True, "output": output}
|
|
except asyncio.TimeoutError:
|
|
if proc and proc.returncode is None:
|
|
proc.kill()
|
|
await proc.wait()
|
|
return {"success": False, "error": f"Timeout after {timeout}s"}
|
|
except FileNotFoundError:
|
|
return {"success": False, "error": f"esptool not found at {self.config.esptool_path}"}
|
|
except Exception as e:
|
|
if proc and proc.returncode is None:
|
|
proc.kill()
|
|
await proc.wait()
|
|
return {"success": False, "error": str(e)}
|
|
|
|
def _register_tools(self) -> None:
|
|
"""Register partition management tools"""
|
|
|
|
@self.app.tool("esp_partition_create_ota")
|
|
async def create_ota_partition(
|
|
context: Context,
|
|
flash_size: str = "4MB",
|
|
app_size: str = "1MB",
|
|
) -> dict[str, Any]:
|
|
"""Create OTA-enabled partition table.
|
|
|
|
Generates a partition table CSV with two OTA app slots, NVS storage,
|
|
OTA data partition, and PHY calibration data. The layout follows
|
|
Espressif's recommended OTA structure.
|
|
|
|
The generated CSV can be converted to binary with gen_esp32part.py
|
|
(from ESP-IDF) and flashed to the partition table offset (typically 0x8000).
|
|
|
|
Args:
|
|
flash_size: Total flash size (e.g. "4MB", "8MB", "16MB", default: "4MB")
|
|
app_size: Size for each OTA app slot (e.g. "1MB", "1536K", default: "1MB")
|
|
"""
|
|
return await self._create_ota_impl(context, flash_size, app_size)
|
|
|
|
@self.app.tool("esp_partition_custom")
|
|
async def create_custom_partition(
|
|
context: Context,
|
|
partition_config: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
"""Create custom partition table from a configuration dict.
|
|
|
|
Accepts a list of partition entries and generates a valid ESP
|
|
partition table CSV. Each entry needs: name, type, subtype, size.
|
|
Offset is auto-calculated if omitted.
|
|
|
|
Example partition_config:
|
|
{
|
|
"partitions": [
|
|
{"name": "nvs", "type": "data", "subtype": "nvs", "size": "24K"},
|
|
{"name": "factory", "type": "app", "subtype": "factory", "size": "1MB"},
|
|
{"name": "storage", "type": "data", "subtype": "spiffs", "size": "512K"}
|
|
]
|
|
}
|
|
|
|
Args:
|
|
partition_config: Dict with "partitions" key containing list of entries
|
|
"""
|
|
return await self._create_custom_impl(context, partition_config)
|
|
|
|
@self.app.tool("esp_partition_analyze")
|
|
async def analyze_partitions(
|
|
context: Context,
|
|
port: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Analyze current partition table on a connected ESP device.
|
|
|
|
Reads the partition table from flash (at offset 0x8000, 0xC00 bytes)
|
|
and parses the binary format into a human-readable table. Shows
|
|
partition names, types, offsets, sizes, and flags.
|
|
|
|
Works with physical devices and QEMU virtual devices.
|
|
|
|
Args:
|
|
port: Serial port or socket:// URI (required)
|
|
"""
|
|
return await self._analyze_impl(context, port)
|
|
|
|
async def _create_ota_impl(
|
|
self,
|
|
context: Context,
|
|
flash_size: str,
|
|
app_size: str,
|
|
) -> dict[str, Any]:
|
|
"""Generate an OTA-capable partition table."""
|
|
|
|
try:
|
|
total_bytes = _parse_size_spec(flash_size)
|
|
app_bytes = _parse_size_spec(app_size)
|
|
except (ValueError, TypeError) as e:
|
|
return {"success": False, "error": f"Invalid size: {e}"}
|
|
|
|
# Standard layout:
|
|
# 0x9000 - nvs (24KB)
|
|
# 0xf000 - otadata (8KB)
|
|
# 0x11000 - phy_init (4KB)
|
|
# 0x12000 - ota_0 (app_size)
|
|
# ota_0 + app_size - ota_1 (app_size)
|
|
|
|
nvs_size = 24 * 1024
|
|
otadata_size = 8 * 1024
|
|
phy_size = 4 * 1024
|
|
|
|
# Check it all fits (partition table at 0x8000 + 0x1000)
|
|
overhead = 0x9000 + nvs_size + otadata_size + phy_size # Before first app
|
|
needed = overhead + (2 * app_bytes)
|
|
if needed > total_bytes:
|
|
return {
|
|
"success": False,
|
|
"error": (
|
|
f"Layout requires {_format_size(needed)} but flash is {flash_size}. "
|
|
f"Reduce app_size or increase flash_size."
|
|
),
|
|
}
|
|
|
|
partitions = [
|
|
("nvs", "data", "nvs", "0x9000", _format_size(nvs_size)),
|
|
("otadata", "data", "ota", f"0x{0x9000 + nvs_size:x}", _format_size(otadata_size)),
|
|
("phy_init", "data", "phy", f"0x{0x9000 + nvs_size + otadata_size:x}", _format_size(phy_size)),
|
|
("ota_0", "app", "ota_0", f"0x{overhead:x}", _format_size(app_bytes)),
|
|
("ota_1", "app", "ota_1", f"0x{overhead + app_bytes:x}", _format_size(app_bytes)),
|
|
]
|
|
|
|
# Remaining space for storage
|
|
used = overhead + (2 * app_bytes)
|
|
remaining = total_bytes - used
|
|
if remaining >= 4096:
|
|
partitions.append(
|
|
("storage", "data", "spiffs", f"0x{used:x}", _format_size(remaining))
|
|
)
|
|
|
|
# Generate CSV
|
|
csv_lines = ["# ESP-IDF Partition Table (OTA layout)", "# Name, Type, SubType, Offset, Size, Flags"]
|
|
for name, ptype, subtype, offset, size in partitions:
|
|
csv_lines.append(f"{name}, {ptype}, {subtype}, {offset}, {size},")
|
|
|
|
csv_text = "\n".join(csv_lines) + "\n"
|
|
|
|
return {
|
|
"success": True,
|
|
"flash_size": flash_size,
|
|
"app_size": app_size,
|
|
"partition_csv": csv_text,
|
|
"partitions": [
|
|
{"name": p[0], "type": p[1], "subtype": p[2], "offset": p[3], "size": p[4]}
|
|
for p in partitions
|
|
],
|
|
"space_remaining": _format_size(remaining) if remaining >= 4096 else "0",
|
|
"note": "Flash this CSV with gen_esp32part.py to binary, then write to 0x8000",
|
|
}
|
|
|
|
async def _create_custom_impl(
|
|
self,
|
|
context: Context,
|
|
partition_config: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
"""Generate a custom partition table from config."""
|
|
|
|
partitions_input = partition_config.get("partitions", [])
|
|
if not partitions_input:
|
|
return {"success": False, "error": "partition_config must have a 'partitions' list"}
|
|
|
|
# Auto-calculate offsets starting after partition table (0x9000)
|
|
current_offset = 0x9000
|
|
partitions = []
|
|
errors = []
|
|
|
|
for i, entry in enumerate(partitions_input):
|
|
name = entry.get("name")
|
|
ptype = entry.get("type")
|
|
subtype = entry.get("subtype")
|
|
size_str = entry.get("size")
|
|
|
|
if not all([name, ptype, subtype, size_str]):
|
|
errors.append(f"Partition {i}: requires name, type, subtype, size")
|
|
continue
|
|
|
|
# Validate type
|
|
if ptype not in PARTITION_TYPES:
|
|
errors.append(f"Partition '{name}': invalid type '{ptype}' (use: {list(PARTITION_TYPES.keys())})")
|
|
continue
|
|
|
|
# Validate subtype
|
|
valid_subtypes = APP_SUBTYPES if ptype == "app" else DATA_SUBTYPES
|
|
if subtype not in valid_subtypes:
|
|
errors.append(f"Partition '{name}': invalid subtype '{subtype}' (use: {list(valid_subtypes.keys())})")
|
|
continue
|
|
|
|
try:
|
|
size_bytes = _parse_size_spec(size_str)
|
|
except (ValueError, TypeError):
|
|
errors.append(f"Partition '{name}': invalid size '{size_str}'")
|
|
continue
|
|
|
|
# Use explicit offset if provided, otherwise auto-calculate
|
|
offset = entry.get("offset")
|
|
if offset:
|
|
current_offset = int(offset, 0) if isinstance(offset, str) else offset
|
|
|
|
# App partitions must be 64KB aligned
|
|
if ptype == "app" and current_offset % 0x10000 != 0:
|
|
current_offset = (current_offset + 0xFFFF) & ~0xFFFF
|
|
|
|
partitions.append({
|
|
"name": name,
|
|
"type": ptype,
|
|
"subtype": subtype,
|
|
"offset": f"0x{current_offset:x}",
|
|
"size": _format_size(size_bytes),
|
|
"size_bytes": size_bytes,
|
|
})
|
|
|
|
current_offset += size_bytes
|
|
|
|
if errors:
|
|
return {"success": False, "errors": errors}
|
|
|
|
# Generate CSV
|
|
csv_lines = ["# ESP-IDF Partition Table (custom layout)", "# Name, Type, SubType, Offset, Size, Flags"]
|
|
for p in partitions:
|
|
csv_lines.append(f"{p['name']}, {p['type']}, {p['subtype']}, {p['offset']}, {p['size']},")
|
|
|
|
csv_text = "\n".join(csv_lines) + "\n"
|
|
|
|
return {
|
|
"success": True,
|
|
"partition_csv": csv_text,
|
|
"partitions": [{k: v for k, v in p.items() if k != "size_bytes"} for p in partitions],
|
|
"total_size": _format_size(sum(p["size_bytes"] for p in partitions)),
|
|
"note": "Flash this CSV with gen_esp32part.py to binary, then write to 0x8000",
|
|
}
|
|
|
|
async def _analyze_impl(
|
|
self,
|
|
context: Context,
|
|
port: str | None,
|
|
) -> dict[str, Any]:
|
|
"""Read and parse partition table from a connected device."""
|
|
|
|
if not port:
|
|
return {"success": False, "error": "Port is required for partition analysis"}
|
|
|
|
import tempfile
|
|
|
|
# Partition table is at 0x8000, max size 0xC00 (3KB)
|
|
with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as tmp:
|
|
tmp_path = tmp.name
|
|
|
|
try:
|
|
result = await self._run_esptool(
|
|
port,
|
|
["read-flash", "0x8000", "0xC00", tmp_path],
|
|
timeout=60.0,
|
|
)
|
|
|
|
if not result["success"]:
|
|
return {"success": False, "error": result["error"], "port": port}
|
|
|
|
raw = Path(tmp_path).read_bytes()
|
|
partitions = self._parse_partition_table_binary(raw)
|
|
|
|
if not partitions:
|
|
return {
|
|
"success": True,
|
|
"port": port,
|
|
"partitions": [],
|
|
"note": "No valid partition entries found (flash may be blank or erased)",
|
|
}
|
|
|
|
return {
|
|
"success": True,
|
|
"port": port,
|
|
"partition_count": len(partitions),
|
|
"partitions": partitions,
|
|
}
|
|
|
|
finally:
|
|
import os
|
|
try:
|
|
os.unlink(tmp_path)
|
|
except OSError:
|
|
pass
|
|
|
|
def _parse_partition_table_binary(self, raw: bytes) -> list[dict[str, Any]]:
|
|
"""Parse ESP32 binary partition table format.
|
|
|
|
Each entry is 32 bytes:
|
|
- 2 bytes: magic (0xAA50)
|
|
- 1 byte: type
|
|
- 1 byte: subtype
|
|
- 4 bytes: offset (LE)
|
|
- 4 bytes: size (LE)
|
|
- 16 bytes: name (null-terminated)
|
|
- 4 bytes: flags
|
|
"""
|
|
import struct
|
|
|
|
entry_size = 32
|
|
magic_expected = 0x50AA # Little-endian
|
|
|
|
# Reverse lookup tables
|
|
type_names = {v: k for k, v in PARTITION_TYPES.items()}
|
|
app_subtype_names = {v: k for k, v in APP_SUBTYPES.items()}
|
|
data_subtype_names = {v: k for k, v in DATA_SUBTYPES.items()}
|
|
|
|
partitions = []
|
|
for i in range(0, len(raw) - entry_size + 1, entry_size):
|
|
entry = raw[i : i + entry_size]
|
|
magic = struct.unpack_from("<H", entry, 0)[0]
|
|
|
|
if magic == 0xFFFF:
|
|
# End of table (erased flash)
|
|
break
|
|
if magic != magic_expected:
|
|
continue
|
|
|
|
ptype = entry[2]
|
|
subtype = entry[3]
|
|
offset = struct.unpack_from("<I", entry, 4)[0]
|
|
size = struct.unpack_from("<I", entry, 8)[0]
|
|
name = entry[12:28].split(b"\x00")[0].decode("ascii", errors="replace")
|
|
flags = struct.unpack_from("<I", entry, 28)[0]
|
|
|
|
type_name = type_names.get(ptype, f"0x{ptype:02x}")
|
|
if ptype == 0x00:
|
|
subtype_name = app_subtype_names.get(subtype, f"0x{subtype:02x}")
|
|
elif ptype == 0x01:
|
|
subtype_name = data_subtype_names.get(subtype, f"0x{subtype:02x}")
|
|
else:
|
|
subtype_name = f"0x{subtype:02x}"
|
|
|
|
partitions.append({
|
|
"name": name,
|
|
"type": type_name,
|
|
"subtype": subtype_name,
|
|
"offset": f"0x{offset:x}",
|
|
"size": _format_size(size),
|
|
"size_bytes": size,
|
|
"encrypted": bool(flags & 1),
|
|
})
|
|
|
|
return partitions
|
|
|
|
async def health_check(self) -> dict[str, Any]:
|
|
"""Component health check"""
|
|
return {"status": "healthy", "note": "Partition manager ready"}
|