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