rescue_upload.py

#!/usr/bin/env python3
"""
python rescue_upload.py --host 192.168.1.79 --password 'ota-password'

It reuses the protocol helpers from esphome.espota2 (2026.3.1) and adds knobs
the CLI does not expose:

  --chunk-size      TCP write size (stock client: 8192)
  --chunk-delay-ms  pause after each device ACK, so the ESP8266 can feed its
                    watchdog between flash sector erases
  --no-compress     offer no feature bits; device answers HEADER_OK and the
                    image is streamed uncompressed (no gzip path at all)
  --no-sha256       do not offer SHA256 auth (device falls back to MD5)
  --timeout         per-receive timeout (stock client uses 90 s)

The device ACKs every OTA_BLOCK_SIZE = 8192 received bytes regardless of
--chunk-size (esphome/components/esphome/ota/ota_esphome.cpp), so this script
waits for one ACK per 8192 bytes sent.
"""

from __future__ import annotations

import argparse
import gzip
import hashlib
import secrets
import socket
import sys
import time
from pathlib import Path

from esphome.espota2 import (
    FEATURE_SUPPORTS_COMPRESSION,
    FEATURE_SUPPORTS_SHA256_AUTH,
    MAGIC_BYTES,
    OTA_VERSION_1_0,
    OTA_VERSION_2_0,
    RESPONSE_AUTH_OK,
    RESPONSE_BIN_MD5_OK,
    RESPONSE_CHUNK_OK,
    RESPONSE_OK,
    RESPONSE_RECEIVE_OK,
    RESPONSE_REQUEST_AUTH,
    RESPONSE_REQUEST_SHA256_AUTH,
    RESPONSE_SUPPORTS_COMPRESSION,
    RESPONSE_UPDATE_END_OK,
    RESPONSE_UPDATE_PREPARE_OK,
    OTAError,
    _AUTH_METHODS,
    receive_exactly,
    send_check,
)

DEVICE_ACK_BLOCK = 8192  # OTA_BLOCK_SIZE in ota_esphome.cpp


def find_firmware() -> Path:
    candidates = sorted(
        Path(".esphome/build").glob("*/.pioenvs/*/firmware.bin"),
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )
    if not candidates:
        raise SystemExit("No firmware.bin under .esphome/build - pass --file")
    return candidates[0]


def authenticate(sock: socket.socket, password: str | None) -> None:
    (auth,) = receive_exactly(
        sock,
        1,
        "auth",
        [RESPONSE_REQUEST_AUTH, RESPONSE_REQUEST_SHA256_AUTH, RESPONSE_AUTH_OK],
    )
    if auth == RESPONSE_AUTH_OK:
        return
    if password is None:
        raise OTAError("ESP requests password, but no password given!")
    hash_func, nonce_size, hash_name = _AUTH_METHODS[auth]
    nonce = receive_exactly(
        sock, nonce_size, f"{hash_name} authentication nonce", [], decode=False
    ).decode()
    cnonce = secrets.token_hex(nonce_size // 2)
    send_check(sock, cnonce, "auth cnonce")
    hasher = hash_func()
    hasher.update(password.encode("utf-8"))
    hasher.update(nonce.encode())
    hasher.update(cnonce.encode())
    send_check(sock, hasher.hexdigest(), "auth result")
    receive_exactly(sock, 1, "auth result", RESPONSE_AUTH_OK)


def upload(args: argparse.Namespace) -> int:
    contents = args.file.read_bytes()
    sock = socket.create_connection((args.host, args.port), timeout=20.0)
    try:
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        send_check(sock, MAGIC_BYTES, "magic bytes")

        _, version = receive_exactly(sock, 2, "version", RESPONSE_OK)
        print(f"[i] device reports OTA protocol version {version}")
        if version not in (OTA_VERSION_1_0, OTA_VERSION_2_0):
            raise OTAError(f"unsupported device OTA version {version}")

        features = 0
        if not args.no_sha256:
            features |= FEATURE_SUPPORTS_SHA256_AUTH
        if not args.no_compress:
            features |= FEATURE_SUPPORTS_COMPRESSION
        send_check(sock, features, "features")
        response = receive_exactly(sock, 1, "features", None)[0]
        if response == RESPONSE_SUPPORTS_COMPRESSION:
            payload = gzip.compress(contents, compresslevel=9)
            print(f"[i] gzip: {len(contents)} -> {len(payload)} bytes")
        else:
            payload = contents
            print(f"[i] uncompressed: {len(payload)} bytes")

        authenticate(sock, args.password)

        size = len(payload)
        send_check(
            sock,
            [
                (size >> 24) & 0xFF,
                (size >> 16) & 0xFF,
                (size >> 8) & 0xFF,
                size & 0xFF,
            ],
            "binary size",
        )
        receive_exactly(sock, 1, "binary size", RESPONSE_UPDATE_PREPARE_OK)

        send_check(sock, hashlib.md5(payload).hexdigest(), "file checksum")
        receive_exactly(sock, 1, "file checksum", RESPONSE_BIN_MD5_OK)

        sock.settimeout(args.timeout)
        delay = args.chunk_delay_ms / 1000.0
        step = max(size // 10, 1)
        next_report = step
        sent = 0
        acked = 0
        while sent < size:
            chunk = payload[sent : sent + args.chunk_size]
            sock.sendall(chunk)
            sent += len(chunk)

            if version >= OTA_VERSION_2_0:
                while acked + DEVICE_ACK_BLOCK <= sent:
                    receive_exactly(sock, 1, "chunk OK", RESPONSE_CHUNK_OK)
                    acked += DEVICE_ACK_BLOCK
                    if delay:
                        time.sleep(delay)
            elif delay:
                time.sleep(delay)

            if sent >= next_report:
                print(f"[i] {sent * 100 // size}% ({sent}/{size})")
                next_report += step

        if version >= OTA_VERSION_2_0 and acked < size:
            receive_exactly(sock, 1, "chunk OK", RESPONSE_CHUNK_OK)

        receive_exactly(sock, 1, "receive OK", RESPONSE_RECEIVE_OK)
        receive_exactly(sock, 1, "Update end", RESPONSE_UPDATE_END_OK)
        send_check(sock, RESPONSE_OK, "end acknowledgement")
        print("[i] OTA successful")
        return 0
    except OTAError as err:
        print(f"[!] {err}", file=sys.stderr)
        return 1
    finally:
        sock.close()


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--host", required=True, help="device IP, e.g. 192.168.1.79")
    parser.add_argument("--port", type=int, default=8266)
    parser.add_argument("--password", default=None, help="OTA password, if any")
    parser.add_argument("--file", type=Path, default=None, help="firmware.bin")
    parser.add_argument("--chunk-size", type=int, default=2048)
    parser.add_argument("--chunk-delay-ms", type=float, default=5.0)
    parser.add_argument("--timeout", type=float, default=180.0)
    parser.add_argument("--no-compress", action="store_true")
    parser.add_argument("--no-sha256", action="store_true")
    args = parser.parse_args()
    if args.file is None:
        args.file = find_firmware()
    print(f"[i] flashing {args.file} to {args.host}:{args.port}")
    return upload(args)


if __name__ == "__main__":
    raise SystemExit(main())