swapforth

1
#!/usr/bin/env python
2
#
3
# ESP8266 ROM Bootloader Utility
4
# https://github.com/themadinventor/esptool
5
#
6
# Copyright (C) 2014 Fredrik Ahlberg
7
#
8
# This program is free software; you can redistribute it and/or modify it under
9
# the terms of the GNU General Public License as published by the Free Software
10
# Foundation; either version 2 of the License, or (at your option) any later version.
11
#
12
# This program is distributed in the hope that it will be useful, but WITHOUT
13
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
14
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
15
#
16
# You should have received a copy of the GNU General Public License along with
17
# this program; if not, write to the Free Software Foundation, Inc., 51 Franklin
18
# Street, Fifth Floor, Boston, MA 02110-1301 USA.
19

20
import sys
21
import struct
22
import serial
23
import time
24
import argparse
25
import os
26
import subprocess
27
import tempfile
28

29

30
sys.path.append("../shell")
31
import swapforth
32

33
class TetheredESP(swapforth.TetheredTarget):
34

35
    def open_ser(self, port, speed):
36
        self.ser = port
37

38
    def reset(self):
39
        while 1:
40
            c = self.ser.read(1)
41
            if c == b'\x1e':
42
                break
43
            sys.stdout.write(c)
44
            sys.stdout.flush()
45
        print 'ESCAPED'
46
        return
47
        # time.sleep(1)
48
        self.ser.write("words\n1 2 + .x\nwords\n")
49
        self.ser.flush();
50
        while 1:
51
            c = self.ser.read(1)
52
            sys.stdout.write(c)
53
            sys.stdout.flush()
54

55
    def boot(self, bootfile = None):
56
        sys.stdout.write('Contacting... ')
57
        self.reset()
58
        print('established')
59

60
    def interrupt(self):
61
        self.reset()
62

63
    def serialize(self):
64
        l = self.command_response('0 here dump')
65
        lines = l.strip().replace('\r', '').split('\n')
66
        s = []
67
        for l in lines:
68
            l = l.split()
69
            s += [int(b, 16) for b in l[1:17]]
70
        s = array.array('B', s).tostring().ljust(32768, chr(0xff))
71
        return array.array('i', s)
72

73

74
class ESPROM:
75
    # These are the currently known commands supported by the ROM
76
    ESP_FLASH_BEGIN = 0x02
77
    ESP_FLASH_DATA  = 0x03
78
    ESP_FLASH_END   = 0x04
79
    ESP_MEM_BEGIN   = 0x05
80
    ESP_MEM_END     = 0x06
81
    ESP_MEM_DATA    = 0x07
82
    ESP_SYNC        = 0x08
83
    ESP_WRITE_REG   = 0x09
84
    ESP_READ_REG    = 0x0a
85

86
    # Maximum block sized for RAM and Flash writes, respectively.
87
    ESP_RAM_BLOCK   = 0x1800
88
    ESP_FLASH_BLOCK = 0x400
89

90
    # Default baudrate. The ROM auto-bauds, so we can use more or less whatever we want.
91
    ESP_ROM_BAUD    = 115200
92

93
    # First byte of the application image
94
    ESP_IMAGE_MAGIC = 0xe9
95

96
    # Initial state for the checksum routine
97
    ESP_CHECKSUM_MAGIC = 0xef
98

99
    # OTP ROM addresses
100
    ESP_OTP_MAC0    = 0x3ff00050
101
    ESP_OTP_MAC1    = 0x3ff00054
102

103
    # Sflash stub: an assembly routine to read from spi flash and send to host
104
    SFLASH_STUB     = "\x80\x3c\x00\x40\x1c\x4b\x00\x40\x21\x11\x00\x40\x00\x80" \
105
                      "\xfe\x3f\xc1\xfb\xff\xd1\xf8\xff\x2d\x0d\x31\xfd\xff\x41\xf7\xff\x4a" \
106
                      "\xdd\x51\xf9\xff\xc0\x05\x00\x21\xf9\xff\x31\xf3\xff\x41\xf5\xff\xc0" \
107
                      "\x04\x00\x0b\xcc\x56\xec\xfd\x06\xff\xff\x00\x00"
108

109
    def __init__(self, port=0, baud=ESP_ROM_BAUD):
110
        self._port = serial.Serial(port)
111
        # setting baud rate in a separate step is a workaround for
112
        # CH341 driver on some Linux versions (this opens at 9600 then
