commit bb73388acf41cf7e94e485dfad6e549b1710b518
author Tao Bao <tbao@google.com> Wed Jul 24 23:31:19 2019 -0700
committer Tao Bao <tbao@google.com> Wed Jul 24 23:34:25 2019 -0700
tree ec3ac9b6fe6e057861dfa93c26e0539a81539b68
parent 6722de127788be23c051b17424ae4d239da3ac6e

releasetools: Fix the use of StringIO.

Based on the actual semantics, it actually wants an in-memory _bytes_
buffer (io.BytesIO), especially when running with Python 3. This CL
fixes the issue and adds a unittest.

Bug: 131631303
Test: python -m unittest test_sign_target_files_apks
Test: python3 -m unittest test_sign_target_files_apks
Change-Id: I3fb067acc26713f1842e831225607779fd0d1b7e

tools/releasetools/sign_target_files_apks.py

diff --git a/tools/releasetools/sign_target_files_apks.py b/tools/releasetools/sign_target_files_apks.py
index 4cb3a37..2f3642b 100755
--- a/tools/releasetools/sign_target_files_apks.py
+++ b/tools/releasetools/sign_target_files_apks.py
@@ -111,6 +111,7 @@
 import copy
 import errno
 import gzip
+import io
 import itertools
 import logging
 import os
@@ -746,12 +747,7 @@
     filename: The archive name in the output zip.
     keys: A list of public keys to use during OTA package verification.
   """
-
-  try:
-    from StringIO import StringIO
-  except ImportError:
-    from io import StringIO
-  temp_file = StringIO()
+  temp_file = io.BytesIO()
   certs_zip = zipfile.ZipFile(temp_file, "w")
   for k in keys:
     common.ZipWrite(certs_zip, k)

tools/releasetools/test_sign_target_files_apks.py

diff --git a/tools/releasetools/test_sign_target_files_apks.py b/tools/releasetools/test_sign_target_files_apks.py
index 0100729..70c147e 100644
--- a/tools/releasetools/test_sign_target_files_apks.py
+++ b/tools/releasetools/test_sign_target_files_apks.py
@@ -15,6 +15,7 @@
 #
 
 import base64
+import io
 import os.path
 import zipfile
 
@@ -22,7 +23,7 @@
 import test_utils
 from sign_target_files_apks import (
     CheckApkAndApexKeysAvailable, EditTags, GetApkFileInfo, ReadApexKeysInfo,
-    ReplaceCerts, ReplaceVerityKeyId, RewriteProps)
+    ReplaceCerts, ReplaceVerityKeyId, RewriteProps, WriteOtacerts)
 
 
 class SignTargetFilesApksTest(test_utils.ReleaseToolsTestCase):
@@ -236,6 +237,22 @@
     }
     self.assertEqual(output_xml, ReplaceCerts(input_xml))
 
+  def test_WriteOtacerts(self):
+    certs = [
+        os.path.join(self.testdata_dir, 'platform.x509.pem'),
+        os.path.join(self.testdata_dir, 'media.x509.pem'),
+        os.path.join(self.testdata_dir, 'testkey.x509.pem'),
+    ]
+    entry_name = 'SYSTEM/etc/security/otacerts.zip'
+    output_file = common.MakeTempFile(suffix='.zip')
+    with zipfile.ZipFile(output_file, 'w') as output_zip:
+      WriteOtacerts(output_zip, entry_name, certs)
+    with zipfile.ZipFile(output_file) as input_zip:
+      self.assertIn(entry_name, input_zip.namelist())
+      otacerts_file = io.BytesIO(input_zip.read(entry_name))
+      with zipfile.ZipFile(otacerts_file) as otacerts_zip:
+        self.assertEqual(3, len(otacerts_zip.namelist()))
+
   def test_CheckApkAndApexKeysAvailable(self):
     input_file = common.MakeTempFile(suffix='.zip')
     with zipfile.ZipFile(input_file, 'w') as input_zip: