Version 1.0.9.
[gnuk/neug.git] / tool / neug_upgrade.py
index 35e6723..fc43b8a 100755 (executable)
@@ -3,7 +3,7 @@
 """
 neug_upgrade.py - a tool to upgrade firmware of Gnuk Token / NeuG device
 
-Copyright (C) 2012 Free Software Initiative of Japan
+Copyright (C) 2012, 2015 Free Software Initiative of Japan
 Author: NIIBE Yutaka <gniibe@fsij.org>
 
 This file is a part of NeuG, a TRNG implementation.
@@ -24,8 +24,11 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 from struct import *
 import sys, time, os, binascii, string
+from getpass import getpass
 
-# INPUT: binary file
+DEFAULT_PW3 = "12345678"
+
+# INPUT: <regnual binary file> <new firmware binary file>
 
 # Assume only single NeuG device is attached to computer
 
@@ -42,98 +45,92 @@ class regnual(object):
         intf_alt = conf.interfaces[0]
         intf = intf_alt[0]
         if intf.interfaceClass != 0xff:
-            raise ValueError, "Wrong interface class"
+            raise ValueError("Wrong interface class")
         self.__devhandle = dev.open()
-        try:
-            self.__devhandle.setConfiguration(conf)
-        except:
-            pass
         self.__devhandle.claimInterface(intf)
-        self.__devhandle.setAltInterface(intf)
 
     def mem_info(self):
         mem = self.__devhandle.controlMsg(requestType = 0xc0, request = 0,
-                                          value = 0, index = 0, buffer = 8,
+                                          buffer = 8, value = 0, index = 0,
                                           timeout = 10000)
         start = ((mem[3]*256 + mem[2])*256 + mem[1])*256 + mem[0]
         end = ((mem[7]*256 + mem[6])*256 + mem[5])*256 + mem[4]
         return (start, end)
 
-    def download(self, start, data):
+    def download(self, start, data, verbose=False):
         addr = start
         addr_end = (start + len(data)) & 0xffffff00
-        i = (addr - 0x08000000) / 0x100
+        i = int((addr - 0x08000000) / 0x100)
         j = 0
-        print "start %08x" % addr
-        print "end   %08x" % addr_end
+        print("start %08x" % addr)
+        print("end   %08x" % addr_end)
         while addr < addr_end:
-            print "# %08x: %d: %d : %d" % (addr, i, j, 256)
+            if verbose:
+                print("# %08x: %d: %d : %d" % (addr, i, j, 256))
             self.__devhandle.controlMsg(requestType = 0x40, request = 1,
-                                        value = 0, index = 0,
                                         buffer = data[j*256:j*256+256],
+                                        value = 0, index = 0,
                                         timeout = 10000)
             crc32code = crc32(data[j*256:j*256+256])
             res = self.__devhandle.controlMsg(requestType = 0xc0, request = 2,
-                                              value = 0, index = 0, buffer = 4,
+                                              buffer = 4, value = 0, index = 0,
                                               timeout = 10000)
             r_value = ((res[3]*256 + res[2])*256 + res[1])*256 + res[0]
             if (crc32code ^ r_value) != 0xffffffff:
-                print "failure"
+                print("failure")
             self.__devhandle.controlMsg(requestType = 0x40, request = 3,
-                                        value = i, index = 0,
-                                        buffer = None,
+                                        buffer = None, value = i, index = 0,
                                         timeout = 10000)
             time.sleep(0.010)
             res = self.__devhandle.controlMsg(requestType = 0xc0, request = 2,
-                                              value = 0, index = 0, buffer = 4,
+                                              buffer = 4, value = 0, index = 0,
                                               timeout = 10000)
             r_value = ((res[3]*256 + res[2])*256 + res[1])*256 + res[0]
             if r_value == 0:
-                print "failure"
+                print("failure")
             i = i+1
             j = j+1
             addr = addr + 256
         residue = len(data) % 256
         if residue != 0:
-            print "# %08x: %d : %d" % (addr, i, residue)
+            if verbose:
+                print("# %08x: %d : %d" % (addr, i, residue))
             self.__devhandle.controlMsg(requestType = 0x40, request = 1,
-                                        value = 0, index = 0,
                                         buffer = data[j*256:],
-                                        timeout = 10000)
-            crc32code = crc32(data[j*256:].ljust(256,chr(255)))
+                                        value = 0, index = 0, timeout = 10000)
+            crc32code = crc32(data[j*256:].ljust(256,b'\xff'))
             res = self.__devhandle.controlMsg(requestType = 0xc0, request = 2,
-                                              value = 0, index = 0, buffer = 4,
+                                              buffer = 4, value = 0, index = 0,
                                               timeout = 10000)
             r_value = ((res[3]*256 + res[2])*256 + res[1])*256 + res[0]
             if (crc32code ^ r_value) != 0xffffffff:
-                print "failure"
+                print("failure")
             self.__devhandle.controlMsg(requestType = 0x40, request = 3,
-                                        value = i, index = 0,
-                                        buffer = None,
+                                        buffer = None, value = i, index = 0,
                                         timeout = 10000)
             time.sleep(0.010)
             res = self.__devhandle.controlMsg(requestType = 0xc0, request = 2,
-                                              value = 0, index = 0, buffer = 4,
+                                              buffer = 4, value = 0, index = 0,
                                               timeout = 10000)
             r_value = ((res[3]*256 + res[2])*256 + res[1])*256 + res[0]
             if r_value == 0:
-                print "failure"
+                print("failure")
 
     def protect(self):
         self.__devhandle.controlMsg(requestType = 0x40, request = 4,
-                                    value = 0, index = 0, buffer = None,
+                                    buffer = None, value = 0, index = 0,
                                     timeout = 10000)
         time.sleep(0.100)
         res = self.__devhandle.controlMsg(requestType = 0xc0, request = 2,
-                                          value = 0, index = 0, buffer = 4,
+                                          buffer = 4, value = 0, index = 0,
                                           timeout = 10000)
         r_value = ((res[3]*256 + res[2])*256 + res[1])*256 + res[0]
         if r_value == 0:
-            print "protection failure"
+            print("protection failure")
 
     def finish(self):
         self.__devhandle.controlMsg(requestType = 0x40, request = 5,
-                                    value = 0, index = 0, buffer = None,
+                                    buffer = None, value = 0, index = 0,
                                     timeout = 10000)
 
     def reset_device(self):
@@ -152,13 +149,10 @@ class neug(object):
         interface: usb.Interface object representing the interface and altenate setting.
         """
         if interface.interfaceClass !=COM_CLASS:
