• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python3
2#
3# Copyright (C) 2017 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18"""Send an A/B update to an Android device over adb."""
19
20from __future__ import print_function
21from __future__ import absolute_import
22
23import argparse
24import binascii
25import logging
26import os
27import re
28import socket
29import subprocess
30import sys
31import struct
32import tempfile
33import time
34import threading
35import zipfile
36import shutil
37
38from six.moves import BaseHTTPServer
39
40
41# The path used to store the OTA package when applying the package from a file.
42OTA_PACKAGE_PATH = '/data/ota_package'
43
44# The path to the payload public key on the device.
45PAYLOAD_KEY_PATH = '/etc/update_engine/update-payload-key.pub.pem'
46
47# The port on the device that update_engine should connect to.
48DEVICE_PORT = 1234
49
50
51def CopyFileObjLength(fsrc, fdst, buffer_size=128 * 1024, copy_length=None, speed_limit=None):
52  """Copy from a file object to another.
53
54  This function is similar to shutil.copyfileobj except that it allows to copy
55  less than the full source file.
56
57  Args:
58    fsrc: source file object where to read from.
59    fdst: destination file object where to write to.
60    buffer_size: size of the copy buffer in memory.
61    copy_length: maximum number of bytes to copy, or None to copy everything.
62    speed_limit: upper limit for copying speed, in bytes per second.
63
64  Returns:
65    the number of bytes copied.
66  """
67  # If buffer size significantly bigger than speed limit
68  # traffic would seem extremely spiky to the client.
69  if speed_limit:
70    print(f"Applying speed limit: {speed_limit}")
71    buffer_size = min(speed_limit//32, buffer_size)
72
73  start_time = time.time()
74  copied = 0
75  while True:
76    chunk_size = buffer_size
77    if copy_length is not None:
78      chunk_size = min(chunk_size, copy_length - copied)
79      if not chunk_size:
80        break
81    buf = fsrc.read(chunk_size)
82    if not buf:
83      break
84    if speed_limit:
85      expected_duration = copied/speed_limit
86      actual_duration = time.time() - start_time
87      if actual_duration < expected_duration:
88        time.sleep(expected_duration-actual_duration)
89    fdst.write(buf)
90    copied += len(buf)
91  return copied
92
93
94class AndroidOTAPackage(object):
95  """Android update payload using the .zip format.
96
97  Android OTA packages traditionally used a .zip file to store the payload. When
98  applying A/B updates over the network, a payload binary is stored RAW inside
99  this .zip file which is used by update_engine to apply the payload. To do
100  this, an offset and size inside the .zip file are provided.
101  """
102
103  # Android OTA package file paths.
104  OTA_PAYLOAD_BIN = 'payload.bin'
105  OTA_PAYLOAD_PROPERTIES_TXT = 'payload_properties.txt'
106  SECONDARY_OTA_PAYLOAD_BIN = 'secondary/payload.bin'
107  SECONDARY_OTA_PAYLOAD_PROPERTIES_TXT = 'secondary/payload_properties.txt'
108  PAYLOAD_MAGIC_HEADER = b'CrAU'
109
110  def __init__(self, otafilename, secondary_payload=False):
111    self.otafilename = otafilename
112
113    otazip = zipfile.ZipFile(otafilename, 'r')
114    payload_entry = (self.SECONDARY_OTA_PAYLOAD_BIN if secondary_payload else
115                     self.OTA_PAYLOAD_BIN)
116    payload_info = otazip.getinfo(payload_entry)
117
118    if payload_info.compress_type != 0:
119      logging.error(
120          "Expected payload to be uncompressed, got compression method %d",
121          payload_info.compress_type)
122    # Don't use len(payload_info.extra). Because that returns size of extra
123    # fields in central directory. We need to look at local file directory,
124    # as these two might have different sizes.
125    with open(otafilename, "rb") as fp:
126      fp.seek(payload_info.header_offset)
127      data = fp.read(zipfile.sizeFileHeader)
128      fheader = struct.unpack(zipfile.structFileHeader, data)
129      # Last two fields of local file header are filename length and
130      # extra length
131      filename_len = fheader[-2]
132      extra_len = fheader[-1]
133      self.offset = payload_info.header_offset
134      self.offset += zipfile.sizeFileHeader
135      self.offset += filename_len + extra_len
136      self.size = payload_info.file_size
137      fp.seek(self.offset)
138      payload_header = fp.read(4)
139      if payload_header != self.PAYLOAD_MAGIC_HEADER:
140        logging.warning(
141            "Invalid header, expected %s, got %s."
142            "Either the offset is not correct, or payload is corrupted",
143            binascii.hexlify(self.PAYLOAD_MAGIC_HEADER),
144            binascii.hexlify(payload_header))
145
146    property_entry = (self.SECONDARY_OTA_PAYLOAD_PROPERTIES_TXT if
147                      secondary_payload else self.OTA_PAYLOAD_PROPERTIES_TXT)
148    self.properties = otazip.read(property_entry)
149
150
151class UpdateHandler(BaseHTTPServer.BaseHTTPRequestHandler):
152  """A HTTPServer that supports single-range requests.
153
154  Attributes:
155    serving_payload: path to the only payload file we are serving.
156    serving_range: the start offset and size tuple of the payload.
157  """
158
159  @staticmethod
160  def _parse_range(range_str, file_size):
161    """Parse an HTTP range string.
162
163    Args:
164      range_str: HTTP Range header in the request, not including "Header:".
165      file_size: total size of the serving file.
166
167    Returns:
168      A tuple (start_range, end_range) with the range of bytes requested.
169    """
170    start_range = 0
171    end_range = file_size
172
173    if range_str:
174      range_str = range_str.split('=', 1)[1]
175      s, e = range_str.split('-', 1)
176      if s:
177        start_range = int(s)
178        if e:
179          end_range = int(e) + 1
180      elif e:
181        if int(e) < file_size:
182          start_range = file_size - int(e)
183    return start_range, end_range
184
185  def do_GET(self):  # pylint: disable=invalid-name
186    """Reply with the requested payload file."""
187    if self.path != '/payload':
188      self.send_error(404, 'Unknown request')
189      return
190
191    if not self.serving_payload:
192      self.send_error(500, 'No serving payload set')
193      return
194
195    try:
196      f = open(self.serving_payload, 'rb')
197    except IOError:
198      self.send_error(404, 'File not found')
199      return
200    # Handle the range request.
201    if 'Range' in self.headers:
202      self.send_response(206)
203    else:
204      self.send_response(200)
205
206    serving_start, serving_size = self.serving_range
207    start_range, end_range = self._parse_range(self.headers.get('range'),
208                                               serving_size)
209    logging.info('Serving request for %s from %s [%d, %d) length: %d',
210                 self.path, self.serving_payload, serving_start + start_range,
211                 serving_start + end_range, end_range - start_range)
212
213    self.send_header('Accept-Ranges', 'bytes')
214    self.send_header('Content-Range',
215                     'bytes ' + str(start_range) + '-' + str(end_range - 1) +
216                     '/' + str(end_range - start_range))
217    self.send_header('Content-Length', end_range - start_range)
218
219    stat = os.fstat(f.fileno())
220    self.send_header('Last-Modified', self.date_time_string(stat.st_mtime))
221    self.send_header('Content-type', 'application/octet-stream')
222    self.end_headers()
223
224    f.seek(serving_start + start_range)
225    CopyFileObjLength(f, self.wfile, copy_length=end_range -
226                      start_range, speed_limit=self.speed_limit)
227
228
229class ServerThread(threading.Thread):
230  """A thread for serving HTTP requests."""
231
232  def __init__(self, ota_filename, serving_range, speed_limit):
233    threading.Thread.__init__(self)
234    # serving_payload and serving_range are class attributes and the
235    # UpdateHandler class is instantiated with every request.
236    UpdateHandler.serving_payload = ota_filename
237    UpdateHandler.serving_range = serving_range
238    UpdateHandler.speed_limit = speed_limit
239    self._httpd = BaseHTTPServer.HTTPServer(('127.0.0.1', 0), UpdateHandler)
240    self.port = self._httpd.server_port
241
242  def run(self):
243    try:
244      self._httpd.serve_forever()
245    except (KeyboardInterrupt, socket.error):
246      pass
247    logging.info('Server Terminated')
248
249  def StopServer(self):
250    self._httpd.shutdown()
251    self._httpd.socket.close()
252
253
254def StartServer(ota_filename, serving_range, speed_limit):
255  t = ServerThread(ota_filename, serving_range, speed_limit)
256  t.start()
257  return t
258
259
260def AndroidUpdateCommand(ota_filename, secondary, payload_url, extra_headers):
261  """Return the command to run to start the update in the Android device."""
262  ota = AndroidOTAPackage(ota_filename, secondary)
263  headers = ota.properties
264  headers += b'USER_AGENT=Dalvik (something, something)\n'
265  headers += b'NETWORK_ID=0\n'
266  headers += extra_headers.encode()
267
268  return ['update_engine_client', '--update', '--follow',
269          '--payload=%s' % payload_url, '--offset=%d' % ota.offset,
270          '--size=%d' % ota.size, '--headers="%s"' % headers.decode()]
271
272
273class AdbHost(object):
274  """Represents a device connected via ADB."""
275
276  def __init__(self, device_serial=None):
277    """Construct an instance.
278
279    Args:
280        device_serial: options string serial number of attached device.
281    """
282    self._device_serial = device_serial
283    self._command_prefix = ['adb']
284    if self._device_serial:
285      self._command_prefix += ['-s', self._device_serial]
286
287  def adb(self, command, timeout_seconds: float = None):
288    """Run an ADB command like "adb push".
289
290    Args:
291      command: list of strings containing command and arguments to run
292
293    Returns:
294      the program's return code.
295
296    Raises:
297      subprocess.CalledProcessError on command exit != 0.
298    """
299    command = self._command_prefix + command
300    logging.info('Running: %s', ' '.join(str(x) for x in command))
301    p = subprocess.Popen(command, universal_newlines=True)
302    p.wait(timeout_seconds)
303    return p.returncode
304
305  def adb_output(self, command):
306    """Run an ADB command like "adb push" and return the output.
307
308    Args:
309      command: list of strings containing command and arguments to run
310
311    Returns:
312      the program's output as a string.
313
314    Raises:
315      subprocess.CalledProcessError on command exit != 0.
316    """
317    command = self._command_prefix + command
318    logging.info('Running: %s', ' '.join(str(x) for x in command))
319    return subprocess.check_output(command, universal_newlines=True)
320
321
322def PushMetadata(dut, otafile, metadata_path):
323  header_format = ">4sQQL"
324  with tempfile.TemporaryDirectory() as tmpdir:
325    with zipfile.ZipFile(otafile, "r") as zfp:
326      extracted_path = os.path.join(tmpdir, "payload.bin")
327      with zfp.open("payload.bin") as payload_fp, \
328              open(extracted_path, "wb") as output_fp:
329        # Only extract the first |data_offset| bytes from the payload.
330        # This is because allocateSpaceForPayload only needs to see
331        # the manifest, not the entire payload.
332        # Extracting the entire payload works, but is slow for full
333        # OTA.
334        header = payload_fp.read(struct.calcsize(header_format))
335        magic, major_version, manifest_size, metadata_signature_size = struct.unpack(header_format, header)
336        assert magic == b"CrAU", "Invalid magic {}, expected CrAU".format(magic)
337        assert major_version == 2, "Invalid major version {}, only version 2 is supported".format(major_version)
338        output_fp.write(header)
339        output_fp.write(payload_fp.read(manifest_size + metadata_signature_size))
340
341      return dut.adb([
342          "push",
343          extracted_path,
344          metadata_path
345      ]) == 0
346
347
348def ParseSpeedLimit(arg: str) -> int:
349  arg = arg.strip().upper()
350  if not re.match(r"\d+[KkMmGgTt]?", arg):
351    raise argparse.ArgumentError(
352        "Wrong speed limit format, expected format is number followed by unit, such as 10K, 5m, 3G (case insensitive)")
353  unit = 1
354  if arg[-1].isalpha():
355    if arg[-1] == "K":
356      unit = 1024
357    elif arg[-1] == "M":
358      unit = 1024 * 1024
359    elif arg[-1] == "G":
360      unit = 1024 * 1024 * 1024
361    elif arg[-1] == "T":
362      unit = 1024 * 1024 * 1024 * 1024
363    else:
364      raise argparse.ArgumentError(
365          f"Unsupported unit for download speed: {arg[-1]}, supported units are K,M,G,T (case insensitive)")
366  return int(float(arg[:-1]) * unit)
367
368
369def main():
370  parser = argparse.ArgumentParser(description='Android A/B OTA helper.')
371  parser.add_argument('otafile', metavar='PAYLOAD', type=str,
372                      help='the OTA package file (a .zip file) or raw payload \
373                      if device uses Omaha.')
374  parser.add_argument('--file', action='store_true',
375                      help='Push the file to the device before updating.')
376  parser.add_argument('--no-push', action='store_true',
377                      help='Skip the "push" command when using --file')
378  parser.add_argument('-s', type=str, default='', metavar='DEVICE',
379                      help='The specific device to use.')
380  parser.add_argument('--no-verbose', action='store_true',
381                      help='Less verbose output')
382  parser.add_argument('--public-key', type=str, default='',
383                      help='Override the public key used to verify payload.')
384  parser.add_argument('--extra-headers', type=str, default='',
385                      help='Extra headers to pass to the device.')
386  parser.add_argument('--secondary', action='store_true',
387                      help='Update with the secondary payload in the package.')
388  parser.add_argument('--no-slot-switch', action='store_true',
389                      help='Do not perform slot switch after the update.')
390  parser.add_argument('--no-postinstall', action='store_true',
391                      help='Do not execute postinstall scripts after the update.')
392  parser.add_argument('--allocate-only', action='store_true',
393                      help='Allocate space for this OTA, instead of actually \
394                        applying the OTA.')
395  parser.add_argument('--verify-only', action='store_true',
396                      help='Verify metadata then exit, instead of applying the OTA.')
397  parser.add_argument('--no-care-map', action='store_true',
398                      help='Do not push care_map.pb to device.')
399  parser.add_argument('--perform-slot-switch', action='store_true',
400                      help='Perform slot switch for this OTA package')
401  parser.add_argument('--perform-reset-slot-switch', action='store_true',
402                      help='Perform reset slot switch for this OTA package')
403  parser.add_argument('--wipe-user-data', action='store_true',
404                      help='Wipe userdata after installing OTA')
405  parser.add_argument('--vabc-none', action='store_true',
406                      help='Set Virtual AB Compression algorithm to none, but still use Android COW format')
407  parser.add_argument('--disable-vabc', action='store_true',
408                      help='Option to enable or disable vabc. If set to false, will fall back on A/B')
409  parser.add_argument('--enable-threading', action='store_true',
410                      help='Enable multi-threaded compression for VABC')
411  parser.add_argument('--disable-threading', action='store_true',
412                      help='Disable multi-threaded compression for VABC')
413  parser.add_argument('--batched-writes', action='store_true',
414                      help='Enable batched writes for VABC')
415  parser.add_argument('--speed-limit', type=str,
416                      help='Speed limit for serving payloads over HTTP. For '
417                      'example: 10K, 5m, 1G, input is case insensitive')
418
419  args = parser.parse_args()
420  if args.speed_limit:
421    args.speed_limit = ParseSpeedLimit(args.speed_limit)
422
423  logging.basicConfig(
424      level=logging.WARNING if args.no_verbose else logging.INFO)
425
426  start_time = time.perf_counter()
427
428  dut = AdbHost(args.s)
429
430  server_thread = None
431  # List of commands to execute on exit.
432  finalize_cmds = []
433  # Commands to execute when canceling an update.
434  cancel_cmd = ['shell', 'su', '0', 'update_engine_client', '--cancel']
435  # List of commands to perform the update.
436  cmds = []
437
438  help_cmd = ['shell', 'su', '0', 'update_engine_client', '--help']
439
440  metadata_path = "/data/ota_package/metadata"
441  if args.allocate_only:
442    with zipfile.ZipFile(args.otafile, "r") as zfp:
443      headers = zfp.read("payload_properties.txt").decode()
444    if PushMetadata(dut, args.otafile, metadata_path):
445      dut.adb([
446          "shell", "update_engine_client", "--allocate",
447          "--metadata={} --headers='{}'".format(metadata_path, headers)])
448    # Return 0, as we are executing ADB commands here, no work needed after
449    # this point
450    return 0
451  if args.verify_only:
452    if PushMetadata(dut, args.otafile, metadata_path):
453      dut.adb([
454          "shell", "update_engine_client", "--verify",
455          "--metadata={}".format(metadata_path)])
456    # Return 0, as we are executing ADB commands here, no work needed after
457    # this point
458    return 0
459  if args.perform_slot_switch:
460    assert PushMetadata(dut, args.otafile, metadata_path)
461    dut.adb(["shell", "update_engine_client",
462            "--switch_slot=true", "--metadata={}".format(metadata_path), "--follow"])
463    return 0
464  if args.perform_reset_slot_switch:
465    assert PushMetadata(dut, args.otafile, metadata_path)
466    dut.adb(["shell", "update_engine_client",
467            "--switch_slot=false", "--metadata={}".format(metadata_path)])
468    return 0
469
470  if args.no_slot_switch:
471    args.extra_headers += "\nSWITCH_SLOT_ON_REBOOT=0"
472  if args.no_postinstall:
473    args.extra_headers += "\nRUN_POST_INSTALL=0"
474  if args.wipe_user_data:
475    args.extra_headers += "\nPOWERWASH=1"
476  if args.vabc_none:
477    args.extra_headers += "\nVABC_NONE=1"
478  if args.disable_vabc:
479    args.extra_headers += "\nDISABLE_VABC=1"
480  if args.enable_threading:
481    args.extra_headers += "\nENABLE_THREADING=1"
482  elif args.disable_threading:
483    args.extra_headers += "\nENABLE_THREADING=0"
484  if args.batched_writes:
485    args.extra_headers += "\nBATCHED_WRITES=1"
486
487  with zipfile.ZipFile(args.otafile) as zfp:
488    CARE_MAP_ENTRY_NAME = "care_map.pb"
489    if CARE_MAP_ENTRY_NAME in zfp.namelist() and not args.no_care_map:
490      # Need root permission to push to /data
491      dut.adb(["root"])
492      with tempfile.NamedTemporaryFile() as care_map_fp:
493        care_map_fp.write(zfp.read(CARE_MAP_ENTRY_NAME))
494        care_map_fp.flush()
495        dut.adb(["push", care_map_fp.name,
496                "/data/ota_package/" + CARE_MAP_ENTRY_NAME])
497
498  if args.file:
499    # Update via pushing a file to /data.
500    device_ota_file = os.path.join(OTA_PACKAGE_PATH, 'debug.zip')
501    payload_url = 'file://' + device_ota_file
502    if not args.no_push:
503      data_local_tmp_file = '/data/local/tmp/debug.zip'
504      cmds.append(['push', args.otafile, data_local_tmp_file])
505      cmds.append(['shell', 'su', '0', 'mv', data_local_tmp_file,
506                   device_ota_file])
507      cmds.append(['shell', 'su', '0', 'chcon',
508                   'u:object_r:ota_package_file:s0', device_ota_file])
509    cmds.append(['shell', 'su', '0', 'chown', 'system:cache', device_ota_file])
510    cmds.append(['shell', 'su', '0', 'chmod', '0660', device_ota_file])
511  else:
512    # Update via sending the payload over the network with an "adb reverse"
513    # command.
514    payload_url = 'http://127.0.0.1:%d/payload' % DEVICE_PORT
515    serving_range = (0, os.stat(args.otafile).st_size)
516    server_thread = StartServer(args.otafile, serving_range, args.speed_limit)
517    cmds.append(
518        ['reverse', 'tcp:%d' % DEVICE_PORT, 'tcp:%d' % server_thread.port])
519    finalize_cmds.append(['reverse', '--remove', 'tcp:%d' % DEVICE_PORT])
520
521  if args.public_key:
522    payload_key_dir = os.path.dirname(PAYLOAD_KEY_PATH)
523    cmds.append(
524        ['shell', 'su', '0', 'mount', '-t', 'tmpfs', 'tmpfs', payload_key_dir])
525    # Allow adb push to payload_key_dir
526    cmds.append(['shell', 'su', '0', 'chcon', 'u:object_r:shell_data_file:s0',
527                 payload_key_dir])
528    cmds.append(['push', args.public_key, PAYLOAD_KEY_PATH])
529    # Allow update_engine to read it.
530    cmds.append(['shell', 'su', '0', 'chcon', '-R', 'u:object_r:system_file:s0',
531                 payload_key_dir])
532    finalize_cmds.append(['shell', 'su', '0', 'umount', payload_key_dir])
533
534  try:
535    # The main update command using the configured payload_url.
536    update_cmd = AndroidUpdateCommand(args.otafile, args.secondary,
537                                        payload_url, args.extra_headers)
538    cmds.append(['shell', 'su', '0'] + update_cmd)
539
540    for cmd in cmds:
541      dut.adb(cmd)
542  except KeyboardInterrupt:
543    dut.adb(cancel_cmd)
544  finally:
545    if server_thread:
546      server_thread.StopServer()
547    for cmd in finalize_cmds:
548      dut.adb(cmd, 5)
549
550  logging.info('Update took %.3f seconds', (time.perf_counter() - start_time))
551  return 0
552
553
554if __name__ == '__main__':
555  sys.exit(main())
556