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         # use different command
  73         tn.write("flash_layout\n")
  74         buf = tn.read_until("Returned 0", 3)
  75         i = buf.rfind('Range ')
  76         if i > 0:
  77             return int(buf[i+17:].split()[0],16)
  78         print "Can't determine flash size!"
  79     else:
  80         print "Unable to obtain flash size!"
  81     sys.exit(2)
  82
  83 def image_dump(tn, dumpfile):
  84     if not dumpfile:
  85         tn.write("ver\n");
  86         buf = tn.read_until("Returned 0",2)
  87         i = buf.find("Platform:")
  88         if i < 0:
  89             platform="jungo"
  90         else:
  91             line=buf[i+9:]
  92             i=line.find('\n')
  93             platform=line[:i].split()[-1]
  94
  95         tn.write("rg_conf_print /dev/%s/mac\n" % device);
  96         buf = tn.read_until("Returned 0",3)
  97
  98         i = buf.find("mac(")
  99         if i > 0:
 100             i += 4
 101         else:
 102             print "No MAC address found! (use -f option)"
 103             sys.exit(1)
 104         dumpfile = "%s-%s.bin" % (platform, buf[i:i+17].replace(':',''))
 105     else:
 106         tn.write("\n")
 107
 108     print "Dumping flash contents (%dMB) to %s" % (flashsize/1048576, dumpfile)
 109     f = open(dumpfile, "wb")
 110
 111     t=flashsize/dumplen
 112     for addr in range(t):
 113         if verbose:
 114             sys.stdout.write('\r%d%%'%(100*addr/t))
 115             sys.stdout.flush()
 116
 117         tn.write("flash_dump -r 0x%x -l %d -4\n" % (addr*dumplen, dumplen))
 118         tn.read_until("\n")
 119
 120         count = addr*dumplen
 121         while 1:
 122             buf = tn.read_until("\n")
 123             if buf.strip() == "Returned 0":
 124                 break
 125             s = buf.split()
 126             if s and s[0][-1] == ':':
 127                 a=int(s[0][:-1],16)
 128                 if a != count:
 129                     print "Format error: %x != %x"%(a,count)
 130                     sys.exit(2)
 131                 count += 16
 132                 f.write(binascii.a2b_hex(string.join(s[1:],'')))
 133         tn.read_until(">",1)
 134
 135     f.close()
 136     if verbose:
 137         print ""
 138
 139 def telnet_option(sock,cmd,option):
 140     #print "Option: %d %d" % (ord(cmd), ord(option))
 141     if cmd == telnetlib.DO:
 142         c=telnetlib.WILL
 143     elif cmd == telnetlib.WILL:
 144         c=telnetlib.DO
 145     sock.sendall(telnetlib.IAC + c + option)
 146
 147 def telnet_timeout():
 148     print "Fatal error: telnet timeout!"
 149     sys.exit(1)
 150
 151 def usage():
 152     print __doc__ % os.path.basename(sys.argv[0])
 153
 154 ####################
 155
 156 try:
 157     opts, args = getopt.getopt(sys.argv[1:], "hdf:qp:P:rvV", \
 158         ["help", "dump", "file=", "user=", "pass=", "port=",
 159          "quiet=", "reboot", "verbose", "version"])
 160 except getopt.GetoptError:
 161     # print help information and exit:
 162     usage()
 163     sys.exit(1)
 164
 165 for o, a in opts:
 166     if o in ("-h", "--help"):
 167         usage()
 168         sys.exit(1)
 169     elif o in ("-V", "--version"):
 170         print "%s: 0.10" % sys.argv[0]
 171         sys.exit(1)
 172     elif o in ("-d", "--no-dump"):
 173         do_dump = 1
 174     elif o in ("-f", "--file"):
 175         dumpfile = a
 176     elif o in ("-u", "--user"):
 177         user = a
 178     elif o in ("-p", "--pass"):
 179         password = a
 180     elif o == "--port":
 181         PORT = int(a)
 182     elif o in ("-q", "--quiet"):
 183         verbose = 0
 184     elif o in ("-r", "--reboot"):
 185         reboot = 1
 186     elif o in ("-v", "--verbose"):
 187         verbose = 1
 188
 189 # make sure we have enough arguments
 190 if len(args) > 0:
 191     HOST = args[0]
 192
 193 if len(args) == 2:
 194     if args[1].split(':')[0] in ("tftp", "http", "ftp"):
 195         url = args[1]
 196     else:
 197         imagefile = args[1]
 198 else:
 199     do_dump = 1;
 200
 201 ####################
 202 # create a telnet session to the router
 203 try:
 204     tn = telnetlib.Telnet(HOST)
 205 except socket.error, msg:
 206     print "Unable to establish telnet session to %s: %s" % (HOST, msg)
 207     sys.exit(1)
 208
 209 tn.set_option_negotiation_callback(telnet_option)
 210
 211 buf = tn.read_until("Username: ", 3)
 212 if not buf:
 213     telnet_timeout()
 214 tn.write(user+"\n")
 215 if password:
 216     buf = tn.read_until("Password: ", 3)
 217     if not buf:
 218         telnet_timeout()
 219     tn.write(password+"\n")
 220
 221 # wait for prompt
 222 buf = tn.read_until("> ", 3)
 223 if not buf:
 224     telnet_timeout()
 225
 226 flashsize = get_flash_size()
 227
 228 if do_dump:
 229     image_dump(tn, dumpfile)
 230
 231 if imagefile or url:
 232     splitpath = os.path.split(imagefile)
 233
 234     # create load command
 235     if url:
 236         cmd = "load -u %s -r 0\n" % (url)
 237     else:
 238         server = tn.get_socket().getsockname()[0]
 239         cmd = "load -u http://%s:%d/%s -r 0\n" % (server, PORT, splitpath[1])
 240
 241         if not os.access(imagefile, os.R_OK):
 242             print "File access error: %s" % (imagefile)
 243             sys.exit(3)
 244
 245         # make sure we're in the directory where the image is located
 246         if splitpath[0]:
 247             os.chdir(splitpath[0])
 248
 249         start_server(server)
 250
 251     if verbose:
 252         print "Unlocking flash..."
 253     tn.write("unlock 0 0x%x\n" % flashsize)
 254     buf = tn.read_until("Returned 0",5)
 255
 256     if verbose:
 257         print "Writing new image..."
 258     print cmd,
 259     tn.write(cmd)
 260     buf = tn.read_until("Returned 0",10)
 261
 262     # wait till the transfer completed
 263     buf = tn.read_until("Download completed successfully",20)
 264     if buf:
 265         print "Flash update complete!"
 266         if reboot:
 267             tn.write("reboot\n")
 268             print "Rebooting..."
 269
 270 tn.write("exit\n")
 271 tn.close()
 272