113
        # sets), shouldn't matter for other platforms/drivers. See
114
        # https://github.com/themadinventor/esptool/issues/44#issuecomment-107094446
115
        self._port.baudrate = baud
116

117
    """ Read bytes from the serial port while performing SLIP unescaping """
118
    def read(self, length=1):
119
        b = ''
120
        while len(b) < length:
121
            c = self._port.read(1)
122
            if c == '\xdb':
123
                c = self._port.read(1)
124
                if c == '\xdc':
125
                    b = b + '\xc0'
126
                elif c == '\xdd':
127
                    b = b + '\xdb'
128
                else:
129
                    raise FatalError('Invalid SLIP escape')
130
            else:
131
                b = b + c
132
        return b
133

134
    """ Write bytes to the serial port while performing SLIP escaping """
135
    def write(self, packet):
136
        buf = '\xc0' \
137
              + (packet.replace('\xdb','\xdb\xdd').replace('\xc0','\xdb\xdc')) \
138
              + '\xc0'
139
        self._port.write(buf)
140

141
    """ Calculate checksum of a blob, as it is defined by the ROM """
142
    @staticmethod
143
    def checksum(data, state=ESP_CHECKSUM_MAGIC):
144
        for b in data:
145
            state ^= ord(b)
146
        return state
147

148
    """ Send a request and read the response """
149
    def command(self, op=None, data=None, chk=0):
150
        if op:
151
            pkt = struct.pack('<BBHI', 0x00, op, len(data), chk) + data
152
            self.write(pkt)
153

154
        # tries to get a response until that response has the
155
        # same operation as the request or a retries limit has
156
        # exceeded. This is needed for some esp8266s that
157
        # reply with more sync responses than expected.
158
        retries = 100
159
        while retries > 0:
160
            (op_ret, val, body) = self.receive_response()
161
            if op is None or op_ret == op:
162
                return val, body  # valid response received
163
            retries = retries - 1
164

165
        raise FatalError("Response doesn't match request")
166

167
    """ Receive a response to a command """
168
    def receive_response(self):
169
        # Read header of response and parse
170
        if self._port.read(1) != '\xc0':
171
            raise FatalError('Invalid head of packet')
172
        hdr = self.read(8)
173
        (resp, op_ret, len_ret, val) = struct.unpack('<BBHI', hdr)
174
        if resp != 0x01:
175
            raise FatalError('Invalid response 0x%02x" to command' % resp)
176

177
        # The variable-length body
178
        body = self.read(len_ret)
179

180
        # Terminating byte
181
        if self._port.read(1) != chr(0xc0):
182
            raise FatalError('Invalid end of packet')
183

184
        return op_ret, val, body
185

186
    """ Perform a connection test """
187
    def sync(self):
188
        self.command(ESPROM.ESP_SYNC, '\x07\x07\x12\x20' + 32 * '\x55')
189
        for i in xrange(7):
190
            self.command()
191

192
    """ Try connecting repeatedly until successful, or giving up """
193
    def connect(self):
194
        print 'Connecting...'
195

196
        for _ in xrange(4):
197
            # issue reset-to-bootloader:
198
            # RTS = either CH_PD or nRESET (both active low = chip in reset)
199
            # DTR = GPIO0 (active low = boot to flasher)
200
            self._port.setDTR(False)
201
            self._port.setRTS(True)
202
            time.sleep(0.05)
203
            self._port.setDTR(True)
204
            self._port.setRTS(False)
205
            time.sleep(0.05)
206
            self._port.setDTR(False)
207

208
            # worst-case latency timer should be 255ms (probably <20ms)
209
            self._port.timeout = 0.3
210
            for _ in xrange(4):
211
                try:
212
                    self._port.flushInput()
213
                    self._port.flushOutput()
214
                    self.sync()
215
                    self._port.timeout = 5
216
                    return
217
                except:
218
                    time.sleep(0.05)
219
        raise FatalError('Failed to connect to ESP8266')
220

221
    """ Read memory address in target """
222
    def read_reg(self, addr):
223
        res = self.command(ESPROM.ESP_READ_REG, struct.pack('<I', addr))
224
        if res[1] != "\0\0":
225
            raise FatalError('Failed to read target memory')
226
        return res[0]
227

228
    """ Write to memory address in target """
229
    def write_reg(self, addr, value, mask, delay_us=0):
