1 files changed, 301 insertions, 0 deletions
diff --git a/scripts/update_device.py b/scripts/update_device.py
new file mode 100755
index 00000000..75c58a7b
--- /dev/null
+++ b/scripts/update_device.py
@@ -0,0 +1,301 @@
+#!/usr/bin/env python
+#
+# Copyright (C) 2017 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Send an A/B update to an Android device over adb."""
+
+import argparse
+import BaseHTTPServer
+import logging
+import os
+import socket
+import subprocess
+import sys
+import threading
+import zipfile
+
+
+# The path used to store the OTA package when applying the package from a file.
+OTA_PACKAGE_PATH = '/data/ota_package'
+
+
+def CopyFileObjLength(fsrc, fdst, buffer_size=128 * 1024, copy_length=None):
+  """Copy from a file object to another.
+
+  This function is similar to shutil.copyfileobj except that it allows to copy
+  less than the full source file.
+
+  Args:
+    fsrc: source file object where to read from.
+    fdst: destination file object where to write to.
+    buffer_size: size of the copy buffer in memory.
+    copy_length: maximum number of bytes to copy, or None to copy everything.
+
+  Returns:
+    the number of bytes copied.
+  """
+  copied = 0
+  while True:
+    chunk_size = buffer_size
+    if copy_length is not None:
+      chunk_size = min(chunk_size, copy_length - copied)
+      if not chunk_size:
+        break
+    buf = fsrc.read(chunk_size)
+    if not buf:
+      break
+    fdst.write(buf)
+    copied += len(buf)
+  return copied
+
+
+class AndroidOTAPackage(object):
+  """Android update payload using the .zip format.
+
+  Android OTA packages traditionally used a .zip file to store the payload. When
+  applying A/B updates over the network, a payload binary is stored RAW inside
+  this .zip file which is used by update_engine to apply the payload. To do
+  this, an offset and size inside the .zip file are provided.
+  """
+
+  # Android OTA package file paths.
+  OTA_PAYLOAD_BIN = 'payload.bin'
+  OTA_PAYLOAD_PROPERTIES_TXT = 'payload_properties.txt'
+
+  def __init__(self, otafilename):
+    self.otafilename = otafilename
+
+    otazip = zipfile.ZipFile(otafilename, 'r')
+    payload_info = otazip.getinfo(self.OTA_PAYLOAD_BIN)
+    self.offset = payload_info.header_offset + len(payload_info.FileHeader())
+    self.size = payload_info.file_size
+    self.properties = otazip.read(self.OTA_PAYLOAD_PROPERTIES_TXT)
+
+
+class UpdateHandler(BaseHTTPServer.BaseHTTPRequestHandler):
+  """A HTTPServer that supports single-range requests.
+
+  Attributes:
+    serving_payload: path to the only payload file we are serving.
+  """
+
+  @staticmethod
+  def _ParseRange(range_str, file_size):
+    """Parse an HTTP range string.
+
+    Args:
+      range_str: HTTP Range header in the request, not including "Header:".
+      file_size: total size of the serving file.
+
+    Returns:
+      A tuple (start_range, end_range) with the range of bytes requested.
+    """
+    start_range = 0
+    end_range = file_size
+
+    if range_str:
+      range_str = range_str.split('=', 1)[1]
+      s, e = range_str.split('-', 1)
+      if s:
+        start_range = int(s)
+        if e:
+          end_range = int(e) + 1
+      elif e:
+        if int(e) < file_size:
+          start_range = file_size - int(e)
+    return start_range, end_range
+
+
+  def do_GET(self):  # pylint: disable=invalid-name
+    """Reply with the requested payload file."""
+    if self.path != '/payload':
+      self.send_error(404, 'Unknown request')
+      return
+
+    if not self.serving_payload:
+      self.send_error(500, 'No serving payload set')
+      return
+
+    try:
+      f = open(self.serving_payload, 'rb')
+    except IOError:
+      self.send_error(404, 'File not found')
+      return
+    # Handle the range request.
+    if 'Range' in self.headers:
+      self.send_response(206)
+    else:
+      self.send_response(200)
+
+    stat = os.fstat(f.fileno())
+    start_range, end_range = self._ParseRange(self.headers.get('range'),
+                                              stat.st_size)
+    logging.info('Serving request for %s from %s [%d, %d) length: %d',
+             self.path, self.serving_payload, start_range, end_range,
+             end_range - start_range)
+
+    self.send_header('Accept-Ranges', 'bytes')
+    self.send_header('Content-Range',
+                     'bytes ' + str(start_range) + '-' + str(end_range - 1) +
+                     '/' + str(end_range - start_range))
+    self.send_header('Content-Length', end_range - start_range)
+
+    self.send_header('Last-Modified', self.date_time_string(stat.st_mtime))
+    self.send_header('Content-type', 'application/octet-stream')
+    self.end_headers()
+
+    f.seek(start_range)
+    CopyFileObjLength(f, self.wfile, copy_length=end_range - start_range)
+
+
+class ServerThread(threading.Thread):
+  """A thread for serving HTTP requests."""
+
+  def __init__(self, ota_filename):
+    threading.Thread.__init__(self)
+    # serving_payload is a class attribute and the UpdateHandler class is
+    # instantiated with every request.
+    UpdateHandler.serving_payload = ota_filename
+    self._httpd = BaseHTTPServer.HTTPServer(('127.0.0.1', 0), UpdateHandler)
+    self.port = self._httpd.server_port
+
+  def run(self):
+    try:
+      self._httpd.serve_forever()
+    except (KeyboardInterrupt, socket.error):
+      pass
+    logging.info('Server Terminated')
+
+  def StopServer(self):
+    self._httpd.socket.close()
+
+
+def StartServer(ota_filename):
+  t = ServerThread(ota_filename)
+  t.start()
+  return t
+
+
+def AndroidUpdateCommand(ota_filename, payload_url):
+  """Return the command to run to start the update in the Android device."""
+  ota = AndroidOTAPackage(ota_filename)
+  headers = ota.properties
+  headers += 'USER_AGENT=Dalvik (something, something)\n'
+
+  # headers += 'POWERWASH=1\n'
+  headers += 'NETWORK_ID=0\n'
+
+  return ['update_engine_client', '--update', '--follow',
+          '--payload=%s' % payload_url, '--offset=%d' % ota.offset,
+          '--size=%d' % ota.size, '--headers="%s"' % headers]
+
+
+class AdbHost(object):
+  """Represents a device connected via ADB."""
+
+  def __init__(self, device_serial=None):
+    """Construct an instance.
+
+    Args:
+        device_serial: options string serial number of attached device.
+    """
+    self._device_serial = device_serial
+    self._command_prefix = ['adb']
+    if self._device_serial:
+      self._command_prefix += ['-s', self._device_serial]
+
+  def adb(self, command):
+    """Run an ADB command like "adb push".
+
+    Args:
+      command: list of strings containing command and arguments to run
+
+    Returns:
+      the program's return code.
+
+    Raises:
+      subprocess.CalledProcessError on command exit != 0.
+    """
+    command = self._command_prefix + command
+    logging.info('Running: %s', ' '.join(str(x) for x in command))
+    p = subprocess.Popen(command, universal_newlines=True)
+    p.wait()
+    return p.returncode
+
+
+def main():
+  parser = argparse.ArgumentParser(description='Android A/B OTA helper.')
+  parser.add_argument('otafile', metavar='ZIP', type=str,
+                      help='the OTA package file (a .zip file).')
+  parser.add_argument('--file', action='store_true',
+                      help='Push the file to the device before updating.')
+  parser.add_argument('--no-push', action='store_true',
+                      help='Skip the "push" command when using --file')
+  parser.add_argument('-s', type=str, default='', metavar='DEVICE',
+                      help='The specific device to use.')
+  parser.add_argument('--no-verbose', action='store_true',
+                      help='Less verbose output')
+  args = parser.parse_args()
+  logging.basicConfig(
+      level=logging.WARNING if args.no_verbose else logging.INFO)
+
+  dut = AdbHost(args.s)
+
+  server_thread = None
+  # List of commands to execute on exit.
+  finalize_cmds = []
+  # Commands to execute when canceling an update.
+  cancel_cmd = ['shell', 'su', '0', 'update_engine_client', '--cancel']
+  # List of commands to perform the update.
+  cmds = []
+
+  if args.file:
+    # Update via pushing a file to /data.
+    device_ota_file = os.path.join(OTA_PACKAGE_PATH, 'debug.zip')
+    payload_url = 'file://' + device_ota_file
+    if not args.no_push:
+      cmds.append(['push', args.otafile, device_ota_file])
+    cmds.append(['shell', 'su', '0', 'chown', 'system:cache', device_ota_file])
+    cmds.append(['shell', 'su', '0', 'chmod', '0660', device_ota_file])
+  else:
+    # Update via sending the payload over the network with an "adb reverse"
+    # command.
+    device_port = 1234
+    payload_url = 'http://127.0.0.1:%d/payload' % device_port
+    server_thread = StartServer(args.otafile)
+    cmds.append(
+        ['reverse', 'tcp:%d' % device_port, 'tcp:%d' % server_thread.port])
+    finalize_cmds.append(['reverse', '--remove', 'tcp:%d' % device_port])
+
+  try:
+    # The main update command using the configured payload_url.
+    update_cmd = AndroidUpdateCommand(args.otafile, payload_url)
+    cmds.append(['shell', 'su', '0'] + update_cmd)
+
+    for cmd in cmds:
+      dut.adb(cmd)
+  except KeyboardInterrupt:
+    dut.adb(cancel_cmd)
+  finally:
+    if server_thread:
+      server_thread.StopServer()
+    for cmd in finalize_cmds:
+      dut.adb(cmd)
+
+  return 0
+
+if __name__ == '__main__':
+  sys.exit(main())