f19bfd52e16d74a24c576caf100b322824241aea
[gnuk/neug.git] / tool / neug_upgrade.py
1 #! /usr/bin/python
2
3 """
4 neug_upgrade.py - a tool to upgrade firmware of Gnuk Token / NeuG device
5
6 Copyright (C) 2012, 2015 Free Software Initiative of Japan
7 Author: NIIBE Yutaka <gniibe@fsij.org>
8
9 This file is a part of NeuG, a TRNG 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, os, binascii, string
27 from getpass import getpass
28
29 DEFAULT_PW3 = "12345678"
30
31 # INPUT: <regnual binary file> <new firmware binary file>
32
33 # Assume only single NeuG device is attached to computer
34
35 import usb
36
37 # USB class, subclass, protocol
38 COM_CLASS = 0x0a
39 COM_SUBCLASS = 0x00
40 COM_PROTOCOL_0 = 0x00
41
42 class regnual(object):
43     def __init__(self, dev):
44         conf = dev.configurations[0]
45         intf_alt = conf.interfaces[0]
46         intf = intf_alt[0]
47         if intf.interfaceClass != 0xff:
48             raise ValueError("Wrong interface class")
49         self.__devhandle = dev.open()
50         self.__devhandle.claimInterface(intf)
51
52     def mem_info(self):
53         mem = self.__devhandle.controlMsg(requestType = 0xc0, request = 0,
54                                           buffer = 8, value = 0, index = 0,
55                                           timeout = 10000)
56         start = ((mem[3]*256 + mem[2])*256 + mem[1])*256 + mem[0]
57         end = ((mem[7]*256 + mem[6])*256 + mem[5])*256 + mem[4]
58         return (start, end)
59
60     def download(self, start, data, verbose=False):
61         addr = start
62         addr_end = (start + len(data)) & 0xffffff00
63         i = int((addr - 0x08000000) / 0x100)
64         j = 0
65         print("start %08x" % addr)
66         print("end   %08x" % addr_end)
67         while addr < addr_end:
68             if verbose:
69                 print("# %08x: %d: %d : %d" % (addr, i, j, 256))
70             self.__devhandle.controlMsg(requestType = 0x40, request = 1,
71                                         buffer = data[j*256:j*256+256],
72                                         value = 0, index = 0,
73                                         timeout = 10000)
74             crc32code = crc32(data[j*256:j*256+256])
75             res = self.__devhandle.controlMsg(requestType = 0xc0, request = 2,
76                                               buffer = 4, value = 0, index = 0,
77                                               timeout = 10000)
78             r_value = ((res[3]*256 + res[2])*256 + res[1])*256 + res[0]
79             if (crc32code ^ r_value) != 0xffffffff:
80                 print("failure")
81             self.__devhandle.controlMsg(requestType = 0x40, request = 3,
82                                         buffer = None, value = i, index = 0,
83                                         timeout = 10000)
84             time.sleep(0.010)
85             res = self.__devhandle.controlMsg(requestType = 0xc0, request = 2,
86                                               buffer = 4, value = 0, index = 0,
87                                               timeout = 10000)
88             r_value = ((res[3]*256 + res[2])*256 + res[1])*256 + res[0]
89             if r_value == 0:
90                 print("failure")
91             i = i+1
92             j = j+1
93             addr = addr + 256
94         residue = len(data) % 256
95         if residue != 0:
96             if verbose:
97                 print("# %08x: %d : %d" % (addr, i, residue))
98             self.__devhandle.controlMsg(requestType = 0x40, request = 1,
99                                         buffer = data[j*256:],
100                                         value = 0, index = 0, timeout = 10000)
101             crc32code = crc32(data[j*256:].ljust(256,b'\xff'))
102             res = self.__devhandle.controlMsg(requestType = 0xc0, request = 2,
103                                               buffer = 4, value = 0, index = 0,
104                                               timeout = 10000)
105             r_value = ((res[3]*256 + res[2])*256 + res[1])*256 + res[0]
106             if (crc32code ^ r_value) != 0xffffffff:
107                 print("failure")
108             self.__devhandle.controlMsg(requestType = 0x40, request = 3,
109                                         buffer = None, value = i, index = 0,
110                                         timeout = 10000)
111             time.sleep(0.010)
112             res = self.__devhandle.controlMsg(requestType = 0xc0, request = 2,
113                                               buffer = 4, value = 0, index = 0,
114                                               timeout = 10000)
115             r_value = ((res[3]*256 + res[2])*256 + res[1])*256 + res[0]
116             if r_value == 0:
117                 print("failure")
118
119     def protect(self):
120         self.__devhandle.controlMsg(requestType = 0x40, request = 4,
121                                     buffer = None, value = 0, index = 0,
122                                     timeout = 10000)
123         time.sleep(0.100)
124         res = self.__devhandle.controlMsg(requestType = 0xc0, request = 2,
125                                           buffer = 4, value = 0, index = 0,
126                                           timeout = 10000)
127         r_value = ((res[3]*256 + res[2])*256 + res[1])*256 + res[0]
128         if r_value == 0:
129             print("protection failure")
130
131     def finish(self):
132         self.__devhandle.controlMsg(requestType = 0x40, request = 5,
133                                     buffer = None, value = 0, index = 0,
134                                     timeout = 10000)
135
136     def reset_device(self):
137         try:
138             self.__devhandle.reset()
139         except:
140             pass
141
142 class neug(object):
143     def __init__(self, device, configuration, interface):
144         """
145         __init__(device, configuration, interface) -> None
146         Initialize the device.
147         device: usb.Device object.
148         configuration: configuration number.
149         interface: usb.Interface object representing the interface and altenate setting.
150         """
151         if interface.interfaceClass !=COM_CLASS:
152             raise ValueError("Wrong interface class")
153         if interface.interfaceSubClass != COM_SUBCLASS:
154             raise ValueError("Wrong interface sub class")
155         self.__devhandle = device.open()
156
157         self.__intf = interface.interfaceNumber
158         self.__alt = interface.alternateSetting
159         self.__conf = configuration
160
161         self.__timeout = 10000
162
163     def detach_driver(self):
164         self.__devhandle.detachKernelDriver(self.__intf)
165
166     def reset_device(self):
167         try:
168             self.__devhandle.reset()
169         except:
170             pass
171
172     def set_passwd(self, passwd):
173         self.__devhandle.controlMsg(requestType = 0x40, request = 253,
174                                     buffer = passwd, value = 0, index = 0,
175                                     timeout = 1000)
176         return
177
178     def stop_neug(self, passwd):
179         self.__devhandle.controlMsg(requestType = 0x40, request = 255,
180                                     buffer = passwd, value = 0, index = 0,
181                                     timeout = 1000)
182         # self.__devhandle.releaseInterface()
183         # self.__devhandle.setConfiguration(0)
184         return
185
186     def mem_info(self):
187         mem = self.__devhandle.controlMsg(requestType = 0xc0, request = 0,
188                                           buffer = 8, value = 0, index = 0,
189                                           timeout = 1000)
190         start = ((mem[3]*256 + mem[2])*256 + mem[1])*256 + mem[0]
191         end = ((mem[7]*256 + mem[6])*256 + mem[5])*256 + mem[4]
192         return (start, end)
193
194     def download(self, start, data, verbose=False):
195         addr = start
196         addr_end = (start + len(data)) & 0xffffff00
197         i = int((addr - 0x20000000) / 0x100)
198         j = 0
199         print("start %08x" % addr)
200         print("end   %08x" % addr_end)
201         while addr < addr_end:
202             if verbose:
203                 print("# %08x: %d : %d" % (addr, i, 256))
204             self.__devhandle.controlMsg(requestType = 0x40, request = 1,
205                                         buffer = data[j*256:j*256+256],
206                                         value = i, index = 0, timeout = 10)
207             i = i+1
208             j = j+1
209             addr = addr + 256
210         residue = len(data) % 256
211         if residue != 0:
212             print("# %08x: %d : %d" % (addr, i, residue))
213             self.__devhandle.controlMsg(requestType = 0x40, request = 1,
214                                         buffer = data[j*256:],
215                                         value = i, index = 0, timeout = 10)
216
217     def execute(self, last_addr):
218         i = int((last_addr - 0x20000000) / 0x100)
219         o = (last_addr - 0x20000000) % 0x100
220         self.__devhandle.controlMsg(requestType = 0x40, request = 2,
221                                     buffer = None, value = i, index = o,
222                                     timeout = 10)
223
224 def compare(data_original, data_in_device):
225     i = 0 
226     for d in data_original:
227         if ord(d) != data_in_device[i]:
228             raise ValueError("verify failed at %08x" % i)
229         i += 1
230
231 def com_devices():
232     busses = usb.busses()
233     for bus in busses:
234         devices = bus.devices
235         for dev in devices:
236             for config in dev.configurations:
237                 for intf in config.interfaces:
238                     for alt in intf:
239                         if alt.interfaceClass == COM_CLASS and \
240                                 alt.interfaceSubClass == COM_SUBCLASS and \
241                                 alt.interfaceProtocol == COM_PROTOCOL_0:
242                             yield dev, config, alt
243
244 USB_VENDOR_FSIJ=0x234b
245 USB_PRODUCT_GNUK=0x0000
246
247 def gnuk_devices_by_vidpid():
248     busses = usb.busses()
249     for bus in busses:
250         devices = bus.devices
251         for dev in devices:
252             if dev.idVendor != USB_VENDOR_FSIJ:
253                 continue
254             if dev.idProduct != USB_PRODUCT_GNUK:
255                 continue
256             yield dev
257
258 def to_string(t):
259     result = ""
260     for c in t:
261         result += chr(c)
262     return result
263
264 def UNSIGNED(n):
265     return n & 0xffffffff
266
267 def crc32(bytestr):
268     crc = binascii.crc32(bytestr)
269     return UNSIGNED(crc)
270
271 def main(wait_e, passwd, data_regnual, data_upgrade):
272     l = len(data_regnual)
273     if (l & 0x03) != 0:
274         data_regnual = data_regnual.ljust(l + 4 - (l & 0x03), b'\x00')
275     crc32code = crc32(data_regnual)
276     print("CRC32: %04x\n" % crc32code)
277     data_regnual += pack('<I', crc32code)
278     com = None
279     for (dev, config, intf) in com_devices():
280         try:
281             com = neug(dev, config, intf)
282             print("Device: %s" % dev.filename)
283             print("Configuration: %d" % config.value)
284             print("Interface: %d" % intf.interfaceNumber)
285             break
286         except:
287             pass
288     if not com:
289         raise ValueError("No NeuG Device Present")
290     com.stop_neug(passwd)
291     time.sleep(1.500)
292     mem_info = com.mem_info()
293     print("%08x:%08x" % mem_info)
294     com.detach_driver()
295     time.sleep(1.500)
296     print("Downloading flash upgrade program...")
297     com.download(mem_info[0], data_regnual)
298     print("Run flash upgrade program...")
299     com.execute(mem_info[0] + len(data_regnual) - 4)
300     #
301     time.sleep(3)
302     com.reset_device()
303     del com
304     com = None
305     #
306     reg = None
307     while reg == None:
308         print("Wait {} second{}...".format(wait_e, 's' if wait_e > 1 else ''))
309         time.sleep(wait_e)
310         for dev in gnuk_devices_by_vidpid():
311             try:
312                 reg = regnual(dev)
313                 print("Device: %s" % dev.filename)
314                 break
315             except:
316                 pass
317     # Then, send upgrade program...
318     mem_info = reg.mem_info()
319     print("%08x:%08x" % mem_info)
320     print("Downloading the program")
321     reg.download(mem_info[0], data_upgrade)
322     reg.protect()
323     reg.finish()
324     reg.reset_device()
325     return 0
326
327
328 # This should be event driven, not guessing some period.
329 DEFAULT_WAIT_FOR_REENUMERATION=3
330
331 if __name__ == '__main__':
332     passwd = None
333     wait_e = DEFAULT_WAIT_FOR_REENUMERATION
334     if len(sys.argv) == 2 and sys.argv[1] == '-s': # S for set passwd
335         passwd = getpass("Admin password: ")
336         com = None
337         for (dev, config, intf) in com_devices():
338             try:
339                 com = neug(dev, config, intf)
340                 print("Device: %s" % dev.filename)
341                 print("Configuration: %d" % config.value)
342                 print("Interface: %d" % intf.interfaceNumber)
343                 break
344             except:
345                 pass
346         if not com:
347             raise ValueError("No NeuG Device Present")
348         com.set_passwd(passwd)
349         exit(0)
350     while len(sys.argv) > 3:
351         option = sys.argv[1]
352         sys.argv.pop(1)
353         if option == '-f':      # F for Factory setting
354             passwd = DEFAULT_PW3
355         elif option == '-e':    # E for Enumeration
356             wait_e = int(sys.argv[1])
357             sys.argv.pop(1)
358         else:
359             raise ValueError("unknown option", option)
360     if not passwd:
361         passwd = getpass("Admin password: ")
362     filename_regnual = sys.argv[1]
363     filename_upgrade = sys.argv[2]
364     f = open(filename_regnual, "rb")
365     data_regnual = f.read()
366     f.close()
367     print("%s: %d" % (filename_regnual, len(data_regnual)))
368     f = open(filename_upgrade, "rb")
369     data_upgrade = f.read()
370     f.close()
371     print("%s: %d" % (filename_upgrade, len(data_upgrade)))
372     main(wait_e, passwd, data_regnual, data_upgrade[4096:])