cros_image_to_target handles test images, less verbose

bin/cros_image_to_target.py

 #!/usr/bin/python
 #
-# Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
+# Copyright (c) 2010 The Chromium OS Authors. All rights reserved
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Create and copy update image to target host.
 
 auto-update and devserver change out from beneath us often enough
 that despite having to duplicate a litte code, it seems that the
 right thing to do here is to start over and do something that is
 simple enough and easy enough to understand so that when more
 stuff breaks, at least we can solve them faster.
 """
 
 import BaseHTTPServer
 import cgi
 import errno
 import optparse
 import os
 import signal
 import subprocess
 import sys
 import tempfile
 import time
 import traceback
 
 from xml.dom import minidom
 
 
 # This is the default filename within the image directory to load updates from
 DEFAULT_IMAGE_NAME = 'chromiumos_image.bin'
+DEFAULT_IMAGE_NAME_TEST = 'chromiumos_test_image.bin'
 
 # The filenames we provide to clients to pull updates
 UPDATE_FILENAME = 'update.gz'
-STATEFUL_FILENAME = 'stateful.image.gz'
+STATEFUL_FILENAME = 'stateful.gz'
 
 # How long do we wait for the server to start before launching client
 SERVER_STARTUP_WAIT = 1
 
 
 class Command(object):
   """Shell command ease-ups for Python."""
 
   def __init__(self, env):
     self.env = env
 
   def RunPipe(self, pipeline, infile=None, outfile=None,
-              capture=False, oneline=False):
-    """Perform a command pipeline, with optional input/output filenames."""
+              capture=False, oneline=False, hide_stderr=False):
+    """
+    Perform a command pipeline, with optional input/output filenames.
+
+    hide_stderr     Don't allow output of stderr (default False)
+    """
 
     last_pipe = None
     while pipeline:
       cmd = pipeline.pop(0)
       kwargs = {}
       if last_pipe is not None:
         kwargs['stdin'] = last_pipe.stdout
       elif infile:
         kwargs['stdin'] = open(infile, 'rb')
       if pipeline or capture:
         kwargs['stdout'] = subprocess.PIPE
       elif outfile:
         kwargs['stdout'] = open(outfile, 'wb')
+      if hide_stderr:
+        kwargs['stderr'] = open('/dev/null', 'wb')
 
-      self.env.Info('Running: %s' % ' '.join(cmd))
+      self.env.Debug('Running: %s' % ' '.join(cmd))
       last_pipe = subprocess.Popen(cmd, **kwargs)
 
     if capture:
       ret = last_pipe.communicate()[0]
       if not ret:
         return None
       elif oneline:
         return ret.rstrip('\r\n')
       else:
         return ret
@@ skipping 27 matching lines @@
     self.known_hosts = os.path.join(self.ssh_dir, 'known-hosts')
 
   def Cleanup(self):
     Command.RunPipe(self, [['rm', '-rf', self.ssh_dir]])
     self.ssh_dir = None
 
   def GetArgs(self):
     if not self.ssh_dir:
       self.Setup()
 
+    # allow up to 3 seconds to connect
     return ['-o', 'Compression=no',
             '-o', 'ConnectTimeout=%d' % self.CONNECT_TIMEOUT,
             '-o', 'StrictHostKeyChecking=no',
+            '-o', 'ConnectTimeout=3',

[Paul Stewart, 2010/11/17 21:37:36]

Look 2 lines above this one.  I even used a constant. :-)

[sjg, 2010/11/23 18:49:14]

Sorry! It was hanging at some point and I was looking for why...

             '-o', 'UserKnownHostsFile=%s' % self.known_hosts,
             '-i', self.identity]
 
   def RunPipe(self, pipeline, **kwargs):
     args = ['ssh'] + self.GetArgs()
     if 'remote_tunnel' in kwargs:
       ports = kwargs.pop('remote_tunnel')
       args += ['-R %d:localhost:%d' % ports]
     pipeline[0] = args + ['root@%s' % self.remote] + list(pipeline[0])
     return Command.RunPipe(self, pipeline, **kwargs)
 
   def Reset(self):
     os.unlink(self.known_hosts)
 
   def Copy(self, src, dest):
     return Command.RunPipe(self, [['scp'] + self.GetArgs() +
                                   [src, 'root@%s:%s' %
                                    (self.remote, dest)]])
 
 
 class CrosEnv(object):
   """Encapsulates the ChromeOS build system environment functionality."""
 
   REBOOT_START_WAIT = 5
   REBOOT_WAIT_TIME = 60
 
-  def __init__(self, verbose=False):
+  def __init__(self, verbose=0):
     self.cros_root = os.path.dirname(os.path.abspath(sys.argv[0]))
     parent = os.path.dirname(self.cros_root)
     if os.path.exists(os.path.join(parent, 'chromeos-common.sh')):
       self.cros_root = parent
     self.cmd = Command(self)
     self.verbose = verbose
 
+    # do we have the pv progress tool? (sudo apt-get install pv)
+    self.have_pv = (os.system('pv --help >/dev/null 2>&1') == 0)

[Paul Stewart, 2010/11/17 21:37:36]

I'd much prefer if you used self.cmd.Run() instead

+
   def Error(self, msg):
     print >> sys.stderr, 'ERROR: %s' % msg
 
   def Fatal(self, msg=None):
     if msg:
       self.Error(msg)
     sys.exit(1)
 
   def Info(self, msg):
-    if self.verbose:
+    if self.verbose > 0:
       print 'INFO: %s' % msg
 
+  def Debug(self, msg):
+    if self.verbose > 1:
+      print 'DEBUG: %s' % msg
+
   def CrosUtilsPath(self, filename):
     return os.path.join(self.cros_root, filename)
 
   def ChrootPath(self, filename):
     return self.CrosUtilsPath(os.path.join('..', '..', 'chroot',
                                            filename.strip(os.path.sep)))
 
   def FileOneLine(self, filename):
     return file(filename).read().rstrip('\r\n')
 
@@ skipping 20 matching lines @@
 
     return True
 
   def BuildStateful(self, src, dst):
     """Create a stateful partition update image."""
 
     if self.GetCached(src, dst):
       self.Info('Using cached stateful %s' % dst)
       return True
 
+    self.Info('Building stateful')
     cgpt = self.ChrootPath('/usr/bin/cgpt')
     offset = self.cmd.OutputOneLine(cgpt, 'show', '-b', '-i', '1', src)
     size = self.cmd.OutputOneLine(cgpt, 'show', '-s', '-i', '1', src)
     if None in (size, offset):
       self.Error('Unable to use cgpt to get image geometry')
       return False
 
-    return self.cmd.RunPipe([['dd', 'if=%s' % src, 'bs=512',
-                              'skip=%s' % offset, 'count=%s' % size],
-                             ['gzip', '-c']], outfile=dst)
+    return self.cmd.RunPipe([
+        ['dd', 'if=%s' % src, 'bs=512', 'skip=%s' % offset, 'count=%s' % size],
+        self.have_pv and ['pv', '-s', '%d' % (int (size) * 512)] or ['cat'],
+        ['gzip', '-c']], outfile=dst, hide_stderr=True)
 
   def GetSize(self, filename):
     return os.path.getsize(filename)
 
   def GetHash(self, filename):
     return self.cmd.RunPipe([['openssl', 'sha1', '-binary'],
                              ['openssl', 'base64']],
                             infile=filename,
                             capture=True, oneline=True)
 
@@ skipping 33 matching lines @@
   def CreateServer(self, port, update_file, stateful_file):
     """Start the devserver clone."""
 
     PingUpdateResponse.Setup(self.GetHash(update_file),
                              self.GetSha256(update_file),
                              self.GetSize(update_file))
 
     UpdateHandler.SetupUrl('/update', PingUpdateResponse())
     UpdateHandler.SetupUrl('/%s' % UPDATE_FILENAME,
                            FileUpdateResponse(update_file,
-                                              verbose=self.verbose))
+                                              verbose=self.verbose, have_pv=self.have_pv))

[adlr, 2010/11/17 20:43:25]

80 cols here, and a couple lines down as well

     UpdateHandler.SetupUrl('/%s' % STATEFUL_FILENAME,
                            FileUpdateResponse(stateful_file,
-                                              verbose=self.verbose))
+                                              verbose=self.verbose, have_pv=self.have_pv))
 
     self.http_server = BaseHTTPServer.HTTPServer(('', port), UpdateHandler)
 
   def StartServer(self):
     self.Info('Starting http server')
     self.http_server.serve_forever()
 
   def GetUpdateStatus(self):
     status = self.ssh_cmd.Output('/usr/bin/update_engine_client', '--status')
     if not status:
@@ skipping 18 matching lines @@
       self.Info('Client has not yet restarted (try %d).  Waiting...' % attempt)
       wait_time = SSHCommand.CONNECT_TIMEOUT - (time.time() - start)
       if wait_time > 0:
         time.sleep(wait_time)
 
     return False
 
   def StartClient(self, port):
     """Ask the client machine to update from our server."""
 
+    self.Info ("Starting client...")

[adlr, 2010/11/17 20:43:25]

i think style says no spaces between method name and (. Also, I think single
quotes are preferred in python

     status = self.GetUpdateStatus()
     if status != 'UPDATE_STATUS_IDLE':
       self.Error('Client update status is not IDLE: %s' % status)
       return False
 
     url_base = 'http://localhost:%d' % port
     update_url = '%s/update' % url_base
     fd, update_log = tempfile.mkstemp(prefix='image-to-target-')
     self.Info('Starting update on client.  Client output stored to %s' %
               update_log)
     self.ssh_cmd.Run('/usr/bin/update_engine_client', '--update',
                      '--omaha_url', update_url, remote_tunnel=(port, port),
                      outfile=update_log)
 
     if self.GetUpdateStatus() != 'UPDATE_STATUS_UPDATED_NEED_REBOOT':
       self.Error('Client update failed')
       return False
 
+    self.Info('Update complete - running update script on client')
     self.ssh_cmd.Copy(self.CrosUtilsPath('../platform/dev/stateful_update'),
                       '/tmp')
     if not self.ssh_cmd.Run('/tmp/stateful_update', url_base,
                             remote_tunnel=(port, port)):
       self.Error('Client stateful update failed')
       return False
 
     self.Info('Rebooting client')
     if not self.ClientReboot():
       self.Error('Client may not have successfully rebooted...')
       return False
 
-    print 'Client update completed successfully!'
+    self.Info ('Client update completed successfully!')

[adlr, 2010/11/17 20:43:25]

remove space before (

     return True
 
 
 class UpdateResponse(object):
   """Default response is the 404 error response."""
 
   def Reply(self, handler, send_content=True, post_data=None):
-    handler.send_Error(404, 'File not found')
+    handler.send_error(404, 'File not found')
     return None
 
 
 class FileUpdateResponse(UpdateResponse):
   """Respond by sending the contents of a file."""
 
   def __init__(self, filename, content_type='application/octet-stream',
-               verbose=False, blocksize=16*1024):
+               verbose=False, blocksize=16*1024, have_pv=False):
     self.filename = filename
     self.content_type = content_type
     self.verbose = verbose
     self.blocksize = blocksize
+    self.have_pv = have_pv
 
   def Reply(self, handler, send_content=True, post_data=None):
     """Return file contents to the client.  Optionally display progress."""
 
     try:
       f = open(self.filename, 'rb')
     except IOError:
       return UpdateResponse.Reply(self, handler)
 
     handler.send_response(200)
     handler.send_header('Content-type', self.content_type)
     filestat = os.fstat(f.fileno())
     filesize = filestat[6]
     handler.send_header('Content-Length', str(filesize))
     handler.send_header('Last-Modified',
                         handler.date_time_string(filestat.st_mtime))
     handler.end_headers()
 
-    if not send_content:
-      return
+    if send_content:
+      try:
+        fp = False
+        sent_size = 0
+        sent_percentage = None
 
-    if filesize <= self.blocksize:
-      handler.wfile.write(f.read())
-    else:
-      sent_size = 0
-      sent_percentage = None
-      while True:
-        buf = f.read(self.blocksize)
-        if not buf:
-          break
-        handler.wfile.write(buf)
-        if self.verbose:
-          sent_size += len(buf)
-          percentage = int(100 * sent_size / filesize)
-          if sent_percentage != percentage:
-            sent_percentage = percentage
-            print '\rSent %d%%' % sent_percentage,
-            sys.stdout.flush()
+        # sadly pv gets caught up in the server, not sure why
+        if False: # self.verbose and self.have_pv:
+          fp = subprocess.Popen('pv -s %d >/dev/null' % filesize,
+                shell=True, stdin=subprocess.PIPE, close_fds=True).stdin
+        while True:
+          buf = f.read(self.blocksize)
+          if not buf:
+            break
+          handler.wfile.write(buf)
+          if fp:
+            fp.write(buf)
+          if self.verbose:
+            sent_size += len(buf)
+            percentage = int(100 * sent_size / filesize)
+            if sent_percentage != percentage:
+              sent_percentage = percentage
+              print '\rSent %d%%' % sent_percentage,
+              sys.stdout.flush()
+      finally:
+        if fp:
+          fp.close()

[Paul Stewart, 2010/11/17 21:37:36]

If you can't get it working, please revert this change.

       if self.verbose:
         print '\n'
     f.close()
 
 
 class StringUpdateResponse(UpdateResponse):
   """Respond by sending the contents of a string."""
 
   def __init__(self, string, content_type='text/plain'):
     self.string = string
@@ skipping 143 matching lines @@
   usage = 'usage: %prog [options]'
   parser = optparse.OptionParser(usage=usage)
   parser.add_option('--board', dest='board', default=None,
                     help='Board platform type')
   parser.add_option('--force-mismatch', dest='force_mismatch', default=False,
                     action='store_true',
                     help='Upgrade even if client arch does not match')
   parser.add_option('--from', dest='src', default=None,
                     help='Source image to install')
   parser.add_option('--image-name', dest='image_name',
-                    default=DEFAULT_IMAGE_NAME,
                     help='Filename within image directory to load')
   parser.add_option('--port', dest='port', default=8081, type='int',
                     help='TCP port to serve from and tunnel through')
   parser.add_option('--remote', dest='remote', default=None,
                     help='Remote device-under-test IP address')
   parser.add_option('--server-only', dest='server_only', default=False,
                     action='store_true', help='Do not start client')
   parser.add_option('--verbose', dest='verbose', default=False,
+                    action='store_true', help='Display progress')
+  parser.add_option('--debug', dest='debug', default=False,
                     action='store_true', help='Display running commands')
+  parser.add_option('--test', dest='test', default=False,
+                    action='store_true', help='Select test image')
 
   (options, args) = parser.parse_args(argv)
 
-  cros_env = CrosEnv(verbose=options.verbose)
+  # we can build the test image if it doesn't exist, so remember if we want to
+  maybe_build_test = False

[adlr, 2010/11/17 20:43:25]

it would be nice to have a variable name that more accurately describes its use
(the comment doesn't seem to). would build_test be a better name?

+
+  verbosity = 0
+  if options.verbose:
+      verbosity = 1
+  if options.debug:
+      verbosity = 2
+  cros_env = CrosEnv(verbose=verbosity)
 
   if not options.board:
     options.board = cros_env.GetDefaultBoard()
 
   if not options.src:
     options.src = cros_env.GetLatestImage(options.board)
     if options.src is None:
       parser.error('No --from argument given and no default image found')
 
   cros_env.Info('Performing update from %s' % options.src)
 
   if not os.path.exists(options.src):
     parser.error('Path %s does not exist' % options.src)
 
+  if not options.image_name:
+    # auto-select the correct image
+    if options.test:
+      options.image_name = DEFAULT_IMAGE_NAME_TEST
+
+      # we will build the test image if not found
+      maybe_build_test = True
+    else:
+      options.image_name = DEFAULT_IMAGE_NAME
+
   if os.path.isdir(options.src):
     image_directory = options.src
     image_file = os.path.join(options.src, options.image_name)
 
     if not os.path.exists(image_file):
+      if maybe_build_test:
+        # we want a test image but it doesn't exist
+        # try to build it if we can
+        cros_env.Info('Creating test image')
+        test_output = cros_env.cmd.Output(cros_env.CrosUtilsPath('enter_chroot.sh'),

[adlr, 2010/11/17 20:43:25]

a few lines around here (this, 646, 649) are over 80 cols

+                '--', './mod_image_for_test.sh', '--board=%s' % options.board, '-y')
+        if not os.path.exists(image_file):
+          print test_output
+          cros_env.Fatal('Failed to create test image - please run ./mod_image_for_test.sh manually inside the chroot')
       parser.error('Image file %s does not exist' % image_file)
   else:
     image_file = options.src
     image_directory = os.path.dirname(options.src)
 
+  update_file = os.path.join(image_directory, UPDATE_FILENAME)
+  stateful_file = os.path.join(image_directory, STATEFUL_FILENAME)
+  cros_env.Debug("Image file %s" % image_file)
+  cros_env.Debug("Update file %s" % update_file)
+  cros_env.Debug("Stateful file %s" % stateful_file)
+
   if options.remote:
+    cros_env.Info('Contacting client %s' % options.remote)
     cros_env.SetRemote(options.remote)
     rel = cros_env.GetRemoteRelease()
     if not rel:
       cros_env.Fatal('Could not retrieve remote lsb-release')
     board = rel.get('CHROMEOS_RELEASE_BOARD', '(None)')
     if board != options.board and not options.force_mismatch:
       cros_env.Error('Board %s does not match expected %s' %
                      (board, options.board))
       cros_env.Error('(Use --force-mismatch option to override this)')
       cros_env.Fatal()
 
   elif not options.server_only:
     parser.error('Either --server-only must be specified or '
                  '--remote=<client> needs to be given')
 
-  update_file = os.path.join(image_directory, UPDATE_FILENAME)
-  stateful_file = os.path.join(image_directory, STATEFUL_FILENAME)
-
   if (not cros_env.GenerateUpdatePayload(image_file, update_file) or
       not cros_env.BuildStateful(image_file, stateful_file)):
     cros_env.Fatal()
 
   cros_env.CreateServer(options.port, update_file, stateful_file)
 
   exit_status = 1
   if options.server_only:
     child = None
   else:
@@ skipping 31 matching lines @@
 
   if child:
     os.kill(child, 15)
 
   cros_env.Info('Server exiting with status %d' % exit_status)
   sys.exit(exit_status)
 
 
 if __name__ == '__main__':
   main(sys.argv)