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