Fix to wait for internal http server to complete transfer.
[10.03/openwrt.git] / scripts / flashing / jungo-image.py
1 #!/usr/bin/env python
2 #
3 # Copyright 2008, 2009 (C) Jose Vasconcellos <jvasco@verizon.net>
4 #
5 # A script that can communicate with jungo-based routers
6 # (such as MI424-WR, USR8200 and WRV54G) to backup the installed
7 # firmware and replace the boot loader.
8 #
9 # Tested with Python 2.5 on Linux and Windows
10 #
11 """Usage: %s [options] <IP_address> [image.bin | url]
12 Valid options:
13 \t-h | --help: usage statement
14 \t-d | --dump: create a flash dump
15 \t-f | --file: use <filename> to store dump contents
16 \t-u | --user: provide username (default admin)
17 \t-p | --pass: provide password (default password1)
18 \t     --port: set port for http (default 8080)
19 \t-q | --quiet: don't display unnecessary information
20 \t-r | --reboot: reboot target on successful transfer
21 \t-V | --version: display version information
22
23 If no image (or url) is given, a flash dump is created.
24 A built-in http server is used when an image file is provided.
25 """
26
27 import os
28 import sys
29 import getopt
30 import getpass
31 import telnetlib
32 import string
33 import binascii
34 import socket
35 import thread
36 import SocketServer
37 import SimpleHTTPServer
38
39 reboot = 0
40 HOST = "192.168.1.1"
41 PORT = 8080
42 user = "admin"
43 #password = getpass.getpass()
44 password = "password1"
45 proto = "http"
46 url = ""
47 imagefile = ""
48 dumpfile = ""
49 verbose = 1
50 do_dump = 0
51 dumplen = 0x10000
52 flashsize=4*1024*1024
53 #device="br0"
54 device="ixp0"
55
56 ####################
57
58 def start_server(server):
59     httpd = SocketServer.TCPServer((server,PORT),SimpleHTTPServer.SimpleHTTPRequestHandler)
60     thread.start_new_thread(httpd.serve_forever,())
61
62 ####################
63
64 def get_flash_size():
65     tn.write("cat /proc/mtd\n")
66     # wait for prompt
67     buf = tn.read_until("Returned 0", 3)
68     if buf:
69         i = buf.find('mtd0:')
70         if i > 0:
71             return int(buf[i+6:].split()[0],16)
72         print "Can't find mtd0!"
73     else:
74         print "Can't access /proc/mtd!"
75     sys.exit(2)
76
77 def image_dump(tn, dumpfile):
78     if not dumpfile:
79         tn.write("ver\n");
80         buf = tn.read_until("Returned 0")
81         i = buf.find("Platform:")
82         if i < 0:
83             platform="jungo"
84         else:
85             line=buf[i+9:]
86             i=line.find('\n')
87             platform=line[:i].split()[-1]
88
89         tn.write("ifconfig -v %s\n" % device);
90         buf = tn.read_until("Returned 0")
91
92         i = buf.find("mac = 0")
93         if i > 0:
94             i += 6
95         else:
96             print "No MAC address found! (use -f option)"
97             sys.exit(1)
98         dumpfile = "%s-%s.bin" % (platform, buf[i:i+17].replace(':',''))
99     else:
100         tn.write("\n")
101
102     print "Dumping flash contents (%dMB) to %s" % (flashsize/1048576, dumpfile)
103     f = open(dumpfile, "wb")
104
105     t=flashsize/dumplen
106     for addr in range(t):
107         if verbose:
108             sys.stdout.write('\r%d%%'%(100*addr/t))
109             sys.stdout.flush()
110
111         tn.write("flash_dump -r 0x%x -l %d -4\n" % (addr*dumplen, dumplen))
112         tn.read_until("\n")
113
114         count = addr*dumplen
115         while 1:
116             buf = tn.read_until("\n")
117             if buf.strip() == "Returned 0":
118                 break
119             s = buf.split()
120             if s and s[0][-1] == ':':
121                 a=int(s[0][:-1],16)
122                 if a != count:
123                     print "Format error: %x != %x"%(a,count)
124                     sys.exit(2)
125                 count += 16
126                 f.write(binascii.a2b_hex(string.join(s[1:],'')))
127         tn.read_until(">",1)
128
129     f.close()
130     if verbose:
131         print ""
132
133 def telnet_option(sock,cmd,option):
134     #print "Option: %d %d" % (ord(cmd), ord(option))
135     if cmd == telnetlib.DO:
136         c=telnetlib.WILL
137     elif cmd == telnetlib.WILL:
138         c=telnetlib.DO
139     sock.sendall(telnetlib.IAC + c + option)
140
141 def telnet_timeout():
142     print "Fatal error: telnet timeout!"
143     sys.exit(1)
144
145 def usage():
146     print __doc__ % os.path.basename(sys.argv[0])
147
148 ####################
149
150 try:
151     opts, args = getopt.getopt(sys.argv[1:], "hdf:qp:P:rvV", \
152         ["help", "dump", "file=", "user=", "pass=", "port=",
153          "quiet=", "reboot", "verbose", "version"])
154 except getopt.GetoptError:
155     # print help information and exit:
156     usage()
157     sys.exit(1)
158
159 for o, a in opts:
160     if o in ("-h", "--help"):
161         usage()
162         sys.exit(1)
163     elif o in ("-V", "--version"):
164         print "%s: 0.9" % sys.argv[0]
165         sys.exit(1)
166     elif o in ("-d", "--no-dump"):
167         do_dump = 1
168     elif o in ("-f", "--file"):
169         dumpfile = a
170     elif o in ("-u", "--user"):
171         user = a
172     elif o in ("-p", "--pass"):
173         password = a
174     elif o == "--port":
175         PORT = int(a)
176     elif o in ("-q", "--quiet"):
177         verbose = 0
178     elif o in ("-r", "--reboot"):
179         reboot = 1
180     elif o in ("-v", "--verbose"):
181         verbose = 1
182
183 # make sure we have enough arguments
184 if len(args) > 0:
185     HOST = args[0]
186
187 if len(args) == 2:
188     if args[1].split(':')[0] in ("tftp", "http", "ftp"):
189         url = args[1]
190     else:
191         imagefile = args[1]
192 else:
193     do_dump = 1;
194
195 ####################
196 # create a telnet session to the router
197 try:
198     tn = telnetlib.Telnet(HOST)
199 except socket.error, msg:
200     print "Unable to establish telnet session to %s: %s" % (HOST, msg)
201     sys.exit(1)
202
203 tn.set_option_negotiation_callback(telnet_option)
204
205 buf = tn.read_until("Username: ", 3)
206 if not buf:
207     telnet_timeout()
208 tn.write(user+"\n")
209 if password:
210     buf = tn.read_until("Password: ", 3)
211     if not buf:
212         telnet_timeout()
213     tn.write(password+"\n")
214
215 # wait for prompt
216 buf = tn.read_until("> ", 3)
217 if not buf:
218     telnet_timeout()
219
220 flashsize = get_flash_size()
221
222 if do_dump:
223     image_dump(tn, dumpfile)
224
225 if imagefile or url:
226     splitpath = os.path.split(imagefile)
227
228     # create load command
229     if url:
230         cmd = "load -u %s -r 0\n" % (url)
231     else:
232         server = tn.get_socket().getsockname()[0]
233         cmd = "load -u http://%s:%d/%s -r 0\n" % (server, PORT, splitpath[1])
234
235         if not os.access(imagefile, os.R_OK):
236             print "File access error: %s" % (imagefile)
237             sys.exit(3)
238
239         # make sure we're in the directory where the image is located
240         if splitpath[0]:
241             os.chdir(splitpath[0])
242
243         start_server(server)
244
245     if verbose:
246         print "Unlocking flash..."
247     tn.write("unlock 0 0x%x\n" % flashsize)
248     buf = tn.read_until("Returned 0")
249
250     if verbose:
251         print "Writing new image..."
252     print cmd,
253     tn.write(cmd)
254     buf = tn.read_until("Returned 0",10)
255
256     # wait till the transfer completed
257     buf = tn.read_until("Download completed successfully",20)
258     if buf:
259         print "Flash update complete!"
260         if reboot:
261             tn.write("reboot\n")
262             print "Rebooting..."
263
264 tn.write("exit\n")
265 tn.close()
266