mirror of
https://github.com/emsesp/EMS-ESP32.git
synced 2026-10-28 11:14:58 +00:00
232 lines
7.4 KiB
Python
232 lines
7.4 KiB
Python
|
||
# Modified from https://github.com/ayushsharma82/ElegantOTA
|
||
#
|
||
# This is called during the PlatformIO upload process, when the target is 'upload'.
|
||
# Use the file upload_cli.py for manual uploads outside PIO.
|
||
#
|
||
# To use create a pio_local.ini file in the project root and add the following:
|
||
# [env]
|
||
# upload_protocol = custom
|
||
# custom_emsesp_ip = ems-esp.local
|
||
# custom_username = admin
|
||
# custom_password = admin
|
||
#
|
||
# and
|
||
# extra_scripts = scripts/upload.py
|
||
#
|
||
# Prefer a numeric IP over .local — mDNS (especially under WSL2) adds latency
|
||
# and can stall on IPv6 before falling back to IPv4.
|
||
|
||
import hashlib
|
||
import os
|
||
import socket
|
||
import time
|
||
from urllib.parse import urlparse
|
||
|
||
import requests
|
||
from requests.adapters import HTTPAdapter
|
||
|
||
Import("env")
|
||
|
||
try:
|
||
from requests_toolbelt import MultipartEncoder, MultipartEncoderMonitor
|
||
from tqdm import tqdm
|
||
from termcolor import cprint
|
||
except ImportError:
|
||
env.Execute("$PYTHONEXE -m pip install requests_toolbelt")
|
||
env.Execute("$PYTHONEXE -m pip install tqdm")
|
||
env.Execute("$PYTHONEXE -m pip install termcolor")
|
||
from requests_toolbelt import MultipartEncoder, MultipartEncoderMonitor
|
||
from tqdm import tqdm
|
||
from termcolor import cprint
|
||
|
||
# http.client/urllib3 default is 8–16 KiB, which starves a LAN TCP window.
|
||
UPLOAD_BLOCK_SIZE = 64 * 1024
|
||
UPLOAD_SNDBUF = 256 * 1024
|
||
|
||
|
||
def print_success(x):
|
||
cprint(x, 'green')
|
||
|
||
|
||
def print_fail(x):
|
||
cprint(f'Error: {x}', 'red')
|
||
|
||
|
||
def resolve_ipv4(host):
|
||
"""Resolve hostname to IPv4 so we skip mDNS/IPv6 retries on every request."""
|
||
if host.startswith('['):
|
||
return host
|
||
if ':' in host and host.count(':') == 1:
|
||
host = host.rsplit(':', 1)[0]
|
||
try:
|
||
socket.inet_pton(socket.AF_INET, host)
|
||
return host
|
||
except OSError:
|
||
pass
|
||
infos = socket.getaddrinfo(host, 80, socket.AF_INET, socket.SOCK_STREAM)
|
||
if not infos:
|
||
raise socket.gaierror(f'No IPv4 address for {host}')
|
||
return infos[0][4][0]
|
||
|
||
|
||
class FastHTTPAdapter(HTTPAdapter):
|
||
"""Larger write chunks and socket send buffer for streaming POSTs."""
|
||
|
||
def init_poolmanager(self, connections, maxsize, block=False, **pool_kwargs):
|
||
socket_options = [
|
||
(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1),
|
||
(socket.SOL_SOCKET, socket.SO_SNDBUF, UPLOAD_SNDBUF),
|
||
]
|
||
pool_kwargs.setdefault('socket_options', socket_options)
|
||
pool_kwargs.setdefault('blocksize', UPLOAD_BLOCK_SIZE)
|
||
super().init_poolmanager(connections, maxsize, block=block, **pool_kwargs)
|
||
|
||
|
||
def make_session():
|
||
session = requests.Session()
|
||
adapter = FastHTTPAdapter(pool_connections=2, pool_maxsize=2)
|
||
session.mount('http://', adapter)
|
||
return session
|
||
|
||
|
||
def build_headers(host_ip, emsesp_url, content_type='application/json', access_token=None, extra_headers=None):
|
||
"""Build common HTTP headers with optional overrides."""
|
||
headers = {
|
||
'Host': host_ip,
|
||
'User-Agent': 'Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:109.0) Gecko/20100101 Firefox/118.0',
|
||
'Accept': '*/*',
|
||
'Accept-Language': 'en-US',
|
||
'Referer': emsesp_url,
|
||
'Content-Type': content_type,
|
||
'Connection': 'keep-alive'
|
||
}
|
||
|
||
if access_token:
|
||
headers['Authorization'] = f'Bearer {access_token}'
|
||
|
||
if extra_headers:
|
||
headers.update(extra_headers)
|
||
|
||
return headers
|
||
|
||
|
||
def on_upload(source, target, env):
|
||
|
||
# make sure we have set the upload_protocol to custom
|
||
if env.get('UPLOAD_PROTOCOL') != 'custom':
|
||
print_fail(
|
||
"Please set upload_protocol = custom in your pio_local.ini file when using upload.py")
|
||
return
|
||
|
||
# first check authentication
|
||
try:
|
||
username = env.GetProjectOption('custom_username')
|
||
password = env.GetProjectOption('custom_password')
|
||
emsesp_ip = env.GetProjectOption('custom_emsesp_ip')
|
||
except Exception:
|
||
print_fail(f'Missing settings. Add these to your pio_local.ini file:\n\ncustom_username=username\ncustom_password=password\ncustom_emsesp_ip=ems-esp.local\n')
|
||
return
|
||
|
||
parsed_url = urlparse(f"http://{emsesp_ip}")
|
||
host_name = parsed_url.hostname or emsesp_ip
|
||
try:
|
||
host_ip = resolve_ipv4(host_name)
|
||
except OSError as e:
|
||
print_fail(f"Could not resolve {host_name} to IPv4: {e}")
|
||
return
|
||
|
||
if host_ip != host_name:
|
||
print_success(f"Resolved {host_name} -> {host_ip}")
|
||
|
||
emsesp_url = f"http://{host_ip}"
|
||
session = make_session()
|
||
|
||
try:
|
||
signon_url = f"{emsesp_url}/rest/signIn"
|
||
signon_headers = build_headers(host_ip, emsesp_url)
|
||
|
||
response = session.post(
|
||
signon_url,
|
||
json={"username": username, "password": password},
|
||
headers=signon_headers,
|
||
timeout=15)
|
||
|
||
if response.status_code != 200:
|
||
print_fail("Authentication with EMS-ESP failed (code " +
|
||
str(response.status_code) + ")")
|
||
return
|
||
|
||
print_success("Authentication with EMS-ESP successful")
|
||
access_token = response.json().get('access_token')
|
||
|
||
firmware_path = str(source[0])
|
||
with open(firmware_path, 'rb') as firmware_file:
|
||
firmware = firmware_file.read()
|
||
|
||
md5 = hashlib.md5(firmware).hexdigest()
|
||
filename = os.path.basename(firmware_path)
|
||
|
||
encoder = MultipartEncoder(fields={
|
||
'MD5': md5,
|
||
'file': (filename, firmware, 'application/octet-stream')}
|
||
)
|
||
|
||
bar = tqdm(desc='Upload Progress',
|
||
total=encoder.len,
|
||
dynamic_ncols=True,
|
||
unit='B',
|
||
unit_scale=True,
|
||
unit_divisor=1024,
|
||
mininterval=0.2
|
||
)
|
||
|
||
monitor = MultipartEncoderMonitor(
|
||
encoder, lambda monitor: bar.update(monitor.bytes_read - bar.n))
|
||
|
||
post_headers = build_headers(
|
||
host_ip,
|
||
emsesp_url,
|
||
content_type=monitor.content_type,
|
||
access_token=access_token,
|
||
extra_headers={
|
||
'Content-Length': str(monitor.len),
|
||
'Origin': emsesp_url
|
||
}
|
||
)
|
||
|
||
upload_url = f"{emsesp_url}/rest/uploadFile"
|
||
started = time.perf_counter()
|
||
response = session.post(
|
||
upload_url, data=monitor, headers=post_headers, timeout=None)
|
||
elapsed = time.perf_counter() - started
|
||
|
||
bar.close()
|
||
print()
|
||
|
||
if response.status_code != 200:
|
||
print_fail(f"Upload failed (code {response.status_code}).")
|
||
else:
|
||
kibs = (encoder.len / 1024) / elapsed if elapsed > 0 else 0
|
||
print_success(
|
||
f"Upload successful ({kibs:.0f} KiB/s in {elapsed:.1f}s). Rebooting device.")
|
||
restart_headers = build_headers(
|
||
host_ip, emsesp_url, access_token=access_token)
|
||
restart_url = f"{emsesp_url}/api/system/restart"
|
||
try:
|
||
response = session.get(
|
||
restart_url, headers=restart_headers, timeout=10)
|
||
if response.status_code != 200:
|
||
print_fail(f"Restart failed (code {response.status_code})")
|
||
except requests.RequestException:
|
||
# Device often drops the connection as it reboots — that's OK.
|
||
pass
|
||
|
||
print()
|
||
finally:
|
||
session.close()
|
||
|
||
|
||
if env.get('UPLOAD_PROTOCOL') == 'custom':
|
||
env.Replace(UPLOADCMD=on_upload)
|