230
        if self.command(ESPROM.ESP_WRITE_REG,
231
                        struct.pack('<IIII', addr, value, mask, delay_us))[1] != "\0\0":
232
            raise FatalError('Failed to write target memory')
233

234
    """ Start downloading an application image to RAM """
235
    def mem_begin(self, size, blocks, blocksize, offset):
236
        if self.command(ESPROM.ESP_MEM_BEGIN,
237
                        struct.pack('<IIII', size, blocks, blocksize, offset))[1] != "\0\0":
238
            raise FatalError('Failed to enter RAM download mode')
239

240
    """ Send a block of an image to RAM """
241
    def mem_block(self, data, seq):
242
        if self.command(ESPROM.ESP_MEM_DATA,
243
                        struct.pack('<IIII', len(data), seq, 0, 0) + data,
244
                        ESPROM.checksum(data))[1] != "\0\0":
245
            raise FatalError('Failed to write to target RAM')
246

247
    """ Leave download mode and run the application """
248
    def mem_finish(self, entrypoint=0):
249
        if self.command(ESPROM.ESP_MEM_END,
250
                        struct.pack('<II', int(entrypoint == 0), entrypoint))[1] != "\0\0":
251
            raise FatalError('Failed to leave RAM download mode')
252

253
    """ Start downloading to Flash (performs an erase) """
254
    def flash_begin(self, size, offset):
255
        old_tmo = self._port.timeout
256
        num_blocks = (size + ESPROM.ESP_FLASH_BLOCK - 1) / ESPROM.ESP_FLASH_BLOCK
257

258
        sectors_per_block = 16
259
        sector_size = 4096
260
        num_sectors = (size + sector_size - 1) / sector_size
261
        start_sector = offset / sector_size
262

263
        head_sectors = sectors_per_block - (start_sector % sectors_per_block)
264
        if num_sectors < head_sectors:
265
            head_sectors = num_sectors
266

267
        if num_sectors < 2 * head_sectors:
268
            erase_size = (num_sectors + 1) / 2 * sector_size
269
        else:
270
            erase_size = (num_sectors - head_sectors) * sector_size
271

272
        self._port.timeout = 20
273
        t = time.time()
274
        result = self.command(ESPROM.ESP_FLASH_BEGIN,
275
                              struct.pack('<IIII', erase_size, num_blocks, ESPROM.ESP_FLASH_BLOCK, offset))[1]
276
        if size != 0:
277
            print "Took %.2fs to erase flash block" % (time.time() - t)
278
        if result != "\0\0":
279
            raise FatalError.WithResult('Failed to enter Flash download mode (result "%s")', result)
280
        self._port.timeout = old_tmo
281

282
    """ Write block to flash """
283
    def flash_block(self, data, seq):
284
        result = self.command(ESPROM.ESP_FLASH_DATA, struct.pack('<IIII', len(data), seq, 0, 0) + data, ESPROM.checksum(data))[1]
285
        if result != "\0\0":
286
            raise FatalError.WithResult('Failed to write to target Flash after seq %d (got result %%s)' % seq, result)
287

288
    """ Leave flash mode and run/reboot """
289
    def flash_finish(self, reboot=False):
290
        pkt = struct.pack('<I', int(not reboot))
291
        if self.command(ESPROM.ESP_FLASH_END, pkt)[1] != "\0\0":
292
            raise FatalError('Failed to leave Flash mode')
293

294
    """ Run application code in flash """
295
    def run(self, reboot=False):
296
        # Fake flash begin immediately followed by flash end
297
        self.flash_begin(0, 0)
298
        self.flash_finish(reboot)
299

300
    """ Read MAC from OTP ROM """
301
    def read_mac(self):
302
        mac0 = self.read_reg(self.ESP_OTP_MAC0)
303
        mac1 = self.read_reg(self.ESP_OTP_MAC1)
304
        if ((mac1 >> 16) & 0xff) == 0:
305
            oui = (0x18, 0xfe, 0x34)
306
        elif ((mac1 >> 16) & 0xff) == 1:
307
            oui = (0xac, 0xd0, 0x74)
308
        else:
309
            raise FatalError("Unknown OUI")
310
        return oui + ((mac1 >> 8) & 0xff, mac1 & 0xff, (mac0 >> 24) & 0xff)
311

312
    """ Read SPI flash manufacturer and device id """
313
    def flash_id(self):
314
        self.flash_begin(0, 0)
315
        self.write_reg(0x60000240, 0x0, 0xffffffff)