-            raise ValueError, "Wrong interface class"
+            raise ValueError("Wrong interface class")
         if interface.interfaceSubClass != COM_SUBCLASS:
-            raise ValueError, "Wrong interface sub class"
+            raise ValueError("Wrong interface sub class")
         self.__devhandle = device.open()
-        self.__devhandle.detachKernelDriver(interface)
-        # self.__devhandle.claimInterface(interface)
-        # self.__devhandle.setAltInterface(interface)
 
         self.__intf = interface.interfaceNumber
         self.__alt = interface.alternateSetting
@@ -166,64 +160,72 @@ class neug(object):
 
         self.__timeout = 10000
 
+    def detach_driver(self):
+        self.__devhandle.detachKernelDriver(self.__intf)
+
     def reset_device(self):
         try:
             self.__devhandle.reset()
         except:
             pass
 
-    def stop_neug(self):
+    def set_passwd(self, passwd):
+        self.__devhandle.controlMsg(requestType = 0x40, request = 253,
+                                    buffer = passwd, value = 0, index = 0,
+                                    timeout = 1000)
+        return
+
+    def stop_neug(self, passwd):
         self.__devhandle.controlMsg(requestType = 0x40, request = 255,
-                                    value = 0, index = 0, buffer = None,
-                                    timeout = 10)
+                                    buffer = passwd, value = 0, index = 0,
+                                    timeout = 1000)
         # self.__devhandle.releaseInterface()
         # self.__devhandle.setConfiguration(0)
         return
 
     def mem_info(self):
         mem = self.__devhandle.controlMsg(requestType = 0xc0, request = 0,
-                                          value = 0, index = 0, buffer = 8,
-                                          timeout = 10)
+                                          buffer = 8, value = 0, index = 0,
+                                          timeout = 1000)
         start = ((mem[3]*256 + mem[2])*256 + mem[1])*256 + mem[0]
         end = ((mem[7]*256 + mem[6])*256 + mem[5])*256 + mem[4]
         return (start, end)
 
