1#!/usr/bin/env python
2#
3# Copyright 2018 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#
8# You may obtain a copy of the License at
9#
10#      http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing, software
13# distributed under the License is distributed on an "AS IS" BASIS,
14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15# See the License for the specific language governing permissions and
16# limitations under the License.
17"""Helper tool for 'oem at-write-persistent-digest' fastboot command.
18
19This tool generates and stages the correct input data format, based on the
20user-provided inputs, for the 'oem at-write-persistent-digest' fastboot command
21for Android Things devices before running the command itself.
22
23The input format is defined elsewhere to be the following:
24
25  - Name length: 4 bytes (little-endian)
26  - Name: 'name length' bytes
27  - Digest length: 4 bytes (little-endian)
28  - Digest: 'digest length' bytes
29
30Digest length can be zero, indicating that any existing digest with the given
31name should be deleted. This corresponds to the '--clear_digest' option for this
32tool.
33
34Digest names must be prefixed with 'avb.persistent_digest.', and if the
35user-provided name does not include that prefix it is added automatically.
36"""
37
38import sys
39
40ver = sys.version_info
41if (ver[0] < 2) or (ver[0] == 2 and ver[1] < 7) or (ver[0] == 3 and ver[1] < 2):
42  print('This script requires Python 2.7+ or 3.2+')
43  sys.exit(1)
44
45import argparse
46import os
47import shutil
48import struct
49import subprocess
50import tempfile
51
52HELP_DESCRIPTION = """Helper script for 'fastboot oem
53at-write-persistent-digest' that generates and stages the required input data
54format."""
55
56AVB_PERSISTENT_DIGEST_PREFIX = 'avb.persistent_digest.'
57
58
59def WritePersistentDigest(name,
60                          digest=None,
61                          clear_digest=False,
62                          serial=None,
63                          verbose=False):
64  if not name.startswith(AVB_PERSISTENT_DIGEST_PREFIX):
65    print("Automatically adding '{}' prefix to persistent value name".format(
66        AVB_PERSISTENT_DIGEST_PREFIX))
67    name = AVB_PERSISTENT_DIGEST_PREFIX + name
68
69  tempdir = tempfile.mkdtemp()
70  try:
71    digest_data = os.path.join(tempdir, 'digest_data')
72
73    with open(digest_data, 'wb') as out:
74      out.write(struct.pack('<I', len(name)))
75      out.write(name)
76      if clear_digest:
77        out.write(struct.pack('<I', 0))
78      else:
79        digest_bytes = bytearray.fromhex(digest)
80        out.write(struct.pack('<I', len(digest_bytes)))
81        out.write(digest_bytes)
82
83    def fastboot_cmd(args):
84      args = ['fastboot'] + (['-s', serial] if serial else []) + args
85      if verbose:
86        print('$ ' + ' '.join(args))
87
88      try:
89        out = subprocess.check_output(
90            args, stderr=subprocess.STDOUT).decode('utf-8')
91      except subprocess.CalledProcessError as e:
92        print(e.output.decode('utf-8'))
93        print("Command '{}' returned non-zero exit status {}".format(
94            ' '.join(e.cmd), e.returncode))
95        sys.exit(1)
96
97      if verbose:
98        print(out)
99
100    fastboot_cmd(['stage', digest_data])
101    fastboot_cmd(['oem', 'at-write-persistent-digest'])
102
103    print("Persistent value '{}' {}".format(
104        name, 'cleared' if clear_digest else 'written'))
105
106  finally:
107    shutil.rmtree(tempdir)
108
109
110if __name__ == '__main__':
111  parser = argparse.ArgumentParser(description=HELP_DESCRIPTION)
112
113  # Optional arguments
114  parser.add_argument(
115      '-v',
116      '--verbose',
117      action='store_true',
118      help='verbose; prints fastboot commands and their output')
119  parser.add_argument(
120      '-s',
121      '--serial',
122      help=
123      "specify device to unlock, either by serial or any other valid value for fastboot's -s arg"
124  )
125
126  # Required arguments
127  parser.add_argument(
128      '--name',
129      required=True,
130      help=
131      "persistent digest name to write, 'avb.persistent_digest.' prefix will be automatically added if not already present"
132  )
133  group = parser.add_mutually_exclusive_group(required=True)
134  group.add_argument(
135      '--clear_digest',
136      action='store_true',
137      help=
138      'clear any existing persistent digest value, rather than writing a new value'
139  )
140  group.add_argument(
141      '--digest',
142      help='persistent digest value to write, as a hex encoded string')
143
144  # Print help if no args given
145  args = parser.parse_args(args=None if sys.argv[1:] else ['-h'])
146
147  WritePersistentDigest(
148      name=args.name,
149      clear_digest=args.clear_digest,
150      digest=args.digest,
151      serial=args.serial,
152      verbose=args.verbose)
153