more
[gnuk/gnuk.git] / tool / stlinkv2.py
1 #! /usr/bin/python
2
3 """
4 stlinkv2.py - a tool to control ST-Link/V2
5
6 Copyright (C) 2012, 2013 Free Software Initiative of Japan
7 Author: NIIBE Yutaka <gniibe@fsij.org>
8
9 This file is a part of Gnuk, a GnuPG USB Token implementation.
10
11 Gnuk is free software: you can redistribute it and/or modify it
12 under the terms of the GNU General Public License as published by
13 the Free Software Foundation, either version 3 of the License, or
14 (at your option) any later version.
15
16 Gnuk is distributed in the hope that it will be useful, but WITHOUT
17 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
18 or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
19 License for more details.
20
21 You should have received a copy of the GNU General Public License
22 along with this program.  If not, see <http://www.gnu.org/licenses/>.
23 """
24
25 from struct import *
26 import sys, time
27 import usb
28 from colorama import init as colorama_init, Fore, Back, Style
29
30 # INPUT: binary file
31
32 # Assumes only single ST-Link/V2 device is attached to computer.
33
34
35 GPIOA=0x40010800
36 GPIOB=0x40010C00
37 OPTION_BYTES_ADDR=0x1ffff800
38 RDP_KEY=0x00a5                  # Unlock readprotection
39 FLASH_BASE_ADDR=0x40022000
40
41 FLASH_KEYR=    FLASH_BASE_ADDR+0x04
42 FLASH_OPTKEYR= FLASH_BASE_ADDR+0x08
43 FLASH_SR=      FLASH_BASE_ADDR+0x0c
44 FLASH_CR=      FLASH_BASE_ADDR+0x10
45 FLASH_AR=      FLASH_BASE_ADDR+0x14
46 FLASH_OBR=     FLASH_BASE_ADDR+0x1c
47
48 FLASH_KEY1=0x45670123
49 FLASH_KEY2=0xcdef89ab
50
51 FLASH_SR_BSY=      0x0001
52 FLASH_SR_PGERR=    0x0004
53 FLASH_SR_WRPRTERR= 0x0010
54 FLASH_SR_EOP=      0x0020
55
56 FLASH_CR_PG=     0x0001
57 FLASH_CR_PER=    0x0002
58 FLASH_CR_MER=    0x0004
59 FLASH_CR_OPTPG=  0x0010
60 FLASH_CR_OPTER=  0x0020
61 FLASH_CR_STRT=   0x0040
62 FLASH_CR_LOCK=   0x0080
63 FLASH_CR_OPTWRE= 0x0200
64
65
66 SPI1= 0x40013000
67
68 def uint32(v):
69     return v[0] + (v[1]<<8) + (v[2]<<16) + (v[3]<<24)
70
71 ## HERE comes: "movw r2,#SIZE" instruction
72 prog_flash_write_body = "\x0A\x48" + "\x0B\x49" + \
73     "\x08\x4C" + "\x01\x25" + "\x14\x26" + "\x00\x27" + "\x25\x61" + \
74     "\xC3\x5B" + "\xCB\x53" + "\xE3\x68" + "\x2B\x42" + "\xFC\xD1" + \
75     "\x33\x42" + "\x02\xD1" + "\x02\x37" + "\x97\x42" + "\xF5\xD1" + \
76     "\x00\x27" + "\x27\x61" + "\x00\xBE" +  "\x00\x20\x02\x40" + \
77     "\x38\x00\x00\x20"
78 #   .SRC_ADDR: 0x20000038
79 ## HERE comes: target_addr in 4-byte
80 #   .TARGET_ADDR
81
82 def gen_prog_flash_write(addr,size):
83     return pack("<BBBB", (0x40 | (size&0xf000)>>12), (0xf2 | (size&0x0800)>>9),
84                 (size & 0x00ff), (0x02 | ((size&0x0700) >> 4))) + \
85                 prog_flash_write_body + pack("<I", addr)
86
87
88 ## HERE comes: "movw r0,#VAL" instruction
89 prog_option_bytes_write_body = "\x06\x49" + "\x05\x4A" + "\x10\x23" + \
90     "\x01\x24" + "\x13\x61" + "\x08\x80" + "\xD0\x68" + "\x20\x42" + \
91     "\xFC\xD1" + "\x00\x20" + "\x10\x61" + "\x00\xBE" + "\x00\x20\x02\x40"
92 ## HERE comes: target_addr in 4-byte
93 #   .TARGET_ADDR
94
95 def gen_prog_option_bytes_write(addr,val):
96     return pack("<BBBB", (0x40 | (val&0xf000)>>12), (0xf2 | (val&0x0800)>>9),
97                 (val & 0x00ff), (0x00 | ((val&0x0700) >> 4))) + \
98                 prog_option_bytes_write_body + pack("<I", addr)
99
100 prog_blank_check_body = "\x04\x49" + "\x05\x4A" + "\x08\x68" + "\x01\x30" + \
101     "\x02\xD1" + "\x04\x31" + "\x91\x42" + "\xF9\xD1" + "\x00\xBE" + \
102     "\x00\xBF" + "\x00\x00\x00\x08"
103 ## HERE comes: end_addr in 4-byte
104 # .END_ADDR
105
106 def gen_prog_blank_check(size):
107     return prog_blank_check_body + pack("<I", 0x08000000 + size)
108
109
110 SRAM_ADDRESS=0x20000000
111 BLOCK_SIZE=16384                # Should be less than (20KiB - 0x0038)
112 BLOCK_WRITE_TIMEOUT=80          # Increase this when you increase BLOCK_SIZE
113
114
115 class TimeOutError(Exception):
116      def __init__(self, msg):
117          self.msg = msg
118      def __str__(self):
119          return repr(self.msg)
120      def __repr__(self):
121          return "TimeoutError(" + self.msg + ")"
122
123 class OperationFailure(Exception):
124      def __init__(self, msg):
125          self.msg = msg
126      def __str__(self):
127          return repr(self.msg)
128      def __repr__(self):
129          return "OperationFailure(" + self.msg + ")"
130
131
132 class stlinkv2(object):
133     def __init__(self, dev):
134         self.__bulkout = 2
135         self.__bulkin  = 0x81
136
137         self.__timeout = 1000   # 1 second
138         conf = dev.configurations[0]
139         intf_alt = conf.interfaces[0]
140         intf = intf_alt[0]
141         if intf.interfaceClass != 0xff: # Vendor specific
142             raise ValueError("Wrong interface class.", intf.interfaceClass)
143         self.__devhandle = dev.open()
144         try:
145             self.__devhandle.setConfiguration(conf)
146         except:
147             pass
148         self.__devhandle.claimInterface(intf)
149         # self.__devhandle.setAltInterface(intf)  # This is not good for libusb-win32
150
151     def shutdown(self):
152         self.__devhandle.releaseInterface()
153
154     def execute_get(self, cmd, res_len):
155         self.__devhandle.bulkWrite(self.__bulkout, cmd, self.__timeout)
156         res = self.__devhandle.bulkRead(self.__bulkin, res_len, self.__timeout)
157         return res
158
159     def execute_put(self, cmd, data=None):
160         self.__devhandle.bulkWrite(self.__bulkout, cmd, self.__timeout)
161         if (data):
162             self.__devhandle.bulkWrite(self.__bulkout, data, self.__timeout)
163
164     def stl_mode(self):
165         v = self.execute_get("\xf5\x00", 2)
166         return (v[1] * 256 + v[0])
167
168     def exit_from_debug_swd(self):
169         self.execute_put("\xf2\x21\x00")
170         time.sleep(1)
171
172     def exit_from_dfu(self):
173         self.execute_put("\xf3\x07\x00")
174         time.sleep(1)
175
176     def exit_from_debug_swim(self):
177         self.execute_put("\xf4\x01\x00")
178         time.sleep(1)
179
180     def enter_swd(self):
181         self.execute_put("\xf2\x20\xa3")
182         time.sleep(1)
183
184     def get_status(self):
185         v = self.execute_get("\xf2\x01\x00", 2)
186         return (v[1] << 8) + v[0]
187     # RUN:128, HALT:129
188
189     def enter_debug(self):
190         v = self.execute_get("\xf2\x02\x00", 2)
191         return (v[1] << 8) + v[0]
192
193     def exit_debug(self):
194         self.execute_put("\xf2\x21\x00")
195
196     def reset_sys(self):
197         v = self.execute_get("\xf2\x03\x00", 2)
198         return (v[1] << 8) + v[0]
199
200     def read_memory(self, addr, length):
201         return self.execute_get("\xf2\x07" + pack('<IH', addr, length), length)
202
203     def read_memory_u32(self, addr):
204         return uint32(self.execute_get("\xf2\x07" + pack('<IH', addr, 4), 4))
205
206     def write_memory(self, addr, data):
207         return self.execute_put("\xf2\x08" + pack('<IH', addr, len(data)), data)
208
209     def write_memory_u32(self, addr, data):
210         return self.execute_put("\xf2\x08" + pack('<IH', addr, 4),
211                                 pack('<I', data))
212
213     def read_reg(self, regno):
214         return uint32(self.execute_get("\xf2\x05" + pack('<B', regno), 4))
215
216     def write_reg(self, regno, value):
217         return self.execute_get("\xf2\x06" + pack('<BI', regno, value), 2)
218
219     def write_debug_reg(self, addr, value):
220         return self.execute_get("\xf2\x35" + pack('<II', addr, value), 2)
221
222     def run(self):
223         v = self.execute_get("\xf2\x09\x00", 2)
224         return (v[1] << 8) + v[0]
225
226     def core_id(self):
227         v = self.execute_get("\xf2\x22\x00", 4)
228         return v[0] + (v[1]<<8) + (v[2]<<16) + (v[3]<<24)
229
230     def version(self):
231         v = self.execute_get("\xf1", 6)
232         val = (v[0] << 8) + v[1]
233         return ((val >> 12) & 0x0f, (val >> 6) & 0x3f, val & 0x3f)
234
235     # For FST-01-00 and FST-01: LED on, USB connect
236     def setup_gpio(self):
237         apb2enr = self.read_memory_u32(0x40021018)
238         apb2enr = apb2enr | 4 | 8 | 0x1000 # Enable port A, port B, and SPI1
239         self.write_memory_u32(0x40021018, apb2enr)    # RCC->APB2ENR
240         self.write_memory_u32(0x4002100c, 4|8|0x1000) # RCC->APB2RSTR
241         self.write_memory_u32(0x4002100c, 0)
242         self.write_memory_u32(GPIOA+0x0c, 0xffffffff) # ODR
243         self.write_memory_u32(GPIOA+0x04, 0x88888383) # CRH
244         self.write_memory_u32(GPIOA+0x00, 0xBBB38888) # CRL
245         self.write_memory_u32(GPIOB+0x0c, 0xffffffff) # ODR
246         self.write_memory_u32(GPIOB+0x04, 0x88888888) # CRH
247         self.write_memory_u32(GPIOB+0x00, 0x88888883) # CRL
248
249     # For FST-01-00 and FST-01: LED on, USB disconnect
250     def usb_disconnect(self):
251         self.write_memory_u32(GPIOA+0x0c, 0xfffffbff) # ODR
252
253     # For FST-01-00 and FST-01: LED off, USB connect
254     def finish_gpio(self):
255         self.write_memory_u32(GPIOA+0x0c, 0xfffffeff) # ODR
256         self.write_memory_u32(GPIOB+0x0c, 0xfffffffe) # ODR
257         apb2enr = self.read_memory_u32(0x40021018)
258         apb2enr = apb2enr &  ~(4 | 8 | 0x1000)
259         self.write_memory_u32(0x40021018, apb2enr)    # RCC->APB2ENR
260
261     def spi_flash_init(self):
262         self.write_memory_u32(SPI1+0x00, 0x0004); # CR1 <= MSTR
263         i2scfgr = self.read_memory_u32(SPI1+0x1c) # I2SCFGR
264         i2scfgr = i2scfgr & 0xf7ff                # 
265         self.write_memory_u32(SPI1+0x1c, i2scfgr); # I2SCFGR <= SPI mode
266         self.write_memory_u32(SPI1+0x10, 7);       # CRCPR <= 7
267         self.write_memory_u32(SPI1+0x04, 0x04);    # CR2 <= SSOE
268         self.write_memory_u32(SPI1+0x00, 0x0044);  # CR1 <= MSTR | SPE
269
270     def spi_flash_select(self, enable):
271         if enable:
272             self.write_memory_u32(GPIOA+0x0c, 0xffffffef) # ODR
273         else:
274             self.write_memory_u32(GPIOA+0x0c, 0xffffffff) # ODR
275
276     def spi_flash_sendbyte(self, v):
277         i = 0
278         while True:
279             status = self.read_memory_u32(SPI1+0x08) # SR
280             if status & 0x02 != 0:                   # TXE (Data Empty)
281                 break
282             time.sleep(0.01)
283             i = i + 1
284             if i > 10:
285                 raise TimeOutError('spi_flash_sendbyte')
286         self.write_memory_u32(SPI1+0x0c, v) # DR
287         i = 0
288         while True:
289             status = self.read_memory_u32(SPI1+0x08) # SR
290             if status & 0x01 != 0:                   # RXNE (Data Not Empty)
291                 break
292             time.sleep(0.01)
293             i = i + 1
294             if i > 10:
295                 raise TimeOutError('spi_flash_sendbyte')
296         v = self.read_memory_u32(SPI1+0x0c) # DR
297         return v
298
299     def spi_flash_read_id(self):
300         self.spi_flash_select(True)
301         self.spi_flash_sendbyte(0x9f)
302         t0 = self.spi_flash_sendbyte(0xa5)
303         t1 = self.spi_flash_sendbyte(0xa5)
304         t2 = self.spi_flash_sendbyte(0xa5)
305         self.spi_flash_select(False)
306         return (t0 << 16) | (t1 << 8) | t2
307
308     def protection(self):
309         return (self.read_memory_u32(FLASH_OBR) & 0x0002) != 0
310
311     def blank_check(self):
312         prog = gen_prog_blank_check(0x20000) # 128KiB  XXX: table lookup???
313         self.write_memory(SRAM_ADDRESS, prog)
314         self.write_reg(15, SRAM_ADDRESS)
315         self.run()
316         i = 0
317         while self.get_status() == 0x80:
318             time.sleep(0.050)
319             i = i + 1
320             if i >= 10:
321                 raise TimeOutError("blank check")
322
323         r0_value = self.read_reg(0)
324         return r0_value == 0
325
326     def option_bytes_read(self):
327         return self.read_memory_u32(OPTION_BYTES_ADDR)
328
329     def option_bytes_write(self,addr,val):
330         self.write_memory_u32(FLASH_KEYR, FLASH_KEY1)
331         self.write_memory_u32(FLASH_KEYR, FLASH_KEY2)
332         self.write_memory_u32(FLASH_SR, FLASH_SR_EOP | FLASH_SR_WRPRTERR | FLASH_SR_PGERR)
333
334         self.write_memory_u32(FLASH_OPTKEYR, FLASH_KEY1)
335         self.write_memory_u32(FLASH_OPTKEYR, FLASH_KEY2)
336
337         prog = gen_prog_option_bytes_write(addr,val)
338         self.write_memory(SRAM_ADDRESS, prog)
339         self.write_reg(15, SRAM_ADDRESS)
340         self.run()
341         i = 0
342         while self.get_status() == 0x80:
343             time.sleep(0.050)
344             i = i + 1
345             if i >= 10:
346                 raise TimeOutError("option bytes write")
347
348         status = self.read_memory_u32(FLASH_SR)
349         self.write_memory_u32(FLASH_CR, FLASH_CR_LOCK)
350         if (status & FLASH_SR_EOP) == 0:
351             raise OperationFailure("option bytes write")
352
353     def option_bytes_erase(self):
354         self.write_memory_u32(FLASH_KEYR, FLASH_KEY1)
355         self.write_memory_u32(FLASH_KEYR, FLASH_KEY2)
356         self.write_memory_u32(FLASH_SR, FLASH_SR_EOP | FLASH_SR_WRPRTERR | FLASH_SR_PGERR)
357
358         self.write_memory_u32(FLASH_OPTKEYR, FLASH_KEY1)
359         self.write_memory_u32(FLASH_OPTKEYR, FLASH_KEY2)
360
361         self.write_memory_u32(FLASH_CR, FLASH_CR_OPTER)
362         self.write_memory_u32(FLASH_CR, FLASH_CR_STRT | FLASH_CR_OPTER)
363
364         i = 0
365         while True:
366             status = self.read_memory_u32(FLASH_SR)
367             if (status & FLASH_SR_BSY) == 0:
368                 break
369             i = i + 1
370             if i >= 1000:
371                 break
372
373         self.write_memory_u32(FLASH_CR, FLASH_CR_LOCK)
374         if (status & FLASH_SR_EOP) == 0:
375             raise OperationFailure("option bytes erase")
376
377     def flash_write_internal(self, addr, data, off, size):
378         prog = gen_prog_flash_write(addr,size)
379         self.write_memory(SRAM_ADDRESS, prog+data[off:off+size])
380         self.write_reg(15, SRAM_ADDRESS)
381         self.run()
382         i = 0
383         while self.get_status() == 0x80:
384             time.sleep(0.050)
385             i = i + 1
386             if i >= BLOCK_WRITE_TIMEOUT:
387                 raise TimeOutError("flash write")
388         status = self.read_memory_u32(FLASH_SR)
389         if (status & FLASH_SR_PGERR) != 0:
390             raise OperationFailure("flash write: write to not erased part")
391         if (status & FLASH_SR_WRPRTERR) != 0:
392             raise OperationFailure("flash write: write to protected part")
393
394     def flash_write(self, addr, data):
395         self.write_memory_u32(FLASH_KEYR, FLASH_KEY1)
396         self.write_memory_u32(FLASH_KEYR, FLASH_KEY2)
397         self.write_memory_u32(FLASH_SR, FLASH_SR_EOP | FLASH_SR_WRPRTERR | FLASH_SR_PGERR)
398
399         off = 0
400         while True:
401             if len(data[off:]) > BLOCK_SIZE:
402                 size = BLOCK_SIZE
403                 self.flash_write_internal(addr, data, off, size)
404                 off = off + size
405                 addr = addr + size
406             else:
407                 size = len(data[off:])
408                 self.flash_write_internal(addr, data, off, size)
409                 break
410
411         self.write_memory_u32(FLASH_CR, FLASH_CR_LOCK)
412
413     def flash_erase_all(self):
414         self.write_memory_u32(FLASH_KEYR, FLASH_KEY1)
415         self.write_memory_u32(FLASH_KEYR, FLASH_KEY2)
416         self.write_memory_u32(FLASH_SR, FLASH_SR_EOP | FLASH_SR_WRPRTERR | FLASH_SR_PGERR)
417
418         self.write_memory_u32(FLASH_CR, FLASH_CR_MER)
419         self.write_memory_u32(FLASH_CR, FLASH_CR_STRT | FLASH_CR_MER)
420
421         i = 0
422         while True:
423             status = self.read_memory_u32(FLASH_SR)
424             if (status & FLASH_SR_BSY) == 0:
425                 break
426             i = i + 1
427             time.sleep(0.050)
428             if i >= 100:
429                 break
430
431         self.write_memory_u32(FLASH_CR, FLASH_CR_LOCK)
432
433         if (status & FLASH_SR_EOP) == 0:
434             raise OperationFailure("flash erase all")
435
436     def flash_erase_page(self, addr):
437         self.write_memory_u32(FLASH_KEYR, FLASH_KEY1)
438         self.write_memory_u32(FLASH_KEYR, FLASH_KEY2)
439
440         self.write_memory_u32(FLASH_SR, FLASH_SR_EOP | FLASH_SR_WRPRTERR | FLASH_SR_PGERR)
441
442         self.write_memory_u32(FLASH_CR, FLASH_CR_PER)
443         self.write_memory_u32(FLASH_AR, addr)
444         self.write_memory_u32(FLASH_CR, FLASH_CR_STRT | FLASH_CR_PER)
445
446         i = 0
447         while True:
448             status = self.read_memory_u32(FLASH_SR)
449             if (status & FLASH_SR_BSY) == 0:
450                 break
451             i = i + 1
452             if i >= 1000:
453                 break
454
455         self.write_memory_u32(FLASH_CR, FLASH_CR_LOCK)
456
457         if (status & FLASH_SR_EOP) == 0:
458             raise OperationFailure("flash page erase")
459
460     def start(self):
461         mode = self.stl_mode()
462         if mode == 2:
463             self.exit_from_debug_swd()
464         elif mode == 5:
465             self.exit_from_debug_swim()
466         elif mode != 1 and mode != 4:
467             self.exit_from_dfu()
468         new_mode = self.stl_mode()
469         print "Change ST-Link/V2 mode %04x -> %04x" % (mode, new_mode)
470         self.enter_swd()
471         s = self.get_status()
472         if s != 0x0080:
473             print "Status is %04x" % s
474             self.run()
475             s = self.get_status()
476             if s != 0x0080:
477                 raise ValueError("Status of core is not running.", s)
478         mode = self.stl_mode()
479         if mode != 2:
480             raise ValueError("Failed to switch debug mode.", mode)
481
482
483 USB_VENDOR_ST=0x0483            # 0x0483 SGS Thomson Microelectronics
484 USB_VENDOR_STLINKV2=0x3748      # 0x3748 ST-LINK/V2
485
486 def stlinkv2_devices():
487     busses = usb.busses()
488     for bus in busses:
489         devices = bus.devices
490         for dev in devices:
491             if dev.idVendor != USB_VENDOR_ST:
492                 continue
493             if dev.idProduct != USB_VENDOR_STLINKV2:
494                 continue
495             yield dev
496
497 def compare(data_original, data_in_device):
498     i = 0 
499     for d in data_original:
500         if ord(d) != data_in_device[i]:
501             raise ValueError("Verify failed at:", i)
502         i += 1
503
504 def open_stlinkv2():
505     for d in stlinkv2_devices():
506         try:
507             stl = stlinkv2(d)
508             return stl
509         except:
510             pass
511     return None
512
513 def help():
514     print "stlinkv2.py [-h]: Show this help message"
515     print "stlinkv2.py [-e]: Erase flash ROM"
516     print "stlinkv2.py [-u]: Unlock flash ROM"
517     print "stlinkv2.py [-s]: Show status"
518     print "stlinkv2.py [-b] [-n] [-r] [-i] FILE: Write content of FILE to flash ROM"
519     print "    -b: Blank check before write (auto erase when not blank)"
520     print "    -n: Don't enable read protection after write"
521     print "    -r: Don't reset after write"
522     print "    -i: Don't test SPI flash"
523
524
525 def main(show_help, erase_only, no_protect, spi_flash_check,
526          reset_after_successful_write,
527          skip_blank_check, status_only, unlock, data):
528     if show_help or len(sys.argv) != 1:
529         help()
530         return 1
531
532     stl = open_stlinkv2()
533     if not stl:
534         raise ValueError("No ST-Link/V2 device found.", None)
535
536     print "ST-Link/V2 version info: %d %d %d" % stl.version()
537     stl.start()
538     core_id = stl.core_id()
539     chip_id = stl.read_memory_u32(0xE0042000)
540
541     # FST-01 chip id: 0x20036410
542     print "CORE: %08x, CHIP_ID: %08x" % (core_id, chip_id)
543     print "Flash ROM read protection:",
544     protection = stl.protection()
545     if protection:
546         print "ON"
547     else:
548         print "off"
549
550     option_bytes = stl.option_bytes_read()
551     print "Option bytes: %08x" % option_bytes
552     if (option_bytes & 0xff) == RDP_KEY:
553         ob_protection_enable = False
554     else:
555         ob_protection_enable = True
556
557     stl.enter_debug()
558     status = stl.get_status()
559     if status != 0x0081:
560         print "Core does not halt, try API V2 halt."
561         #                   DCB_DHCSR    DBGKEY|C_HALT|C_DEBUGEN
562         stl.write_debug_reg(0xE000EDF0, 0xA05F0003)
563         status = stl.get_status()
564         if status != 0x0081:
565             raise ValueError("Status of core is not halt.", status)
566
567     if protection:
568         if status_only:
569             print "The MCU is now stopped."
570             return 0
571         elif not unlock:
572             raise OperationFailure("Flash ROM is protected")
573     else:
574         if not skip_blank_check:
575             stl.reset_sys()
576             blank = stl.blank_check()
577             print "Flash ROM blank check: %s" % blank
578         else:
579             blank = True
580         if status_only:
581             stl.reset_sys()
582             stl.run()
583             stl.exit_debug()
584             return 0
585         elif unlock and not ob_protection_enable:
586             print "No need to unlock.  Protection is not enabled."
587             return 1
588
589     if erase_only:
590         if blank:
591             print "No need to erase"
592             return 0
593
594     stl.setup_gpio()
595
596     if unlock:
597         if option_bytes != 0xff:
598             stl.reset_sys()
599             stl.option_bytes_erase()
600         stl.reset_sys()
601         stl.option_bytes_write(OPTION_BYTES_ADDR,RDP_KEY)
602         stl.usb_disconnect()
603         time.sleep(0.100)
604         stl.finish_gpio()
605         print "Flash ROM read protection disabled.  Reset the board, now."
606         return 0
607
608     if spi_flash_check:
609         stl.spi_flash_init()
610         id = stl.spi_flash_read_id()
611         print "SPI Flash ROM ID: %06x" % id
612         if id != 0xbf254a:
613             raise ValueError("bad spi flash ROM ID")
614
615     if not blank:
616         print "ERASE ALL"
617         stl.reset_sys()
618         stl.flash_erase_all()
619
620     if erase_only:
621         stl.usb_disconnect()
622         time.sleep(0.100)
623         stl.finish_gpio()
624         return 0
625
626     time.sleep(0.100)
627
628     print "WRITE"
629     stl.flash_write(0x08000000, data)
630
631     print "VERIFY"
632     data_received = ()
633     size = len(data)
634     off = 0
635     while size > 0:
636         if size > 1024:
637             blk_size = 1024
638         else:
639             blk_size = size
640         data_received = data_received + stl.read_memory(0x08000000+off, 1024)
641         size = size - blk_size
642         off = off + blk_size
643     compare(data, data_received)
644
645     if not no_protect:
646         print "PROTECT"
647         stl.option_bytes_erase()
648         print "Flash ROM read protection enabled.  Reset the board to enable protection."
649
650     if reset_after_successful_write:
651         stl.usb_disconnect()
652         stl.reset_sys()
653         stl.run()
654         stl.exit_debug()
655     else:
656         stl.finish_gpio()
657
658     stl.shutdown()
659     return 0
660
661 if __name__ == '__main__':
662     show_help = False
663     erase_only = False
664     no_protect = False
665     reset_after_successful_write = True
666     skip_blank_check=True
667     status_only = False
668     unlock = False
669     data = None
670     spi_flash_check = True
671
672     while len(sys.argv) > 1:
673         if sys.argv[1] == '-h':
674             sys.argv.pop(1)
675             show_help = True
676             break
677         elif sys.argv[1] == '-e':
678             sys.argv.pop(1)
679             erase_only = True
680             skip_blank_check=False
681             break
682         elif sys.argv[1] == '-u':
683             sys.argv.pop(1)
684             unlock = True
685             break
686         elif sys.argv[1] == '-s':
687             sys.argv.pop(1)
688             status_only = True
689             skip_blank_check=False
690             break
691         elif sys.argv[1] == '-b':
692             skip_blank_check=False
693         elif sys.argv[1] == '-n':
694             no_protect = True
695         elif sys.argv[1] == '-r':
696             reset_after_successful_write = False
697         elif sys.argv[1] == '-i':
698             spi_flash_check = False
699         else:
700             filename = sys.argv[1]
701             f = open(filename,'rb')
702             data = f.read()
703             f.close()
704         sys.argv.pop(1)
705
706     colorama_init()
707
708     try:
709         r = main(show_help, erase_only, no_protect, spi_flash_check,
710                  reset_after_successful_write,
711                  skip_blank_check, status_only, unlock, data)
712         if r == 0:
713             print Fore.WHITE + Back.BLUE + Style.BRIGHT + "SUCCESS" + Style.RESET_ALL
714         sys.exit(r)
715     except Exception as e:
716         print Back.RED + Style.BRIGHT + repr(e) + Style.RESET_ALL