316
        self.write_reg(0x60000200, 0x10000000, 0xffffffff)
317
        flash_id = self.read_reg(0x60000240)
318
        self.flash_finish(False)
319
        return flash_id
320

321
    """ Read SPI flash """
322
    def flash_read(self, offset, size, count=1):
323
        # Create a custom stub
324
        stub = struct.pack('<III', offset, size, count) + self.SFLASH_STUB
325

326
        # Trick ROM to initialize SFlash
327
        self.flash_begin(0, 0)
328

329
        # Download stub
330
        self.mem_begin(len(stub), 1, len(stub), 0x40100000)
331
        self.mem_block(stub, 0)
332
        self.mem_finish(0x4010001c)
333

334
        # Fetch the data
335
        t = time.time()
336
        data = ''
337
        for _ in xrange(count):
338
            if self._port.read(1) != '\xc0':
339
                raise FatalError('Invalid head of packet (sflash read)')
340

341
            data += self.read(size)
342
            if len(data) % 4096 == 0:
343
                sys.stdout.write(".")
344
                sys.stdout.flush()
345

346
            if self._port.read(1) != chr(0xc0):
347
                raise FatalError('Invalid end of packet (sflash read)')
348

349
        t = time.time() - t
350
        print '\nRead %d bytes in %.2f seconds (%.2f kbit/s)...' % (len(data), t, len(data) / t * 8 / 1000)
351
        return data
352

353
    """ Abuse the loader protocol to force flash to be left in write mode """
354
    def flash_unlock_dio(self):
355
        # Enable flash write mode
356
        self.flash_begin(0, 0)
357
        # Reset the chip rather than call flash_finish(), which would have
358
        # write protected the chip again (why oh why does it do that?!)
359
        self.mem_begin(0,0,0,0x40100000)
360
        self.mem_finish(0x40000080)
361

362
    """ Perform a chip erase of SPI flash """
363
    def flash_erase(self):
364
        # Trick ROM to initialize SFlash
365
        self.flash_begin(0, 0)
366

367
        # This is hacky: we don't have a custom stub, instead we trick
368
        # the bootloader to jump to the SPIEraseChip() routine and then halt/crash
369
        # when it tries to boot an unconfigured system.
370
        self.mem_begin(0,0,0,0x40100000)
371
        self.mem_finish(0x40004984)
372

373
        # Yup - there's no good way to detect if we succeeded.
374
        # It it on the other hand unlikely to fail.
375

376

377
class ESPFirmwareImage:
378

379
    def __init__(self, filename=None):
380
        self.segments = []
381
        self.entrypoint = 0
382
        self.flash_mode = 0
383
        self.flash_size_freq = 0
384

385
        if filename is not None:
386
            f = file(filename, 'rb')
387
            (magic, segments, self.flash_mode, self.flash_size_freq, self.entrypoint) = struct.unpack('<BBBBI', f.read(8))
388

389
            # some sanity check
390
            if magic != ESPROM.ESP_IMAGE_MAGIC or segments > 16:
391
                raise FatalError('Invalid firmware image')
392

393
            for i in xrange(segments):
394
                (offset, size) = struct.unpack('<II', f.read(8))
395
                if offset > 0x40200000 or offset < 0x3ffe0000 or size > 65536:
396
                    raise FatalError('Suspicious segment 0x%x, length %d' % (offset, size))
397
                segment_data = f.read(size)
398
                if len(segment_data) < size:
399
                    raise FatalError('End of file reading segment 0x%x, length %d (actual length %d)' % (offset, size, len(segment_data)))
400
                self.segments.append((offset, size, segment_data))
401

402
            # Skip the padding. The checksum is stored in the last byte so that the
403
            # file is a multiple of 16 bytes.
404
            align = 15 - (f.tell() % 16)
405
            f.seek(align, 1)
406

407
            self.checksum = ord(f.read(1))
408

409
    def add_segment(self, addr, data):
410
        # Data should be aligned on word boundary
411
        l = len(data)
412
        if l % 4:
413
            data += b"\x00" * (4 - l % 4)
414
        if l > 0:
415
            self.segments.append((addr, len(data), data))
416

417
    def save(self, filename):
418
        f = file(filename, 'wb')
419
        f.write(struct.pack('<BBBBI', ESPROM.ESP_IMAGE_MAGIC, len(self.segments),
420
                            self.flash_mode, self.flash_size_freq, self.entrypoint))
