tool/stlinkv2.py

   1 #! /usr/bin/python
   2
   3 """
   4 stlinkv2.py - a tool to control ST-Link/V2
   5
   6 Copyright (C) 2012 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_dfu(self):
 169         self.__devhandle.bulkWrite(self.__bulkout, "\xf3\x07", self.__timeout)
 170         time.sleep(1)
 171
 172     def enter_swd(self):
 173         self.__devhandle.bulkWrite(self.__bulkout, "\xf2\x20\xa3", self.__timeout)
 174         time.sleep(1)
 175
 176     def get_status(self):
 177         v = self.execute_get("\xf2\x01\x00", 2)
 178         return (v[1] << 8) + v[0]
 179     # RUN:128, HALT:129
 180
 181     def enter_debug(self):
 182         v = self.execute_get("\xf2\x02\x00", 2)
 183         return (v[1] << 8) + v[0]
 184
 185     def exit_debug(self):
 186         self.execute_put("\xf2\x21\x00")
 187
 188     def reset_sys(self):
 189         v = self.execute_get("\xf2\x03\x00", 2)
 190         return (v[1] << 8) + v[0]
 191
 192     def read_memory(self, addr, length):
 193         return self.execute_get("\xf2\x07" + pack('<IH', addr, length), length)
 194
 195     def read_memory_u32(self, addr):
 196         return uint32(self.execute_get("\xf2\x07" + pack('<IH', addr, 4), 4))
 197
 198     def write_memory(self, addr, data):
 199         return self.execute_put("\xf2\x08" + pack('<IH', addr, len(data)), data)
 200
 201     def write_memory_u32(self, addr, data):
 202         return self.execute_put("\xf2\x08" + pack('<IH', addr, 4),
 203                                 pack('<I', data))
 204
 205     def read_reg(self, regno):
 206         return uint32(self.execute_get("\xf2\x05" + pack('<B', regno), 4))
 207
 208     def write_reg(self, regno, value):
 209         return self.execute_get("\xf2\x06" + pack('<BI', regno, value), 2)
 210
 211     def run(self):
 212         v = self.execute_get("\xf2\x09\x00", 2)
 213         return (v[1] << 8) + v[0]
 214
 215     def core_id(self):
 216         v = self.execute_get("\xf2\x22\x00", 4)
 217         return v[0] + (v[1]<<8) + (v[2]<<16) + (v[3]<<24)
 218
 219     # For FST-01-00 and FST-01: LED on, USB connect
 220     def setup_gpio(self):
 221         apb2enr = self.read_memory_u32(0x40021018)
 222         apb2enr = apb2enr | 4 | 8 | 0x1000 # Enable port A, port B, and SPI1
 223         self.write_memory_u32(0x40021018, apb2enr)    # RCC->APB2ENR
 224         self.write_memory_u32(0x4002100c, 4|8|0x1000) # RCC->APB2RSTR
 225         self.write_memory_u32(0x4002100c, 0)
 226         self.write_memory_u32(GPIOA+0x0c, 0xffffffff) # ODR
 227         self.write_memory_u32(GPIOA+0x04, 0x88888383) # CRH
 228         self.write_memory_u32(GPIOA+0x00, 0xBBB38888) # CRL
 229         self.write_memory_u32(GPIOB+0x0c, 0xffffffff) # ODR
 230         self.write_memory_u32(GPIOB+0x04, 0x88888888) # CRH
 231         self.write_memory_u32(GPIOB+0x00, 0x88888883) # CRL
 232
 233     # For FST-01-00 and FST-01: LED on, USB disconnect
 234     def usb_disconnect(self):
 235         self.write_memory_u32(GPIOA+0x0c, 0xfffffbff) # ODR
 236
 237     # For FST-01-00 and FST-01: LED off, USB connect
 238     def finish_gpio(self):
 239         self.write_memory_u32(GPIOA+0x0c, 0xfffffeff) # ODR
 240         self.write_memory_u32(GPIOB+0x0c, 0xfffffffe) # ODR
 241         apb2enr = self.read_memory_u32(0x40021018)
 242         apb2enr = apb2enr &  ~(4 | 8 | 0x1000)
 243         self.write_memory_u32(0x40021018, apb2enr)    # RCC->APB2ENR
 244
 245     def spi_flash_init(self):
 246         self.write_memory_u32(SPI1+0x00, 0x0004); # CR1 <= MSTR
 247         i2scfgr = self.read_memory_u32(SPI1+0x1c) # I2SCFGR
 248         i2scfgr = i2scfgr & 0xf7ff                #
 249         self.write_memory_u32(SPI1+0x1c, i2scfgr); # I2SCFGR <= SPI mode
 250         self.write_memory_u32(SPI1+0x10, 7);       # CRCPR <= 7
 251         self.write_memory_u32(SPI1+0x04, 0x04);    # CR2 <= SSOE
 252         self.write_memory_u32(SPI1+0x00, 0x0044);  # CR1 <= MSTR | SPE
 253
 254     def spi_flash_select(self, enable):
 255         if enable:
 256             self.write_memory_u32(GPIOA+0x0c, 0xffffffef) # ODR
 257         else:
 258             self.write_memory_u32(GPIOA+0x0c, 0xffffffff) # ODR
 259
 260     def spi_flash_sendbyte(self, v):
 261         i = 0
 262         while True:
 263             status = self.read_memory_u32(SPI1+0x08) # SR
 264             if status & 0x02 != 0:                   # TXE (Data Empty)
 265                 break
 266             time.sleep(0.01)
 267             i = i + 1
 268             if i > 10:
 269                 raise TimeOutError('spi_flash_sendbyte')
 270         self.write_memory_u32(SPI1+0x0c, v) # DR
 271         i = 0
 272         while True:
 273             status = self.read_memory_u32(SPI1+0x08) # SR
 274             if status & 0x01 != 0:                   # RXNE (Data Not Empty)
 275                 break
 276             time.sleep(0.01)
 277             i = i + 1
 278             if i > 10:
 279                 raise TimeOutError('spi_flash_sendbyte')
 280         v = self.read_memory_u32(SPI1+0x0c) # DR
 281         return v
 282
 283     def spi_flash_read_id(self):
 284         self.spi_flash_select(True)
 285         self.spi_flash_sendbyte(0x9f)
 286         t0 = self.spi_flash_sendbyte(0xa5)
 287         t1 = self.spi_flash_sendbyte(0xa5)
 288         t2 = self.spi_flash_sendbyte(0xa5)
 289         self.spi_flash_select(False)
 290         return (t0 << 16) | (t1 << 8) | t2
 291
 292     def protection(self):
 293         return (self.read_memory_u32(FLASH_OBR) & 0x0002) != 0
 294
 295     def blank_check(self):
 296         prog = gen_prog_blank_check(0x20000) # 128KiB  XXX: table lookup???
 297         self.write_memory(SRAM_ADDRESS, prog)
 298         self.write_reg(15, SRAM_ADDRESS)
 299         self.run()
 300         i = 0
 301         while self.get_status() == 0x80:
 302             time.sleep(0.050)
 303             i = i + 1
 304             if i >= 10:
 305                 raise TimeOutError("blank check")
 306
 307         r0_value = self.read_reg(0)
 308         return r0_value == 0
 309
 310     def option_bytes_read(self):
 311         return self.read_memory_u32(OPTION_BYTES_ADDR)
 312
 313     def option_bytes_write(self,addr,val):
 314         self.write_memory_u32(FLASH_KEYR, FLASH_KEY1)
 315         self.write_memory_u32(FLASH_KEYR, FLASH_KEY2)
 316         self.write_memory_u32(FLASH_SR, FLASH_SR_EOP | FLASH_SR_WRPRTERR | FLASH_SR_PGERR)
 317
 318         self.write_memory_u32(FLASH_OPTKEYR, FLASH_KEY1)
 319         self.write_memory_u32(FLASH_OPTKEYR, FLASH_KEY2)
 320
 321         prog = gen_prog_option_bytes_write(addr,val)
 322         self.write_memory(SRAM_ADDRESS, prog)
 323         self.write_reg(15, SRAM_ADDRESS)
 324         self.run()
 325         i = 0
 326         while self.get_status() == 0x80:
 327             time.sleep(0.050)
 328             i = i + 1
 329             if i >= 10:
 330                 raise TimeOutError("option bytes write")
 331
 332         status = self.read_memory_u32(FLASH_SR)
 333         self.write_memory_u32(FLASH_CR, FLASH_CR_LOCK)
 334         if (status & FLASH_SR_EOP) == 0:
 335             raise OperationFailure("option bytes write")
 336
 337     def option_bytes_erase(self):
 338         self.write_memory_u32(FLASH_KEYR, FLASH_KEY1)
 339         self.write_memory_u32(FLASH_KEYR, FLASH_KEY2)
 340         self.write_memory_u32(FLASH_SR, FLASH_SR_EOP | FLASH_SR_WRPRTERR | FLASH_SR_PGERR)
 341
 342         self.write_memory_u32(FLASH_OPTKEYR, FLASH_KEY1)
 343         self.write_memory_u32(FLASH_OPTKEYR, FLASH_KEY2)
 344
 345         self.write_memory_u32(FLASH_CR, FLASH_CR_OPTER)
 346         self.write_memory_u32(FLASH_CR, FLASH_CR_STRT | FLASH_CR_OPTER)
 347
 348         i = 0
 349         while True:
 350             status = self.read_memory_u32(FLASH_SR)
 351             if (status & FLASH_SR_BSY) == 0:
 352                 break
 353             i = i + 1
 354             if i >= 1000:
 355                 break
 356
 357         self.write_memory_u32(FLASH_CR, FLASH_CR_LOCK)
 358         if (status & FLASH_SR_EOP) == 0:
 359             raise OperationFailure("option bytes erase")
 360
 361     def flash_write_internal(self, addr, data, off, size):
 362         prog = gen_prog_flash_write(addr,size)
 363         self.write_memory(SRAM_ADDRESS, prog+data[off:off+size])
 364         self.write_reg(15, SRAM_ADDRESS)
 365         self.run()
 366         i = 0
 367         while self.get_status() == 0x80:
 368             time.sleep(0.050)
 369             i = i + 1
 370             if i >= BLOCK_WRITE_TIMEOUT:
 371                 raise TimeOutError("flash write")
 372         status = self.read_memory_u32(FLASH_SR)
 373         if (status & FLASH_SR_PGERR) != 0:
 374             raise OperationFailure("flash write: write to not erased part")
 375         if (status & FLASH_SR_WRPRTERR) != 0:
 376             raise OperationFailure("flash write: write to protected part")
 377
 378     def flash_write(self, addr, data):
 379         self.write_memory_u32(FLASH_KEYR, FLASH_KEY1)
 380         self.write_memory_u32(FLASH_KEYR, FLASH_KEY2)
 381         self.write_memory_u32(FLASH_SR, FLASH_SR_EOP | FLASH_SR_WRPRTERR | FLASH_SR_PGERR)
 382
 383         off = 0
 384         while True:
 385             if len(data[off:]) > BLOCK_SIZE:
 386                 size = BLOCK_SIZE
 387                 self.flash_write_internal(addr, data, off, size)
 388                 off = off + size
 389                 addr = addr + size
 390             else:
 391                 size = len(data[off:])
 392                 self.flash_write_internal(addr, data, off, size)
 393                 break
 394
 395         self.write_memory_u32(FLASH_CR, FLASH_CR_LOCK)
 396
 397     def flash_erase_all(self):
 398         self.write_memory_u32(FLASH_KEYR, FLASH_KEY1)
 399         self.write_memory_u32(FLASH_KEYR, FLASH_KEY2)
 400         self.write_memory_u32(FLASH_SR, FLASH_SR_EOP | FLASH_SR_WRPRTERR | FLASH_SR_PGERR)
 401
 402         self.write_memory_u32(FLASH_CR, FLASH_CR_MER)
 403         self.write_memory_u32(FLASH_CR, FLASH_CR_STRT | FLASH_CR_MER)
 404
 405         i = 0
 406         while True:
 407             status = self.read_memory_u32(FLASH_SR)
 408             if (status & FLASH_SR_BSY) == 0:
 409                 break
 410             i = i + 1
 411             time.sleep(0.050)
 412             if i >= 100:
 413                 break
 414
 415         self.write_memory_u32(FLASH_CR, FLASH_CR_LOCK)
 416
 417         if (status & FLASH_SR_EOP) == 0:
 418             raise OperationFailure("flash erase all")
 419
 420     def flash_erase_page(self, addr):
 421         self.write_memory_u32(FLASH_KEYR, FLASH_KEY1)
 422         self.write_memory_u32(FLASH_KEYR, FLASH_KEY2)
 423
 424         self.write_memory_u32(FLASH_SR, FLASH_SR_EOP | FLASH_SR_WRPRTERR | FLASH_SR_PGERR)
 425
 426         self.write_memory_u32(FLASH_CR, FLASH_CR_PER)
 427         self.write_memory_u32(FLASH_AR, addr)
 428         self.write_memory_u32(FLASH_CR, FLASH_CR_STRT | FLASH_CR_PER)
 429
 430         i = 0
 431         while True:
 432             status = self.read_memory_u32(FLASH_SR)
 433             if (status & FLASH_SR_BSY) == 0:
 434                 break
 435             i = i + 1
 436             if i >= 1000:
 437                 break
 438
 439         self.write_memory_u32(FLASH_CR, FLASH_CR_LOCK)
 440
 441         if (status & FLASH_SR_EOP) == 0:
 442             raise OperationFailure("flash page erase")
 443
 444     def start(self):
 445         mode = self.stl_mode()
 446         if mode == 2:
 447             return
 448         elif mode != 1:
 449             self.exit_from_dfu()
 450             mode = self.stl_mode()
 451             print "Change ST-Link/V2 mode to: %04x" % mode
 452         self.enter_swd()
 453         s = self.get_status()
 454         if s != 0x0080:
 455             raise ValueError("Status of core is not running.", s)
 456         mode = self.stl_mode()
 457         if mode != 2:
 458             raise ValueError("Failed to switch debug mode.", mode)
 459
 460
 461 USB_VENDOR_ST=0x0483            # 0x0483 SGS Thomson Microelectronics
 462 USB_VENDOR_STLINKV2=0x3748      # 0x3748 ST-LINK/V2
 463
 464 def stlinkv2_devices():
 465     busses = usb.busses()
 466     for bus in busses:
 467         devices = bus.devices
 468         for dev in devices:
 469             if dev.idVendor != USB_VENDOR_ST:
 470                 continue
 471             if dev.idProduct != USB_VENDOR_STLINKV2:
 472                 continue
 473             yield dev
 474
 475 def compare(data_original, data_in_device):
 476     i = 0
 477     for d in data_original:
 478         if ord(d) != data_in_device[i]:
 479             raise ValueError("Verify failed at:", i)
 480         i += 1
 481
 482 def open_stlinkv2():
 483     for d in stlinkv2_devices():
 484         try:
 485             stl = stlinkv2(d)
 486             return stl
 487         except:
 488             pass
 489     return None
 490
 491 def help():
 492     print "stlinkv2.py [-h]: Show this help message"
 493     print "stlinkv2.py [-e]: Erase flash ROM"
 494     print "stlinkv2.py [-u]: Unlock flash ROM"
 495     print "stlinkv2.py [-s]: Show status"
 496     print "stlinkv2.py [-b] [-n] [-r] [-i] FILE: Write content of FILE to flash ROM"
 497     print "    -b: Blank check before write (auto erase when not blank)"
 498     print "    -n: Don't enable read protection after write"
 499     print "    -r: Don't reset after write"
 500     print "    -i: Don't test SPI flash"
 501
 502
 503 def main(show_help, erase_only, no_protect, spi_flash_check,
 504          reset_after_successful_write,
 505          skip_blank_check, status_only, unlock, data):
 506     if show_help or len(sys.argv) != 1:
 507         help()
 508         return 1
 509
 510     stl = open_stlinkv2()
 511     if not stl:
 512         raise ValueError("No ST-Link/V2 device found.", None)
 513
 514     stl.start()
 515     core_id = stl.core_id()
 516     chip_id = stl.read_memory_u32(0xE0042000)
 517
 518     # FST-01 chip id: 0x20036410
 519     print "CORE: %08x, CHIP_ID: %08x" % (core_id, chip_id)
 520     print "Flash ROM read protection:",
 521     protection = stl.protection()
 522     if protection:
 523         print "ON"
 524     else:
 525         print "off"
 526
 527     option_bytes = stl.option_bytes_read()
 528     print "Option bytes: %08x" % option_bytes
 529     if (option_bytes & 0xff) == RDP_KEY:
 530         ob_protection_enable = False
 531     else:
 532         ob_protection_enable = True
 533
 534     stl.enter_debug()
 535     status = stl.get_status()
 536     if status != 0x0081:
 537         raise ValueError("Status of core is not halt.", status)
 538
 539     if protection:
 540         if status_only:
 541             print "The MCU is now stopped."
 542             return 0
 543         elif not unlock:
 544             raise OperationFailure("Flash ROM is protected")
 545     else:
 546         if not skip_blank_check:
 547             stl.reset_sys()
 548             blank = stl.blank_check()
 549             print "Flash ROM blank check: %s" % blank
 550         else:
 551             blank = True
 552         if status_only:
 553             stl.reset_sys()
 554             stl.run()
 555             stl.exit_debug()
 556             return 0
 557         elif unlock and not ob_protection_enable:
 558             print "No need to unlock.  Protection is not enabled."
 559             return 1
 560
 561     if erase_only:
 562         if blank:
 563             print "No need to erase"
 564             return 0
 565
 566     stl.setup_gpio()
 567
 568     if unlock:
 569         if option_bytes != 0xff:
 570             stl.reset_sys()
 571             stl.option_bytes_erase()
 572         stl.reset_sys()
 573         stl.option_bytes_write(OPTION_BYTES_ADDR,RDP_KEY)
 574         stl.usb_disconnect()
 575         time.sleep(0.100)
 576         stl.finish_gpio()
 577         print "Flash ROM read protection disabled.  Reset the board, now."
 578         return 0
 579
 580     if spi_flash_check:
 581         stl.spi_flash_init()
 582         id = stl.spi_flash_read_id()
 583         print "SPI Flash ROM ID: %06x" % id
 584         if id != 0xbf254a:
 585             raise ValueError("bad spi flash ROM ID")
 586
 587     if not blank:
 588         print "ERASE ALL"
 589         stl.reset_sys()
 590         stl.flash_erase_all()
 591
 592     if erase_only:
 593         stl.usb_disconnect()
 594         time.sleep(0.100)
 595         stl.finish_gpio()
 596         return 0
 597
 598     time.sleep(0.100)
 599
 600     print "WRITE"
 601     stl.flash_write(0x08000000, data)
 602
 603     print "VERIFY"
 604     data_received = ()
 605     size = len(data)
 606     off = 0
 607     while size > 0:
 608         if size > 1024:
 609             blk_size = 1024
 610         else:
 611             blk_size = size
 612         data_received = data_received + stl.read_memory(0x08000000+off, 1024)
 613         size = size - blk_size
 614         off = off + blk_size
 615     compare(data, data_received)
 616
 617     if not no_protect:
 618         print "PROTECT"
 619         stl.option_bytes_erase()
 620         print "Flash ROM read protection enabled.  Reset the board to enable protection."
 621
 622     if reset_after_successful_write:
 623         stl.usb_disconnect()
 624         stl.reset_sys()
 625         stl.run()
 626         stl.exit_debug()
 627     else:
 628         stl.finish_gpio()
 629
 630     stl.shutdown()
 631     return 0
 632
 633 if __name__ == '__main__':
 634     show_help = False
 635     erase_only = False
 636     no_protect = False
 637     reset_after_successful_write = True
 638     skip_blank_check=True
 639     status_only = False
 640     unlock = False
 641     data = None
 642     spi_flash_check = True
 643
 644     while len(sys.argv) > 1:
 645         if sys.argv[1] == '-h':
 646             sys.argv.pop(1)
 647             show_help = True
 648             break
 649         elif sys.argv[1] == '-e':
 650             sys.argv.pop(1)
 651             erase_only = True
 652             skip_blank_check=False
 653             break
 654         elif sys.argv[1] == '-u':
 655             sys.argv.pop(1)
 656             unlock = True
 657             break
 658         elif sys.argv[1] == '-s':
 659             sys.argv.pop(1)
 660             status_only = True
 661             skip_blank_check=False
 662             break
 663         elif sys.argv[1] == '-b':
 664             skip_blank_check=False
 665         elif sys.argv[1] == '-n':
 666             no_protect = True
 667         elif sys.argv[1] == '-r':
 668             reset_after_successful_write = False
 669         elif sys.argv[1] == '-i':
 670             spi_flash_check = False
 671         else:
 672             filename = sys.argv[1]
 673             f = open(filename,'rb')
 674             data = f.read()
 675             f.close()
 676         sys.argv.pop(1)
 677
 678     colorama_init()
 679
 680     try:
 681         r = main(show_help, erase_only, no_protect, spi_flash_check,
 682                  reset_after_successful_write,
 683                  skip_blank_check, status_only, unlock, data)
 684         if r == 0:
 685             print Fore.WHITE + Back.BLUE + Style.BRIGHT + "SUCCESS" + Style.RESET_ALL
 686         sys.exit(r)
 687     except Exception as e:
 688         print Back.RED + Style.BRIGHT + repr(e) + Style.RESET_ALL