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.
 # 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
@@ skipping 12 matching lines @@
 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'

[Paul Stewart, 2010/11/23 21:21:33]

The old stateful.gz is dead, I believe, and removed in tip of trunk.

+STATEFUL_FILENAME_TGZ = 'stateful.tgz'
 
 # 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', '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 = True
+    try:
+      self.cmd.Output('pv', '--help')
+    except OSError:
+      self.have_pv = False
+
   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)
 
   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 23 matching lines @@
         if var:
           ret[var] = val.strip('\t ').rstrip('\t ')
     return ret
 
   def GetRemoteRelease(self):
     lsb_release = self.ssh_cmd.Output('cat', '/etc/lsb-release')
     if not lsb_release:
       return None
     return self.ParseShVars(lsb_release)
 
-  def CreateServer(self, port, update_file, stateful_file):
+  def CreateServer(self, port, update_file, stateful_file, stateful_file_tgz):
     """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))
+                FileUpdateResponse(update_file,
+                        verbose=self.verbose, have_pv=self.have_pv))
     UpdateHandler.SetupUrl('/%s' % STATEFUL_FILENAME,
-                           FileUpdateResponse(stateful_file,
-                                              verbose=self.verbose))
+                FileUpdateResponse(stateful_file,
+                        verbose=self.verbose, have_pv=self.have_pv))
+    UpdateHandler.SetupUrl('/%s' % STATEFUL_FILENAME_TGZ,
+                FileUpdateResponse(stateful_file_tgz,
+                        verbose=self.verbose, have_pv=self.have_pv))

[Paul Stewart, 2010/11/23 21:21:33]

Please rebase these changes on what's currently in tip of trunk.  I think all
you'll need to do is add your have_pv flag here.  The old stateful file is gone.
 I'd suggest you follow style and newline-indent after commas instead of
de-indenting the line above to get it to fit in 80 cols.

 
     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...")
     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)
+
+    # this will make the client read the files we have set up
     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!')
     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 filesize <= self.blocksize:
-      handler.wfile.write(f.read())
-    else:
+    if send_content:
       sent_size = 0
       sent_percentage = None
+
+      #FIXME: this should use pv also

[Paul Stewart, 2010/11/23 21:21:33]

Use # TODO(sgj): ...

       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
@@ skipping 155 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
+  build_test_image = False
+
+  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
+      build_test_image = 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 build_test_image:
+        # 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'),
+                '--', './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)
+  stateful_file_tgz = cros_env.CrosUtilsPath(
+              '../platform/dev/static/stateful.tgz')
+  cros_env.Debug("Image file %s" % image_file)
+  cros_env.Debug("Update file %s" % update_file)
+  cros_env.Debug("Stateful file %s" % stateful_file)
+  cros_env.Debug("Stateful file tgz %s" % stateful_file_tgz)
+
   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)
+  cros_env.CreateServer(options.port, update_file, stateful_file,
+        stateful_file_tgz)
 
   exit_status = 1
   if options.server_only:
     child = None
   else:
     # Start an "image-to-live" instance that will pull bits from the server
     child = os.fork()
     if child:
       signal.signal(signal.SIGCHLD, ChildFinished(child).SigHandler)
     else:
@@ skipping 26 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)