wasp-os/tools/wasptool

489 lines
14 KiB
Python
Executable File

#!/usr/bin/env python3
# SPDX-License-Identifier: LGPL-3.0-or-later
# Copyright (c) 2020 Daniel Thompson
import argparse
import io
import random
import os
import pexpect
import time
import types
import string
import subprocess
import sys
def draw_pbar(percent, quiet=False, end='\r'):
if not quiet:
if percent > 100:
percent = 100
bar = int(percent) // 2
print(f'[{"#"*bar}{"."*(50-bar)}] {round(percent, 1)}% ', end=end, flush=True)
def pbar(iterable, quiet=False):
step = 100 / len(iterable)
for i, v in enumerate(iterable):
draw_pbar(step * i, quiet)
yield v
if not quiet:
draw_pbar(100, quiet, None)
def run_command(c, cmd):
"""Cheap and cheerful command wrapper.
This differs from REPLWrapper because it assumes the remote end will
echo the characters... and that we must eat them before handing over
the results to the caller.
"""
c.sendline(cmd)
c.expect_exact(cmd)
c.expect('>>> ')
return c.before.replace('\r\r\n', '\n').strip('\n')
def sync(c):
"""Stop the watch and synchronize with the command prompt.
Sending a random print ensure the final export (of the prompt)
does not accidentally match a previously issued prompt.
"""
verbose = bool(c.logfile)
tag = ''.join([random.choice(string.ascii_uppercase) for i in range(6)])
try:
if not verbose:
c.logfile = io.StringIO()
c.send('\x03')
c.expect('>>> ')
c.sendline(f'print("{tag[:3]}""{tag[3:]}")')
c.expect(tag)
c.expect('>>> ')
if not verbose:
c.logfile.close()
c.logfile = None
except pexpect.exceptions.EOF:
print("ERROR: Cannot sync with device")
print_log(c.logfile)
sys.exit(1)
def unsync(c):
"""Set the watch running again.
There must be an expect (or a sleep) since if we kill the subordinate
process too early then the sendline will not have completed.
"""
c.sendline('wasp.system.run()')
c.expect(['Watch is running, use Ctrl-C to stop',
'Watch already running in the background'])
c.send('\x18')
def paste(c, f, verbose=False, chunk=None):
docstring = False
tosend = []
for ln in f.readlines():
ln = ln.rstrip()
# This is a bit loose (definitely not PEP-257 compliant) but
# is enough for most code.
if ln.lstrip().startswith('"""'):
docstring = True
if docstring:
if ln.rstrip().endswith('"""'):
docstring = False
continue
if ln.lstrip().startswith('#'):
continue
if ln.strip() == '':
continue
tosend.append(ln)
for ln in pbar(tosend, verbose):
if chunk and ln.startswith('class'):
chunk()
c.sendline(ln)
choice = c.expect(['=== ', 'FATAL: uncaught exception [0-9a-f\r]*\n'])
if choice == 1:
# Capture and display the error message, then exit
if not verbose:
print('\n~~~')
while choice == 1:
if not verbose:
print(c.match.group(0).rstrip(), file=sys.stderr)
choice = c.expect([pexpect.TIMEOUT, '.*\n'], timeout=2)
if not verbose:
print('~~~')
print('\nPlease reboot your device', file=sys.stderr)
sys.exit(16)
def print_log(logfile):
try:
lines = logfile.getvalue().split('\n')
except:
lines = logfile.split('\n')
lines = [ l.strip('\x04\r') for l in lines ]
output = [ l for l in lines if l and l != '>>> ' ]
if output:
print('~~~')
print('\n'.join(output))
print('~~~')
def handle_eval(c, cmd):
verbose = bool(c.logfile)
c.send('\x05')
c.expect('=== ')
c.sendline(cmd)
c.expect('=== ')
if not verbose:
c.logfile = io.StringIO()
c.send('\x04')
c.expect('>>> ')
if not verbose:
print_log(c.logfile)
c.logfile.close()
c.logfile = None
def handle_exec(c, fname):
verbose = bool(c.logfile)
log = io.StringIO()
def chunk():
if not verbose:
c.logfile = log
c.send('\x04')
c.expect('>>> ')
if not verbose:
c.logfile = None
c.send('\x05')
c.expect('=== ')
with open(fname) as f:
if not verbose:
print(f'Preparing to run {fname}:')
c.send('\x05')
c.expect('=== ')
paste(c, f, verbose, chunk)
if not verbose:
c.logfile = log
c.send('\x04')
c.expect('>>> ')
if not verbose:
c.logfile = None
print_log(log)
log.close()
def handle_battery_level(c):
print(f'Battery: {c.run_command("watch.battery.level()")}%')
def handle_memory_free(c):
before_gc = c.run_command("wasp.gc.mem_free()")
c.run_command('wasp.gc.collect()')
after_gc = c.run_command("wasp.gc.mem_free()")
p = subprocess.run(["git", "describe", "--tags"], capture_output=True)
version = p.stdout.decode('UTF-8').strip()
print(f'"{version}",{before_gc},{after_gc}')
def handle_reset(c, ota=False):
cmd = 'reset'
if ota:
cmd = 'enter_ota_dfu'
c.send('\x05')
c.expect('=== ')
c.sendline('import machine')
c.expect('=== ')
c.sendline(f'machine.{cmd}()')
c.expect('=== ')
c.send('\x04')
def handle_rtc(c):
# Wait for the clock to tick over to the next second
now = then = time.localtime()
while now[5] == then[5]:
now = time.localtime()
# Set the time
c.sendline(f'watch.rtc.set_localtime(({now[0]}, {now[1]}, {now[2]}, {now[3]}, {now[4]}, {now[5]}, {now[6]}, {now[7]}))')
c.expect('>>> ')
def check_rtc(c):
c.sendline('print(watch.rtc.get_localtime())')
c.expect(r'\(([0-9]+), ([0-9]+), ([0-9]+), ([0-9]+), ([0-9]+), ([0-9]+), ([0-9]+), ([0-9]+)\)')
t = time.localtime()
watch_hms = (int(c.match[4]), int(c.match[5]), int(c.match[6]))
watch_str = f'{watch_hms[0]:02d}:{watch_hms[1]:02d}:{watch_hms[2]:02d}'
host_hms = (t.tm_hour, t.tm_min, t.tm_sec)
host_str = f'{host_hms[0]:02d}:{host_hms[1]:02d}:{host_hms[2]:02d}'
delta = 3600 * (host_hms[0] - watch_hms[0]) + \
60 * (host_hms[1] - watch_hms[1]) + \
(host_hms[2] - watch_hms[2])
print(f"PC <-> watch: {watch_str} <-> {host_str} (delta {delta})")
c.expect('>>> ')
def handle_binary_download(c, tname, fname):
verbose = bool(c.logfile)
if tname.startswith("/"):
fname = "." + fname
c.run_command('import os')
stat = c.run_command(f'os.stat("{tname}")[6]')
if 'Error' in stat:
print('Watch reported error:')
print(stat)
return
print(f'Downloading {fname}:')
c.run_command(f'f = open("{tname}", "rb")')
sz = eval(stat)
bytes_read = 0
if os.path.exists(fname):
if args.force_pull:
print(f"Will overwrite local file '{fname}'")
else:
raise Exception(f"Local file '{fname}' already exists, use --force argument to overwrite it")
if "/" in fname: # create local directories if needed
pile = "./"
for local_dir in fname.split("/")[:-1]:
if local_dir == ".":
continue
try:
os.mkdir(f"{pile}/{local_dir}/")
print(f"Created local dir {pile}/{local_dir}/")
except FileExistsError: # folder already exists
pass
pile += local_dir + "/"
with open(fname, 'wb') as f:
while True:
draw_pbar(100 * bytes_read / sz, verbose)
reply = c.run_command('repr(f.read(24))')
reply = reply.replace('\\\\', '\\')
if reply.startswith('"'):
# "b'..CONTENT..'"
reply = reply[1:-1]
elif reply.startswith("'"):
# 'b\'..CONTENT..\''
reply = reply[1:-1].replace("\\'", "'")
data = eval(reply)
if len(data) == 0:
break
bytes_read += len(data)
f.write(data)
draw_pbar(100, verbose, end=None)
c.run_command('f.close()')
# Release as much memory as possible
c.run_command('del f')
c.run_command('del os')
def handle_binary_upload(c, fname, tname):
verbose = bool(c.logfile)
if not tname:
tname = os.path.basename(fname)
else:
dname = os.path.dirname(tname)
if dname:
c.run_command('import os')
c.sendline(f'os.mkdir("{dname}")')
c.run_command('del os')
c.run_command(f'f = open("{tname}", "wb")')
# Absorb the file to be uploaded
with open(fname, 'rb') as f:
data = f.read()
chunksz = 24
nchunks = len(data) // chunksz
lastchunk = len(data) % chunksz
if not verbose:
print(f'Uploading {fname}:')
# Send the data
for i in pbar(range(0, chunksz*nchunks, chunksz), verbose):
c.run_command(f'f.write({repr(data[i:i+chunksz])})')
if lastchunk:
c.run_command(f'f.write({repr(data[-lastchunk:])})')
c.run_command('f.close()')
def handle_upload(c, fname, tname):
verbose = bool(c.logfile)
if not tname:
tname = os.path.basename(fname)
else:
dname = os.path.dirname(tname)
if dname:
c.run_command('import os')
c.sendline(f'os.mkdir("{dname}")')
c.run_command('del os')
c.run_command('from shell import upload')
with open(fname) as f:
if not verbose:
print(f'Uploading {fname}:')
c.sendline(f'upload("{tname}")')
c.expect('=== ')
paste(c, f, verbose)
c.send('\x04')
c.expect('>>> ')
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description='Wasp-os command and control client')
parser.add_argument('--as', dest='upload_as', default=None,
help="Filename to use on the target (e.g. wasptool --upload docs/main/chrono.py --as main.py")
parser.add_argument('--battery', action='store_true',
help="Report remaining battery charge")
parser.add_argument('--bootloader', action='store_true',
help="Reboot into the bootloader mode for OTA update")
parser.add_argument('--binary', action='store_true',
help="Enable non-ASCII mode for upload command")
parser.add_argument('--console', action='store_true',
help='Launch a REPL session')
parser.add_argument('--check-rtc', action='store_true',
help='Compare workstation and watch times')
parser.add_argument('--device',
help='Connect only to a specific named device (or MAC address)')
parser.add_argument('--exec',
help='Execute the contents of a file')
parser.add_argument('--eval',
help='Execute the provided python string')
parser.add_argument('--memfree', action='store_true',
help='Report on the current memory usage.')
parser.add_argument('--pull',
help='Fetch a file from the target')
parser.add_argument('--force', action='store_true', dest='force_pull',
help='Force pull to overwrite local file if a file with the same name already exists')
parser.add_argument('--push',
help='Push a file to the target')
parser.add_argument('--reset', action='store_true',
help="Reboot the device (and don't stay in bootloader mode)")
parser.add_argument('--ota',
help="Deliver an OTA update to the device")
parser.add_argument('--rtc', action='store_true',
help='Set the time on the wasp-os device')
parser.add_argument('--upload',
help='Copy the specified file to the wasp-os device')
parser.add_argument('--verbose', action='store_true',
help='Log interaction with the wasp-os device')
args = parser.parse_args()
if args.device:
if ':' in args.device:
device_args = ' --address ' + args.device
else:
device_args = ' --device ' + args.device
else:
device_args = ''
pynus = os.path.dirname(sys.argv[0]) + '/pynus/pynus.py' + device_args
console = pexpect.spawn(pynus, encoding='UTF-8')
console.run_command = types.MethodType(run_command, console)
console.sync = types.MethodType(sync, console)
console.unsync = types.MethodType(unsync, console)
if args.verbose:
console.logfile = sys.stdout
else:
console.logfile = io.StringIO()
try:
console.expect(r'Connect.*\(([0-9A-F:]*)\)')
except pexpect.exceptions.TIMEOUT:
print('ERROR: Cannot find suitable wasp-os device')
if not args.verbose:
print_log(console.logfile)
sys.exit(1)
if not args.verbose:
console.logfile.close()
console.logfile = None
macaddr = console.match.group(1)
console.expect('Exit console using Ctrl-X')
time.sleep(0.5)
console.sync()
if args.rtc:
handle_rtc(console)
if args.check_rtc:
check_rtc(console)
if args.exec:
handle_exec(console, args.exec)
if args.eval:
handle_eval(console, args.eval)
if args.pull:
handle_binary_download(console, args.pull, args.pull)
if args.push:
handle_binary_upload(console, args.push, args.push)
if args.upload:
if args.binary:
handle_binary_upload(console, args.upload, args.upload_as)
else:
handle_upload(console, args.upload, args.upload_as)
if args.memfree:
handle_memory_free(console)
if args.console:
console.close()
argv = pynus.split()
os.execv(argv[0], argv)
if args.ota:
handle_reset(console, ota=True)
time.sleep(2.0)
dfu = os.path.dirname(sys.argv[0]) + '/ota-dfu/dfu.py'
os.execl(dfu, dfu, '-z', args.ota, '-a', macaddr, '--legacy')
if args.reset:
handle_reset(console)
sys.exit(0)
if args.bootloader:
handle_reset(console, ota=True)
sys.exit(0)
if args.battery:
handle_battery_level(console)
console.unsync()