421

422
        checksum = ESPROM.ESP_CHECKSUM_MAGIC
423
        for (offset, size, data) in self.segments:
424
            f.write(struct.pack('<II', offset, size))
425
            f.write(data)
426
            checksum = ESPROM.checksum(data, checksum)
427

428
        align = 15 - (f.tell() % 16)
429
        f.seek(align, 1)
430
        f.write(struct.pack('B', checksum))
431

432

433
class ELFFile:
434

435
    def __init__(self, name):
436
        self.name = name
437
        self.symbols = None
438

439
    def _fetch_symbols(self):
440
        if self.symbols is not None:
441
            return
442
        self.symbols = {}
443
        try:
444
            tool_nm = "xtensa-lx106-elf-nm"
445
            if os.getenv('XTENSA_CORE') == 'lx106':
446
                tool_nm = "xt-nm"
447
            proc = subprocess.Popen([tool_nm, self.name], stdout=subprocess.PIPE)
448
        except OSError:
449
            print "Error calling %s, do you have Xtensa toolchain in PATH?" % tool_nm
450
            sys.exit(1)
451
        for l in proc.stdout:
452
            fields = l.strip().split()
453
            try:
454
                if fields[0] == "U":
455
                    print "Warning: ELF binary has undefined symbol %s" % fields[1]
456
                    continue
457
                self.symbols[fields[2]] = int(fields[0], 16)
458
            except ValueError:
459
                raise FatalError("Failed to strip symbol output from nm: %s" % fields)
460

461
    def get_symbol_addr(self, sym):
462
        self._fetch_symbols()
463
        return self.symbols[sym]
464

465
    def get_entry_point(self):
466
        tool_readelf = "xtensa-lx106-elf-readelf"
467
        if os.getenv('XTENSA_CORE') == 'lx106':
468
            tool_readelf = "xt-readelf"
469
        try:
470
            proc = subprocess.Popen([tool_readelf, "-h", self.name], stdout=subprocess.PIPE)
471
        except OSError:
472
            print "Error calling %s, do you have Xtensa toolchain in PATH?" % tool_readelf
473
            sys.exit(1)
474
        for l in proc.stdout:
475
            fields = l.strip().split()
476
            if fields[0] == "Entry":
477
                return int(fields[3], 0)
478

479
    def load_section(self, section):
480
        tool_objcopy = "xtensa-lx106-elf-objcopy"
481
        if os.getenv('XTENSA_CORE') == 'lx106':
482
            tool_objcopy = "xt-objcopy"
483
        tmpsection = tempfile.mktemp(suffix=".section")
484
        try:
485
            subprocess.check_call([tool_objcopy, "--only-section", section, "-Obinary", self.name, tmpsection])
486
            with open(tmpsection, "rb") as f:
487
                data = f.read()
488
        finally:
489
            os.remove(tmpsection)
490
        return data
491

492

493
def arg_auto_int(x):
494
    return int(x, 0)
495

496

497
def div_roundup(a, b):
498
    """ Return a/b rounded up to nearest integer,
499
    equivalent result to int(math.ceil(float(int(a)) / float(int(b))), only
500
    without possible floating point accuracy errors.
501
    """
502
    return (int(a) + int(b) - 1) / int(b)
503

504

505
class FatalError(RuntimeError):
506
    """
507
    Wrapper class for runtime errors that aren't caused by internal bugs, but by
508
    ESP8266 responses or input content.
509
    """
510
    def __init__(self, message):
511
        RuntimeError.__init__(self, message)
512

513
    @staticmethod
514
    def WithResult(message, result):
515
        """
516
        Return a fatal error object that includes the hex values of
517
        'result' as a string formatted argument.
518
        """
519
        return FatalError(message % ", ".join(hex(ord(x)) for x in result))
520

521

522
def main():
523
    parser = argparse.ArgumentParser(description='ESP8266 ROM Bootloader Utility', prog='esptool')
524

525
    parser.add_argument(
526
        '--port', '-p',
527
        help='Serial port device',
528
        default='/dev/ttyUSB0')
529

530
    parser.add_argument(
531
        '--baud', '-b',
532
        help='Serial port baud rate',
533
        type=arg_auto_int,
534
        default=ESPROM.ESP_ROM_BAUD)
535

536
    subparsers = parser.add_subparsers(
537
        dest='operation',
538
        help='Run esptool {command} -h for additional help')
539