-    def download(self, start, data):
+    def download(self, start, data, verbose=False):
         addr = start
         addr_end = (start + len(data)) & 0xffffff00
-        i = (addr - 0x20000000) / 0x100
+        i = int((addr - 0x20000000) / 0x100)
         j = 0
-        print "start %08x" % addr
-        print "end   %08x" % addr_end
+        print("start %08x" % addr)
+        print("end   %08x" % addr_end)
         while addr < addr_end:
-            print "# %08x: %d : %d" % (addr, i, 256)
+            if verbose:
+                print("# %08x: %d : %d" % (addr, i, 256))
             self.__devhandle.controlMsg(requestType = 0x40, request = 1,
-                                        value = i, index = 0,
                                         buffer = data[j*256:j*256+256],
-                                        timeout = 10)
+                                        value = i, index = 0, timeout = 10)
             i = i+1
             j = j+1
             addr = addr + 256
         residue = len(data) % 256
         if residue != 0:
-            print "# %08x: %d : %d" % (addr, i, residue)
+            print("# %08x: %d : %d" % (addr, i, residue))
             self.__devhandle.controlMsg(requestType = 0x40, request = 1,
-                                        value = i, index = 0,
                                         buffer = data[j*256:],
-                                        timeout = 10)
+                                        value = i, index = 0, timeout = 10)
 
     def execute(self, last_addr):
-        i = (last_addr - 0x20000000) / 0x100
+        i = int((last_addr - 0x20000000) / 0x100)
         o = (last_addr - 0x20000000) % 0x100
         self.__devhandle.controlMsg(requestType = 0x40, request = 2,
-                                    value = i, index = o, buffer = None,
+                                    buffer = None, value = i, index = o,
                                     timeout = 10)
 
 def compare(data_original, data_in_device):
     i = 0 
     for d in data_original:
         if ord(d) != data_in_device[i]:
-            raise ValueError, "verify failed at %08x" % i
+            raise ValueError("verify failed at %08x" % i)
         i += 1
 
 def com_devices():
@@ -242,7 +244,7 @@ def com_devices():
 USB_VENDOR_FSIJ=0x234b
 USB_PRODUCT_GNUK=0x0000
 
-def gnuk_devices():
+def gnuk_devices_by_vidpid():
     busses = usb.busses()
     for bus in busses:
         devices = bus.devices
@@ -266,32 +268,34 @@ def crc32(bytestr):
     crc = binascii.crc32(bytestr)
     return UNSIGNED(crc)
 
-def main(data_regnual, data_upgrade):
+def main(wait_e, passwd, data_regnual, data_upgrade):
     l = len(data_regnual)
     if (l & 0x03) != 0:
-        data_regnual = data_regnual.ljust(l + 4 - (l & 0x03), chr(0))
+        data_regnual = data_regnual.ljust(l + 4 - (l & 0x03), b'\x00')
     crc32code = crc32(data_regnual)
-    print "CRC32: %04x\n" % crc32code
+    print("CRC32: %04x\n" % crc32code)
     data_regnual += pack('<I', crc32code)
     com = None
     for (dev, config, intf) in com_devices():
         try:
             com = neug(dev, config, intf)
-            print "Device: ", dev.filename
-            print "Configuration: ", config.value
-            print "Interface: ", intf.interfaceNumber
+            print("Device: %s" % dev.filename)
+            print("Configuration: %d" % config.value)
+            print("Interface: %d" % intf.interfaceNumber)
             break
         except:
             pass
     if not com:
