Fix for Frauchekey
[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(wait_e, 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     reg = None
314     while reg == None:
315         print("Wait %d seconds..." % wait_e)
316         time.sleep(wait_e)
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     # Then, send upgrade program...
325     mem_info = reg.mem_info()
326     print("%08x:%08x" % mem_info)
327     print("Downloading the program")
328     reg.download(mem_info[0], data_upgrade)
329     reg.protect()
330     reg.finish()
331     reg.reset_device()
332     return 0
333
334
335 # This should be event driven, not guessing some period.
336 DEFAULT_WAIT_FOR_REENUMERATION=3
337
338 if __name__ == '__main__':
339     passwd = None
340     wait_e = DEFAULT_WAIT_FOR_REENUMERATION
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         elif option == '-e':    # E for Enumeration
363             wait_e = int(sys.argv[1])
364             sys.argv.pop(1)
365         else:
366             raise ValueError("unknown option", option)
367     if not passwd:
368         passwd = getpass("Admin password: ")
369     filename_regnual = sys.argv[1]
370     filename_upgrade = sys.argv[2]
371     f = open(filename_regnual, "rb")
372     data_regnual = f.read()
373     f.close()
374     print("%s: %d" % (filename_regnual, len(data_regnual)))
375     f = open(filename_upgrade, "rb")
376     data_upgrade = f.read()
377     f.close()
378     print("%s: %d" % (filename_upgrade, len(data_upgrade)))
379     main(wait_e, passwd, data_regnual, data_upgrade[4096:])