1#!/usr/bin/env python
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
20import argparse
21import BaseHTTPServer
22import logging
23import os
24import socket
25import subprocess
26import sys
27import threading
28import zipfile
29
30
31# The path used to store the OTA package when applying the package from a file.
32OTA_PACKAGE_PATH = '/data/ota_package'
33
34
35def CopyFileObjLength(fsrc, fdst, buffer_size=128 * 1024, copy_length=None):
36  """Copy from a file object to another.
37
38  This function is similar to shutil.copyfileobj except that it allows to copy
39  less than the full source file.
40
41  Args:
42    fsrc: source file object where to read from.
43    fdst: destination file object where to write to.
44    buffer_size: size of the copy buffer in memory.
45    copy_length: maximum number of bytes to copy, or None to copy everything.
46
47  Returns:
48    the number of bytes copied.
49  """
50  copied = 0
51  while True:
52    chunk_size = buffer_size
53    if copy_length is not None:
54      chunk_size = min(chunk_size, copy_length - copied)
55      if not chunk_size:
56        break
57    buf = fsrc.read(chunk_size)
58    if not buf:
59      break
60    fdst.write(buf)
61    copied += len(buf)
62  return copied
63
64
65class AndroidOTAPackage(object):
66  """Android update payload using the .zip format.
67
68  Android OTA packages traditionally used a .zip file to store the payload. When
69  applying A/B updates over the network, a payload binary is stored RAW inside
70  this .zip file which is used by update_engine to apply the payload. To do
71  this, an offset and size inside the .zip file are provided.
72  """
73
74  # Android OTA package file paths.
75  OTA_PAYLOAD_BIN = 'payload.bin'
76  OTA_PAYLOAD_PROPERTIES_TXT = 'payload_properties.txt'
77
78  def __init__(self, otafilename):
79    self.otafilename = otafilename
80
81    otazip = zipfile.ZipFile(otafilename, 'r')
82    payload_info = otazip.getinfo(self.OTA_PAYLOAD_BIN)
83    self.offset = payload_info.header_offset + len(payload_info.FileHeader())
84    self.size = payload_info.file_size
85    self.properties = otazip.read(self.OTA_PAYLOAD_PROPERTIES_TXT)
86
87
88class UpdateHandler(BaseHTTPServer.BaseHTTPRequestHandler):
89  """A HTTPServer that supports single-range requests.
90
91  Attributes:
92    serving_payload: path to the only payload file we are serving.
93  """
94
95  @staticmethod
96  def _ParseRange(range_str, file_size):
97    """Parse an HTTP range string.
98
99    Args:
100      range_str: HTTP Range header in the request, not including "Header:".
101      file_size: total size of the serving file.
102
103    Returns:
104      A tuple (start_range, end_range) with the range of bytes requested.
105    """
106    start_range = 0
107    end_range = file_size
108
109    if range_str:
110      range_str = range_str.split('=', 1)[1]
111      s, e = range_str.split('-', 1)
112      if s:
113        start_range = int(s)
114        if e:
115          end_range = int(e) + 1
116      elif e:
117        if int(e) < file_size:
118          start_range = file_size - int(e)
119    return start_range, end_range
120
121
122  def do_GET(self):  # pylint: disable=invalid-name
123    """Reply with the requested payload file."""
124    if self.path != '/payload':
125      self.send_error(404, 'Unknown request')
126      return
127
128    if not self.serving_payload:
129      self.send_error(500, 'No serving payload set')
130      return
131
132    try:
133      f = open(self.serving_payload, 'rb')
134    except IOError:
135      self.send_error(404, 'File not found')
136      return
137    # Handle the range request.
138    if 'Range' in self.headers:
139      self.send_response(206)
140    else:
141      self.send_response(200)
142
143    stat = os.fstat(f.fileno())
144    start_range, end_range = self._ParseRange(self.headers.get('range'),
145                                              stat.st_size)
146    logging.info('Serving request for %s from %s [%d, %d) length: %d',
147             self.path, self.serving_payload, start_range, end_range,
148             end_range - start_range)
149
150    self.send_header('Accept-Ranges', 'bytes')
151    self.send_header('Content-Range',
152                     'bytes ' + str(start_range) + '-' + str(end_range - 1) +
153                     '/' + str(end_range - start_range))
154    self.send_header('Content-Length', end_range - start_range)
155
156    self.send_header('Last-Modified', self.date_time_string(stat.st_mtime))
157    self.send_header('Content-type', 'application/octet-stream')
158    self.end_headers()
159
160    f.seek(start_range)
161    CopyFileObjLength(f, self.wfile, copy_length=end_range - start_range)
162
163
164class ServerThread(threading.Thread):
165  """A thread for serving HTTP requests."""
166
167  def __init__(self, ota_filename):
168    threading.Thread.__init__(self)
169    # serving_payload is a class attribute and the UpdateHandler class is
170    # instantiated with every request.
171    UpdateHandler.serving_payload = ota_filename
172    self._httpd = BaseHTTPServer.HTTPServer(('127.0.0.1', 0), UpdateHandler)
173    self.port = self._httpd.server_port
174
175  def run(self):
176    try:
177      self._httpd.serve_forever()
178    except (KeyboardInterrupt, socket.error):
179      pass
180    logging.info('Server Terminated')
181
182  def StopServer(self):
183    self._httpd.socket.close()
184
185
186def StartServer(ota_filename):
187  t = ServerThread(ota_filename)
188  t.start()
189  return t
190
191
192def AndroidUpdateCommand(ota_filename, payload_url):
193  """Return the command to run to start the update in the Android device."""
194  ota = AndroidOTAPackage(ota_filename)
195  headers = ota.properties
196  headers += 'USER_AGENT=Dalvik (something, something)\n'
197
198  # headers += 'POWERWASH=1\n'
199  headers += 'NETWORK_ID=0\n'
200
201  return ['update_engine_client', '--update', '--follow',
202          '--payload=%s' % payload_url, '--offset=%d' % ota.offset,
203          '--size=%d' % ota.size, '--headers="%s"' % headers]
204
205
206class AdbHost(object):
207  """Represents a device connected via ADB."""
208
209  def __init__(self, device_serial=None):
210    """Construct an instance.
211
212    Args:
213        device_serial: options string serial number of attached device.
214    """
215    self._device_serial = device_serial
216    self._command_prefix = ['adb']
217    if self._device_serial:
218      self._command_prefix += ['-s', self._device_serial]
219
220  def adb(self, command):
221    """Run an ADB command like "adb push".
222
223    Args:
224      command: list of strings containing command and arguments to run
225
226    Returns:
227      the program's return code.
228
229    Raises:
230      subprocess.CalledProcessError on command exit != 0.
231    """
232    command = self._command_prefix + command
233    logging.info('Running: %s', ' '.join(str(x) for x in command))
234    p = subprocess.Popen(command, universal_newlines=True)
235    p.wait()
236    return p.returncode
237
238
239def main():
240  parser = argparse.ArgumentParser(description='Android A/B OTA helper.')
241  parser.add_argument('otafile', metavar='ZIP', type=str,
242                      help='the OTA package file (a .zip file).')
243  parser.add_argument('--file', action='store_true',
244                      help='Push the file to the device before updating.')
245  parser.add_argument('--no-push', action='store_true',
246                      help='Skip the "push" command when using --file')
247  parser.add_argument('-s', type=str, default='', metavar='DEVICE',
248                      help='The specific device to use.')
249  parser.add_argument('--no-verbose', action='store_true',
250                      help='Less verbose output')
251  args = parser.parse_args()
252  logging.basicConfig(
253      level=logging.WARNING if args.no_verbose else logging.INFO)
254
255  dut = AdbHost(args.s)
256
257  server_thread = None
258  # List of commands to execute on exit.
259  finalize_cmds = []
260  # Commands to execute when canceling an update.
261  cancel_cmd = ['shell', 'su', '0', 'update_engine_client', '--cancel']
262  # List of commands to perform the update.
263  cmds = []
264
265  if args.file:
266    # Update via pushing a file to /data.
267    device_ota_file = os.path.join(OTA_PACKAGE_PATH, 'debug.zip')
268    payload_url = 'file://' + device_ota_file
269    if not args.no_push:
270      cmds.append(['push', args.otafile, device_ota_file])
271    cmds.append(['shell', 'su', '0', 'chown', 'system:cache', device_ota_file])
272    cmds.append(['shell', 'su', '0', 'chmod', '0660', device_ota_file])
273  else:
274    # Update via sending the payload over the network with an "adb reverse"
275    # command.
276    device_port = 1234
277    payload_url = 'http://127.0.0.1:%d/payload' % device_port
278    server_thread = StartServer(args.otafile)
279    cmds.append(
280        ['reverse', 'tcp:%d' % device_port, 'tcp:%d' % server_thread.port])
281    finalize_cmds.append(['reverse', '--remove', 'tcp:%d' % device_port])
282
283  try:
284    # The main update command using the configured payload_url.
285    update_cmd = AndroidUpdateCommand(args.otafile, payload_url)
286    cmds.append(['shell', 'su', '0'] + update_cmd)
287
288    for cmd in cmds:
289      dut.adb(cmd)
290  except KeyboardInterrupt:
291    dut.adb(cancel_cmd)
292  finally:
293    if server_thread:
294      server_thread.StopServer()
295    for cmd in finalize_cmds:
296      dut.adb(cmd)
297
298  return 0
299
300if __name__ == '__main__':
301  sys.exit(main())
302