-        raise ValueError, "No NeuG Device Present"
-    com.stop_neug()
-    time.sleep(0.500)
+        raise ValueError("No NeuG Device Present")
+    com.stop_neug(passwd)
+    time.sleep(1.500)
     mem_info = com.mem_info()
-    print "%08x:%08x" % mem_info
-    print "Downloading flash upgrade program..."
+    print("%08x:%08x" % mem_info)
+    com.detach_driver()
+    time.sleep(1.500)
+    print("Downloading flash upgrade program...")
     com.download(mem_info[0], data_regnual)
-    print "Run flash upgrade program..."
+    print("Run flash upgrade program...")
     com.execute(mem_info[0] + len(data_regnual) - 4)
     #
     time.sleep(3)
@@ -299,20 +303,21 @@ def main(data_regnual, data_upgrade):
     del com
     com = None
     #
-    print "Wait 3 seconds..."
-    time.sleep(3)
-    # Then, send upgrade program...
     reg = None
-    for dev in gnuk_devices():
-        try:
-            reg = regnual(dev)
-            print "Device: ", dev.filename
-            break
-        except:
-            pass
+    while reg == None:
+        print("Wait {} second{}...".format(wait_e, 's' if wait_e > 1 else ''))
+        time.sleep(wait_e)
+        for dev in gnuk_devices_by_vidpid():
+            try:
+                reg = regnual(dev)
+                print("Device: %s" % dev.filename)
+                break
+            except:
+                pass
+    # Then, send upgrade program...
     mem_info = reg.mem_info()
-    print "%08x:%08x" % mem_info
-    print "Downloading the program"
+    print("%08x:%08x" % mem_info)
+    print("Downloading the program")
     reg.download(mem_info[0], data_upgrade)
     reg.protect()
     reg.finish()
@@ -320,15 +325,48 @@ def main(data_regnual, data_upgrade):
     return 0
 
 
+# This should be event driven, not guessing some period.
+DEFAULT_WAIT_FOR_REENUMERATION=1
+
 if __name__ == '__main__':
+    passwd = None
+    wait_e = DEFAULT_WAIT_FOR_REENUMERATION
+    if len(sys.argv) == 2 and sys.argv[1] == '-s': # S for set passwd
+        passwd = getpass("Admin password: ")
+        com = None
+        for (dev, config, intf) in com_devices():
+            try:
+                com = neug(dev, config, intf)
+                print("Device: %s" % dev.filename)
+                print("Configuration: %d" % config.value)
+                print("Interface: %d" % intf.interfaceNumber)
+                break
+            except:
+                pass
+        if not com:
+            raise ValueError("No NeuG Device Present")
+        com.set_passwd(passwd)
+        exit(0)
+    while len(sys.argv) > 3:
+        option = sys.argv[1]
+        sys.argv.pop(1)
+        if option == '-f':      # F for Factory setting
+            passwd = DEFAULT_PW3
+        elif option == '-e':    # E for Enumeration
+            wait_e = int(sys.argv[1])
+            sys.argv.pop(1)
+        else:
+            raise ValueError("unknown option", option)
+    if not passwd:
+        passwd = getpass("Admin password: ")
     filename_regnual = sys.argv[1]
     filename_upgrade = sys.argv[2]
-    f = open(filename_regnual)
+    f = open(filename_regnual, "rb")
     data_regnual = f.read()
     f.close()
-    print "%s: %d" % (filename_regnual, len(data_regnual))
-    f = open(filename_upgrade)
+    print("%s: %d" % (filename_regnual, len(data_regnual)))
+    f = open(filename_upgrade, "rb")
     data_upgrade = f.read()
     f.close()
-    print "%s: %d" % (filename_upgrade, len(data_upgrade))
-    main(data_regnual, data_upgrade[4096:])
+    print("%s: %d" % (filename_upgrade, len(data_upgrade)))
+    main(wait_e, passwd, data_regnual, data_upgrade[4096:])