1
#!/usr/bin/env python3
2
"""Flash firmware and music, or either alone, then reboot. Requires a verified saved backup."""
3
import argparse
4
from contextlib import nullcontext
5
from serial_device import device_port, device_identity
6
from pathlib import Path
7
import subprocess
8
import sys
9
import tempfile
10
from project import ROOT, build_dir
11
from flash_backup import require_backup
12
from music_partition import MUSIC_BYTES, MUSIC_OFFSET, require_music_partition
13
from flash_cache import FlashCache, flash_images, layout_digest
14
 
15
def main():
16
    parser = argparse.ArgumentParser(description=__doc__)
17
    contents = parser.add_mutually_exclusive_group()
18
    contents.add_argument("--firmware-only", action="store_true", help="Keep the previously flashed music partition")
19
    contents.add_argument("--music-only", action="store_true", help="Update only the music and track metadata, then reboot")
20
    parser.add_argument("--full", action="store_true", help="Write the selected images in full")
21
    args = parser.parse_args()
22
    build = build_dir()
23
    try:
24
        require_backup()
25
    except ValueError as error:
26
        raise SystemExit(str(error)) from None
27
    if not args.firmware_only:
28
        music = ROOT / "artifacts/radio/music.bin"
29
        header = music.read_bytes()[:64]
30
        if len(header) != 64 or header[:8] != b"S3MUSIC\0" or music.stat().st_size > MUSIC_BYTES:
31
            raise SystemExit("Prepare a valid music pack with make music before flashing")
32
    images = []
33
    if not args.music_only:
34
        images.extend([(0, build / "bootloader/bootloader.bin"), (0x8000, build / "partition_table/partition-table.bin"), (0x10000, build / "s3_radio.bin")])
35
    if not args.firmware_only:
36
        images.append((MUSIC_OFFSET, music))
37
    if any(not path.is_file() for _, path in images):
38
        raise SystemExit("Build the selected firmware and prepare music before flashing")
39
    port = device_port()
40
    identity = device_identity(port)
41
    cache = FlashCache(identity, "") if identity else None
42
    with cache.locked() if cache else nullcontext():
43
        if args.music_only:
44
            # Read the installed table; a locally built table does not prove migration.
45
            with tempfile.TemporaryDirectory(prefix="radio-partition-") as directory:
46
                table = Path(directory) / "partition.bin"
47
                subprocess.run([sys.executable, "-m", "esptool", "--chip", "esp32s3", "--port", port,
48
                                "--before", "default-reset", "--after", "hard-reset", "read-flash", "0x8000", "0x1000", str(table)], check=True)
49
                installed_table = table.read_bytes()
50
                try: require_music_partition(installed_table)
51
                except ValueError as error: raise SystemExit(str(error)) from None
52
        table = installed_table if args.music_only else (build / "partition_table/partition-table.bin").read_bytes()
53
        require_music_partition(table)
54
        command = [sys.executable, "-m", "esptool", "--chip", "esp32s3", "--port", port,
55
                   "--baud", "460800", "--before", "default-reset", "--after", "hard-reset"]
56
        if cache:
57
            cache.layout = layout_digest(table)
58
        result = flash_images(images, command, subprocess.run, cache, args.full)
59
    raise SystemExit(result)
60
 
61
 
62
if __name__ == "__main__":
63
    main()