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