540
    parser_load_ram = subparsers.add_parser(
541
        'load_ram',
542
        help='Download an image to RAM and execute')
543
    parser_load_ram.add_argument('filename', help='Firmware image')
544

545
    parser_dump_mem = subparsers.add_parser(
546
        'dump_mem',
547
        help='Dump arbitrary memory to disk')
548
    parser_dump_mem.add_argument('address', help='Base address', type=arg_auto_int)
549
    parser_dump_mem.add_argument('size', help='Size of region to dump', type=arg_auto_int)
550
    parser_dump_mem.add_argument('filename', help='Name of binary dump')
551

552
    parser_read_mem = subparsers.add_parser(
553
        'read_mem',
554
        help='Read arbitrary memory location')
555
    parser_read_mem.add_argument('address', help='Address to read', type=arg_auto_int)
556

557
    parser_write_mem = subparsers.add_parser(
558
        'write_mem',
559
        help='Read-modify-write to arbitrary memory location')
560
    parser_write_mem.add_argument('address', help='Address to write', type=arg_auto_int)
561
    parser_write_mem.add_argument('value', help='Value', type=arg_auto_int)
562
    parser_write_mem.add_argument('mask', help='Mask of bits to write', type=arg_auto_int)
563

564
    parser_write_flash = subparsers.add_parser(
565
        'write_flash',
566
        help='Write a binary blob to flash')
567
    parser_write_flash.add_argument('addr_filename', nargs='+', help='Address and binary file to write there, separated by space')
568
    parser_write_flash.add_argument('--flash_freq', '-ff', help='SPI Flash frequency',
569
                                    choices=['40m', '26m', '20m', '80m'], default='40m')
570
    parser_write_flash.add_argument('--flash_mode', '-fm', help='SPI Flash mode',
571
                                    choices=['qio', 'qout', 'dio', 'dout'], default='qio')
572
    parser_write_flash.add_argument('--flash_size', '-fs', help='SPI Flash size in Mbit',
573
                                    choices=['4m', '2m', '8m', '16m', '32m', '16m-c1', '32m-c1', '32m-c2'], default='4m')
574

575
    subparsers.add_parser(
576
        'run',
577
        help='Run application code in flash')
578

579
    parser_image_info = subparsers.add_parser(
580
        'image_info',
581
        help='Dump headers from an application image')
582
    parser_image_info.add_argument('filename', help='Image file to parse')
583

584
    parser_make_image = subparsers.add_parser(
585
        'make_image',
586
        help='Create an application image from binary files')
587
    parser_make_image.add_argument('output', help='Output image file')
588
    parser_make_image.add_argument('--segfile', '-f', action='append', help='Segment input file')
589
    parser_make_image.add_argument('--segaddr', '-a', action='append', help='Segment base address', type=arg_auto_int)
590
    parser_make_image.add_argument('--entrypoint', '-e', help='Address of entry point', type=arg_auto_int, default=0)
591

592
    parser_elf2image = subparsers.add_parser(
593
        'elf2image',
594
        help='Create an application image from ELF file')
595
    parser_elf2image.add_argument('input', help='Input ELF file')
596
    parser_elf2image.add_argument('--output', '-o', help='Output filename prefix', type=str)
597
    parser_elf2image.add_argument('--flash_freq', '-ff', help='SPI Flash frequency',
598
                                  choices=['40m', '26m', '20m', '80m'], default='40m')
599
    parser_elf2image.add_argument('--flash_mode', '-fm', help='SPI Flash mode',
600
                                  choices=['qio', 'qout', 'dio', 'dout'], default='qio')
601
    parser_elf2image.add_argument('--flash_size', '-fs', help='SPI Flash size in Mbit',
602
                                  choices=['4m', '2m', '8m', '16m', '32m', '16m-c1', '32m-c1', '32m-c2'], default='4m')
603

604
    subparsers.add_parser(
605
        'read_mac',
606
        help='Read MAC address from OTP ROM')
607

608
    subparsers.add_parser(
609
        'flash_id',
610
        help='Read SPI flash manufacturer and device ID')
611

612
    parser_read_flash = subparsers.add_parser(
613
        'read_flash',
614
        help='Read SPI flash content')
615
    parser_read_flash.add_argument('address', help='Start address', type=arg_auto_int)
616
    parser_read_flash.add_argument('size', help='Size of region to dump', type=arg_auto_int)
617
    parser_read_flash.add_argument('filename', help='Name of binary dump')
