Add --prefer-integrity option to manifest_fixer.py

If provided (--prefer-integrity=true/false), the script will set the
value in the the manifest.  The script will fail if the value mismatches
the original in the manifest, if any.

Test: scripts/manifest_fixer_test.py
Test: aapt dumps the attribute and observe
Bug: 112037137
Change-Id: I2b333a7c0747dbcbed4d419f1c9ed46d4a4c98e9

2 files changed, 63 insertions, 0 deletions

diff --git a/scripts/manifest_fixer.py b/scripts/manifest_fixer.py
index 9256cb29..07925df6 100755
--- a/scripts/manifest_fixer.py
+++ b/scripts/manifest_fixer.py
@@ -61,6 +61,9 @@ def parse_args():
                       help='specify additional <uses-library> tag to add. android:requred is set to false')
   parser.add_argument('--uses-non-sdk-api', dest='uses_non_sdk_api', action='store_true',
                       help='manifest is for a package built against the platform')
+  parser.add_argument('--prefer-integrity', type=bool, dest='prefer_integrity',
+                      help=('specify if the app prefers strict integrity. Should not be conflict if ' +
+                            'already declared in the manifest.'))
   parser.add_argument('input', help='input AndroidManifest.xml file')
   parser.add_argument('output', help='output AndroidManifest.xml file')
   return parser.parse_args()
@@ -269,6 +272,28 @@ def add_uses_non_sdk_api(doc):
     application.setAttributeNode(attr)
 
 
+def add_prefer_integrity(doc):
+  manifest = parse_manifest(doc)
+  elems = get_children_with_tag(manifest, 'application')
+  application = elems[0] if len(elems) == 1 else None
+  if len(elems) > 1:
+    raise RuntimeError('found multiple <application> tags')
+  elif not elems:
+    application = doc.createElement('application')
+    indent = get_indent(manifest.firstChild, 1)
+    first = manifest.firstChild
+    manifest.insertBefore(doc.createTextNode(indent), first)
+    manifest.insertBefore(application, first)
+
+  attr = application.getAttributeNodeNS(android_ns, 'preferIntegrity')
+  if attr is None:
+    attr = doc.createAttributeNS(android_ns, 'android:preferIntegrity')
+    attr.value = 'true'
+    application.setAttributeNode(attr)
+  elif attr.value != 'true':
+    raise RuntimeError('existing attribute mismatches the option of --prefer-integrity')
+
+
 def write_xml(f, doc):
   f.write('<?xml version="1.0" encoding="utf-8"?>\n')
   for node in doc.childNodes:
@@ -296,6 +321,9 @@ def main():
     if args.uses_non_sdk_api:
       add_uses_non_sdk_api(doc)
 
+    if args.prefer_integrity:
+      add_prefer_integrity(doc)
+
     with open(args.output, 'wb') as f:
       write_xml(f, doc)
 
diff --git a/scripts/manifest_fixer_test.py b/scripts/manifest_fixer_test.py
index d1d401a7..a621445b 100755
--- a/scripts/manifest_fixer_test.py
+++ b/scripts/manifest_fixer_test.py
@@ -346,5 +346,40 @@ class AddUsesNonSdkApiTest(unittest.TestCase):
     self.assertEqual(output, expected)
 
 
+class PreferIntegrityTest(unittest.TestCase):
+  """Unit tests for add_prefer_integrity function."""
+
+  def run_test(self, input_manifest):
+    doc = minidom.parseString(input_manifest)
+    manifest_fixer.add_prefer_integrity(doc)
+    output = StringIO.StringIO()
+    manifest_fixer.write_xml(output, doc)
+    return output.getvalue()
+
+  manifest_tmpl = (
+      '<?xml version="1.0" encoding="utf-8"?>\n'
+      '<manifest xmlns:android="http://schemas.android.com/apk/res/android">\n'
+      '    <application%s/>\n'
+      '</manifest>\n')
+
+  def prefer_integrity(self, value):
+    return ' android:preferIntegrity="%s"' % value
+
+  def test_manifest_with_undeclared_preference(self):
+    manifest_input = self.manifest_tmpl % ''
+    expected = self.manifest_tmpl % self.prefer_integrity('true')
+    output = self.run_test(manifest_input)
+    self.assertEqual(output, expected)
+
+  def test_manifest_with_prefer_integrity(self):
+    manifest_input = self.manifest_tmpl % self.prefer_integrity('true')
+    expected = manifest_input
+    output = self.run_test(manifest_input)
+    self.assertEqual(output, expected)
+
+  def test_manifest_with_not_prefer_integrity(self):
+    manifest_input = self.manifest_tmpl % self.prefer_integrity('false')
+    self.assertRaises(RuntimeError, self.run_test, manifest_input)
+
 if __name__ == '__main__':
   unittest.main()