#!/usr/bin/env python3 # # 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.""" from __future__ import print_function from __future__ import absolute_import import argparse import binascii import logging import os import re import socket import subprocess import sys import struct import tempfile import time import threading import zipfile import shutil from six.moves import BaseHTTPServer # The path used to store the OTA package when applying the package from a file. OTA_PACKAGE_PATH = '/data/ota_package' # The path to the payload public key on the device. PAYLOAD_KEY_PATH = '/etc/update_engine/update-payload-key.pub.pem' # The port on the device that update_engine should connect to. DEVICE_PORT = 1234 def CopyFileObjLength(fsrc, fdst, buffer_size=128 * 1024, copy_length=None, speed_limit=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. speed_limit: upper limit for copying speed, in bytes per second. Returns: the number of bytes copied. """ # If buffer size significantly bigger than speed limit # traffic would seem extremely spiky to the client. if speed_limit: print(f"Applying speed limit: {speed_limit}") buffer_size = min(speed_limit//32, buffer_size) start_time = time.time() 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 if speed_limit: expected_duration = copied/speed_limit actual_duration = time.time() - start_time if actual_duration < expected_duration: time.sleep(expected_duration-actual_duration) 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' SECONDARY_OTA_PAYLOAD_BIN = 'secondary/payload.bin' SECONDARY_OTA_PAYLOAD_PROPERTIES_TXT = 'secondary/payload_properties.txt' PAYLOAD_MAGIC_HEADER = b'CrAU' def __init__(self, otafilename, secondary_payload=False): self.otafilename = otafilename otazip = zipfile.ZipFile(otafilename, 'r') payload_entry = (self.SECONDARY_OTA_PAYLOAD_BIN if secondary_payload else self.OTA_PAYLOAD_BIN) payload_info = otazip.getinfo(payload_entry) if payload_info.compress_type != 0: logging.error( "Expected payload to be uncompressed, got compression method %d", payload_info.compress_type) # Don't use len(payload_info.extra). Because that returns size of extra # fields in central directory. We need to look at local file directory, # as these two might have different sizes. with open(otafilename, "rb") as fp: fp.seek(payload_info.header_offset) data = fp.read(zipfile.sizeFileHeader) fheader = struct.unpack(zipfile.structFileHeader, data) # Last two fields of local file header are filename length and # extra length filename_len = fheader[-2] extra_len = fheader[-1] self.offset = payload_info.header_offset self.offset += zipfile.sizeFileHeader self.offset += filename_len + extra_len self.size = payload_info.file_size fp.seek(self.offset) payload_header = fp.read(4) if payload_header != self.PAYLOAD_MAGIC_HEADER: logging.warning( "Invalid header, expected %s, got %s." "Either the offset is not correct, or payload is corrupted", binascii.hexlify(self.PAYLOAD_MAGIC_HEADER), binascii.hexlify(payload_header)) property_entry = (self.SECONDARY_OTA_PAYLOAD_PROPERTIES_TXT if secondary_payload else self.OTA_PAYLOAD_PROPERTIES_TXT) self.properties = otazip.read(property_entry) class UpdateHandler(BaseHTTPServer.BaseHTTPRequestHandler): """A HTTPServer that supports single-range requests. Attributes: serving_payload: path to the only payload file we are serving. serving_range: the start offset and size tuple of the payload. """ @staticmethod def _parse_range(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) serving_start, serving_size = self.serving_range start_range, end_range = self._parse_range(self.headers.get('range'), serving_size) logging.info('Serving request for %s from %s [%d, %d) length: %d', self.path, self.serving_payload, serving_start + start_range, serving_start + 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) stat = os.fstat(f.fileno()) 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(serving_start + start_range) CopyFileObjLength(f, self.wfile, copy_length=end_range - start_range, speed_limit=self.speed_limit) class ServerThread(threading.Thread): """A thread for serving HTTP requests.""" def __init__(self, ota_filename, serving_range, speed_limit): threading.Thread.__init__(self) # serving_payload and serving_range are class attributes and the # UpdateHandler class is instantiated with every request. UpdateHandler.serving_payload = ota_filename UpdateHandler.serving_range = serving_range UpdateHandler.speed_limit = speed_limit 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.shutdown() self._httpd.socket.close() def StartServer(ota_filename, serving_range, speed_limit): t = ServerThread(ota_filename, serving_range, speed_limit) t.start() return t def AndroidUpdateCommand(ota_filename, secondary, payload_url, extra_headers): """Return the command to run to start the update in the Android device.""" ota = AndroidOTAPackage(ota_filename, secondary) headers = ota.properties headers += b'USER_AGENT=Dalvik (something, something)\n' headers += b'NETWORK_ID=0\n' headers += extra_headers.encode() return ['update_engine_client', '--update', '--follow', '--payload=%s' % payload_url, '--offset=%d' % ota.offset, '--size=%d' % ota.size, '--headers="%s"' % headers.decode()] 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, timeout_seconds: float = None): """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(timeout_seconds) return p.returncode def adb_output(self, command): """Run an ADB command like "adb push" and return the output. Args: command: list of strings containing command and arguments to run Returns: the program's output as a string. Raises: subprocess.CalledProcessError on command exit != 0. """ command = self._command_prefix + command logging.info('Running: %s', ' '.join(str(x) for x in command)) return subprocess.check_output(command, universal_newlines=True) def PushMetadata(dut, otafile, metadata_path): header_format = ">4sQQL" with tempfile.TemporaryDirectory() as tmpdir: with zipfile.ZipFile(otafile, "r") as zfp: extracted_path = os.path.join(tmpdir, "payload.bin") with zfp.open("payload.bin") as payload_fp, \ open(extracted_path, "wb") as output_fp: # Only extract the first |data_offset| bytes from the payload. # This is because allocateSpaceForPayload only needs to see # the manifest, not the entire payload. # Extracting the entire payload works, but is slow for full # OTA. header = payload_fp.read(struct.calcsize(header_format)) magic, major_version, manifest_size, metadata_signature_size = struct.unpack(header_format, header) assert magic == b"CrAU", "Invalid magic {}, expected CrAU".format(magic) assert major_version == 2, "Invalid major version {}, only version 2 is supported".format(major_version) output_fp.write(header) output_fp.write(payload_fp.read(manifest_size + metadata_signature_size)) return dut.adb([ "push", extracted_path, metadata_path ]) == 0 def ParseSpeedLimit(arg: str) -> int: arg = arg.strip().upper() if not re.match(r"\d+[KkMmGgTt]?", arg): raise argparse.ArgumentError( "Wrong speed limit format, expected format is number followed by unit, such as 10K, 5m, 3G (case insensitive)") unit = 1 if arg[-1].isalpha(): if arg[-1] == "K": unit = 1024 elif arg[-1] == "M": unit = 1024 * 1024 elif arg[-1] == "G": unit = 1024 * 1024 * 1024 elif arg[-1] == "T": unit = 1024 * 1024 * 1024 * 1024 else: raise argparse.ArgumentError( f"Unsupported unit for download speed: {arg[-1]}, supported units are K,M,G,T (case insensitive)") return int(float(arg[:-1]) * unit) def main(): parser = argparse.ArgumentParser(description='Android A/B OTA helper.') parser.add_argument('otafile', metavar='PAYLOAD', type=str, help='the OTA package file (a .zip file) or raw payload \ if device uses Omaha.') 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') parser.add_argument('--public-key', type=str, default='', help='Override the public key used to verify payload.') parser.add_argument('--extra-headers', type=str, default='', help='Extra headers to pass to the device.') parser.add_argument('--secondary', action='store_true', help='Update with the secondary payload in the package.') parser.add_argument('--no-slot-switch', action='store_true', help='Do not perform slot switch after the update.') parser.add_argument('--no-postinstall', action='store_true', help='Do not execute postinstall scripts after the update.') parser.add_argument('--allocate-only', action='store_true', help='Allocate space for this OTA, instead of actually \ applying the OTA.') parser.add_argument('--verify-only', action='store_true', help='Verify metadata then exit, instead of applying the OTA.') parser.add_argument('--no-care-map', action='store_true', help='Do not push care_map.pb to device.') parser.add_argument('--perform-slot-switch', action='store_true', help='Perform slot switch for this OTA package') parser.add_argument('--perform-reset-slot-switch', action='store_true', help='Perform reset slot switch for this OTA package') parser.add_argument('--wipe-user-data', action='store_true', help='Wipe userdata after installing OTA') parser.add_argument('--vabc-none', action='store_true', help='Set Virtual AB Compression algorithm to none, but still use Android COW format') parser.add_argument('--disable-vabc', action='store_true', help='Option to enable or disable vabc. If set to false, will fall back on A/B') parser.add_argument('--enable-threading', action='store_true', help='Enable multi-threaded compression for VABC') parser.add_argument('--disable-threading', action='store_true', help='Disable multi-threaded compression for VABC') parser.add_argument('--batched-writes', action='store_true', help='Enable batched writes for VABC') parser.add_argument('--speed-limit', type=str, help='Speed limit for serving payloads over HTTP. For ' 'example: 10K, 5m, 1G, input is case insensitive') args = parser.parse_args() if args.speed_limit: args.speed_limit = ParseSpeedLimit(args.speed_limit) logging.basicConfig( level=logging.WARNING if args.no_verbose else logging.INFO) start_time = time.perf_counter() 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 = [] help_cmd = ['shell', 'su', '0', 'update_engine_client', '--help'] metadata_path = "/data/ota_package/metadata" if args.allocate_only: with zipfile.ZipFile(args.otafile, "r") as zfp: headers = zfp.read("payload_properties.txt").decode() if PushMetadata(dut, args.otafile, metadata_path): dut.adb([ "shell", "update_engine_client", "--allocate", "--metadata={} --headers='{}'".format(metadata_path, headers)]) # Return 0, as we are executing ADB commands here, no work needed after # this point return 0 if args.verify_only: if PushMetadata(dut, args.otafile, metadata_path): dut.adb([ "shell", "update_engine_client", "--verify", "--metadata={}".format(metadata_path)]) # Return 0, as we are executing ADB commands here, no work needed after # this point return 0 if args.perform_slot_switch: assert PushMetadata(dut, args.otafile, metadata_path) dut.adb(["shell", "update_engine_client", "--switch_slot=true", "--metadata={}".format(metadata_path), "--follow"]) return 0 if args.perform_reset_slot_switch: assert PushMetadata(dut, args.otafile, metadata_path) dut.adb(["shell", "update_engine_client", "--switch_slot=false", "--metadata={}".format(metadata_path)]) return 0 if args.no_slot_switch: args.extra_headers += "\nSWITCH_SLOT_ON_REBOOT=0" if args.no_postinstall: args.extra_headers += "\nRUN_POST_INSTALL=0" if args.wipe_user_data: args.extra_headers += "\nPOWERWASH=1" if args.vabc_none: args.extra_headers += "\nVABC_NONE=1" if args.disable_vabc: args.extra_headers += "\nDISABLE_VABC=1" if args.enable_threading: args.extra_headers += "\nENABLE_THREADING=1" elif args.disable_threading: args.extra_headers += "\nENABLE_THREADING=0" if args.batched_writes: args.extra_headers += "\nBATCHED_WRITES=1" with zipfile.ZipFile(args.otafile) as zfp: CARE_MAP_ENTRY_NAME = "care_map.pb" if CARE_MAP_ENTRY_NAME in zfp.namelist() and not args.no_care_map: # Need root permission to push to /data dut.adb(["root"]) with tempfile.NamedTemporaryFile() as care_map_fp: care_map_fp.write(zfp.read(CARE_MAP_ENTRY_NAME)) care_map_fp.flush() dut.adb(["push", care_map_fp.name, "/data/ota_package/" + CARE_MAP_ENTRY_NAME]) 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: data_local_tmp_file = '/data/local/tmp/debug.zip' cmds.append(['push', args.otafile, data_local_tmp_file]) cmds.append(['shell', 'su', '0', 'mv', data_local_tmp_file, device_ota_file]) cmds.append(['shell', 'su', '0', 'chcon', 'u:object_r:ota_package_file:s0', 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. payload_url = 'http://127.0.0.1:%d/payload' % DEVICE_PORT serving_range = (0, os.stat(args.otafile).st_size) server_thread = StartServer(args.otafile, serving_range, args.speed_limit) cmds.append( ['reverse', 'tcp:%d' % DEVICE_PORT, 'tcp:%d' % server_thread.port]) finalize_cmds.append(['reverse', '--remove', 'tcp:%d' % DEVICE_PORT]) if args.public_key: payload_key_dir = os.path.dirname(PAYLOAD_KEY_PATH) cmds.append( ['shell', 'su', '0', 'mount', '-t', 'tmpfs', 'tmpfs', payload_key_dir]) # Allow adb push to payload_key_dir cmds.append(['shell', 'su', '0', 'chcon', 'u:object_r:shell_data_file:s0', payload_key_dir]) cmds.append(['push', args.public_key, PAYLOAD_KEY_PATH]) # Allow update_engine to read it. cmds.append(['shell', 'su', '0', 'chcon', '-R', 'u:object_r:system_file:s0', payload_key_dir]) finalize_cmds.append(['shell', 'su', '0', 'umount', payload_key_dir]) try: # The main update command using the configured payload_url. update_cmd = AndroidUpdateCommand(args.otafile, args.secondary, payload_url, args.extra_headers) 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, 5) logging.info('Update took %.3f seconds', (time.perf_counter() - start_time)) return 0 if __name__ == '__main__': sys.exit(main())