618

619
    subparsers.add_parser(
620
        'erase_flash',
621
        help='Perform Chip Erase on SPI flash')
622

623
    args = parser.parse_args()
624

625
    # Create the ESPROM connection object, if needed
626
    esp = None
627
    if args.operation not in ('image_info','make_image','elf2image'):
628
        esp = ESPROM(args.port, args.baud)
629
        esp.connect()
630

631
    # Do the actual work. Should probably be split into separate functions.
632
    if args.operation == 'load_ram':
633
        image = ESPFirmwareImage(args.filename)
634

635
        print 'RAM boot...'
636
        for (offset, size, data) in image.segments:
637
            print 'Downloading %d bytes at %08x...' % (size, offset),
638
            sys.stdout.flush()
639
            esp.mem_begin(size, div_roundup(size, esp.ESP_RAM_BLOCK), esp.ESP_RAM_BLOCK, offset)
640

641
            seq = 0
642
            while len(data) > 0:
643
                esp.mem_block(data[0:esp.ESP_RAM_BLOCK], seq)
644
                data = data[esp.ESP_RAM_BLOCK:]
645
                seq += 1
646
            print 'done!'
647

648
        print 'All segments done, executing at %08x' % image.entrypoint
649
        esp.mem_finish(image.entrypoint)
650

651
    elif args.operation == 'read_mem':
652
        print '0x%08x = 0x%08x' % (args.address, esp.read_reg(args.address))
653

654
    elif args.operation == 'write_mem':
655
        esp.write_reg(args.address, args.value, args.mask, 0)
656
        print 'Wrote %08x, mask %08x to %08x' % (args.value, args.mask, args.address)
657

658
    elif args.operation == 'dump_mem':
659
        f = file(args.filename, 'wb')
660
        for i in xrange(args.size / 4):
661
            d = esp.read_reg(args.address + (i * 4))
662
            f.write(struct.pack('<I', d))
663
            if f.tell() % 1024 == 0:
664
                print '\r%d bytes read... (%d %%)' % (f.tell(),
665
                                                      f.tell() * 100 / args.size),
666
                sys.stdout.flush()
667
        print 'Done!'
668

669
    elif args.operation == 'write_flash':
670
        assert len(args.addr_filename) % 2 == 0
671

672
        flash_mode = {'qio':0, 'qout':1, 'dio':2, 'dout': 3}[args.flash_mode]
673
        flash_size_freq = {'4m':0x00, '2m':0x10, '8m':0x20, '16m':0x30, '32m':0x40, '16m-c1': 0x50, '32m-c1':0x60, '32m-c2':0x70}[args.flash_size]
674
        flash_size_freq += {'40m':0, '26m':1, '20m':2, '80m': 0xf}[args.flash_freq]
675
        flash_info = struct.pack('BB', flash_mode, flash_size_freq)
676

677
        while args.addr_filename:
678
            address = int(args.addr_filename[0], 0)
679
            filename = args.addr_filename[1]
680
            args.addr_filename = args.addr_filename[2:]
681
            image = file(filename, 'rb').read()
682
            print 'Erasing flash...'
683
            blocks = div_roundup(len(image), esp.ESP_FLASH_BLOCK)
684
            esp.flash_begin(blocks * esp.ESP_FLASH_BLOCK, address)
685
            seq = 0
686
            written = 0
687
            t = time.time()
688
            while len(image) > 0:
689
                print '\rWriting at 0x%08x... (%d %%)' % (address + seq * esp.ESP_FLASH_BLOCK, 100 * (seq + 1) / blocks),
690
                sys.stdout.flush()
691
                block = image[0:esp.ESP_FLASH_BLOCK]
692
                # Fix sflash config data
693
                if address == 0 and seq == 0 and block[0] == '\xe9':
694
                    block = block[0:2] + flash_info + block[4:]
695
                # Pad the last block
696
                block = block + '\xff' * (esp.ESP_FLASH_BLOCK - len(block))
697
                esp.flash_block(block, seq)
698
                image = image[esp.ESP_FLASH_BLOCK:]
699
                seq += 1
700
                written += len(block)
701
            t = time.time() - t
702
            print '\rWrote %d bytes at 0x%08x in %.1f seconds (%.1f kbit/s)...' % (written, address, t, written / t * 8 / 1000)
703
        print '\nLeaving...'
704
        """
705
        if args.flash_mode == 'dio':
706
            esp.flash_unlock_dio()
707
        else:
708
            esp.flash_begin(0, 0)
709
            esp.flash_finish(reboot)
710
        """
711
        esp.run(1)
712
        r = TetheredESP(esp._port)
713
        r.boot()
714

715
        for i in xrange(0):
716
            print i, r.command_response('words')
717
            time.sleep(.1)
718
        r.searchpath += ['../common', '../anstests']
719
        try:
720
            r.include('swapforth.fs')
721
        except swapforth.Bye:
722
            pass
723
        r.shell()
724

725
        # esp._port.write('64 parse\r')
726
        while 0:
727
            sys.stdout.write(esp._port.read(1))
728
            sys.stdout.flush()
729

730
    elif args.operation == 'run':
731
        esp.run()
732
        r = TetheredESP(esp._port)
733
        r.boot()
734
        r.shell()
735

736
    elif args.operation == 'image_info':
737
        image = ESPFirmwareImage(args.filename)
738
        print ('Entry point: %08x' % image.entrypoint) if image.entrypoint != 0 else 'Entry point not set'
739
        print '%d segments' % len(image.segments)
740
        print
741
        checksum = ESPROM.ESP_CHECKSUM_MAGIC
742
        for (idx, (offset, size, data)) in enumerate(image.segments):
743
            print 'Segment %d: %5d bytes at %08x' % (idx + 1, size, offset)
744
            checksum = ESPROM.checksum(data, checksum)
745
        print
746
        print 'Checksum: %02x (%s)' % (image.checksum, 'valid' if image.checksum == checksum else 'invalid!')
747

748
    elif args.operation == 'make_image':
749
        image = ESPFirmwareImage()
750
        if len(args.segfile) == 0:
751
            raise FatalError('No segments specified')
752
        if len(args.segfile) != len(args.segaddr):
753
            raise FatalError('Number of specified files does not match number of specified addresses')
754
        for (seg, addr) in zip(args.segfile, args.segaddr):
755
            data = file(seg, 'rb').read()
756
            image.add_segment(addr, data)
757
        image.entrypoint = args.entrypoint
758
        image.save(args.output)
759

760
    elif args.operation == 'elf2image':
761
        if args.output is None:
762
            args.output = args.input + '-'
763
        e = ELFFile(args.input)
764
        image = ESPFirmwareImage()
765
        image.entrypoint = e.get_entry_point()
766
        for section, start in ((".text", "_text_start"), (".data", "_data_start"), (".rodata", "_rodata_start")):
767
            data = e.load_section(section)
768
            image.add_segment(e.get_symbol_addr(start), data)
769

770
        image.flash_mode = {'qio':0, 'qout':1, 'dio':2, 'dout': 3}[args.flash_mode]
771
        image.flash_size_freq = {'4m':0x00, '2m':0x10, '8m':0x20, '16m':0x30, '32m':0x40, '16m-c1': 0x50, '32m-c1':0x60, '32m-c2':0x70}[args.flash_size]
772
        image.flash_size_freq += {'40m':0, '26m':1, '20m':2, '80m': 0xf}[args.flash_freq]
773

774
        image.save(args.output + "0x00000.bin")
775
        data = e.load_section(".irom0.text")
776
        off = e.get_symbol_addr("_irom0_text_start") - 0x40200000
777
        assert off >= 0
778
        f = open(args.output + "0x%05x.bin" % off, "wb")
779
        f.write(data)
780
        f.close()
781

782
    elif args.operation == 'read_mac':
783
        mac = esp.read_mac()
784
        print 'MAC: %s' % ':'.join(map(lambda x: '%02x' % x, mac))
785

786
    elif args.operation == 'flash_id':
787
        flash_id = esp.flash_id()
788
        print 'Manufacturer: %02x' % (flash_id & 0xff)
789
        print 'Device: %02x%02x' % ((flash_id >> 8) & 0xff, (flash_id >> 16) & 0xff)
790

791
    elif args.operation == 'read_flash':
792
        print 'Please wait...'
793
        file(args.filename, 'wb').write(esp.flash_read(args.address, 1024, div_roundup(args.size, 1024))[:args.size])
794

795
    elif args.operation == 'erase_flash':
796
        esp.flash_erase()
797

798
if __name__ == '__main__':
799
    try:
800
        main()
801
    except FatalError as e:
802
        print '\nA fatal error occurred: %s' % e
803
        sys.exit(2)
804