diff --git a/prebuilt.py b/prebuilt.py index cf03157a9f..d14c801dd0 100755 --- a/prebuilt.py +++ b/prebuilt.py @@ -9,6 +9,7 @@ import optparse import os import re import sys +import tempfile from chromite.lib import cros_build_lib """ @@ -190,6 +191,63 @@ def ShouldFilterPackage(file_path): return False +def _ShouldFilterPackageFileSection(section): + """Return whether an section in the package file should be filtered out. + + Args: + section: The section, as a list of strings. + + Returns: + True if the section should be excluded. + """ + + for line in section: + if line.startswith("CPV: "): + package = line.replace("CPV: ", "").rstrip() + if ShouldFilterPackage(package): + return True + else: + return False + + +def FilterPackagesFile(packages_filename): + """Read a portage Packages file and filter out private packages. + + The new, filtered packages file is written to a temporary file. + + Args: + packages_filename: The filename of the Packages file. + + Returns: + filtered_packages: A filtered Packages file, as a NamedTemporaryFile. + """ + + packages_file = open(packages_filename) + filtered_packages = tempfile.NamedTemporaryFile() + section = [] + for line in packages_file: + if line == "\n": + if not _ShouldFilterPackageFileSection(section): + # Looks like this section doesn't contain a private package. Write it + # out. + filtered_packages.write("".join(section)) + + # Start next section. + section = [] + + section.append(line) + else: + if not _ShouldFilterPackageFileSection(section): + filtered_packages.write("".join(section)) + packages_file.close() + + # Flush contents to disk. + filtered_packages.flush() + filtered_packages.seek(0) + + return filtered_packages + + def _GsUpload(args): """Upload to GS bucket. @@ -203,6 +261,10 @@ def _GsUpload(args): if ShouldFilterPackage(local_file): return + if local_file.endswith("/Packages"): + filtered_packages_file = FilterPackagesFile(local_file) + local_file = filtered_packages_file.name + cmd = '%s cp -a public-read %s %s' % (_GSUTIL_BIN, local_file, remote_file) # TODO(scottz): port to use _Run or similar when it is available in # cros_build_lib. diff --git a/prebuilt_unittest.py b/prebuilt_unittest.py index 29a02faffd..4fe0c91cc9 100755 --- a/prebuilt_unittest.py +++ b/prebuilt_unittest.py @@ -191,5 +191,41 @@ class TestPrebuilt(unittest.TestCase): 'asdfasdf') +class TestPackagesFileFiltering(unittest.TestCase): + + def setUp(self): + self.mox = mox.Mox() + + def tearDown(self): + self.mox.UnsetStubs() + self.mox.VerifyAll() + + def testFilterAllPackages(self): + self.mox.StubOutWithMock(prebuilt, 'ShouldFilterPackage') + prebuilt.ShouldFilterPackage("public1").AndReturn(False) + prebuilt.ShouldFilterPackage("private").AndReturn(True) + prebuilt.ShouldFilterPackage("public2").AndReturn(False) + full_packages_file = [ + "foo: bar\n", "\n", + "CPV: public1\n", "foo: bar1\n", "\n", + "CPV: private\n", "foo: bar2\n", "\n", + "CPV: public2\n", "foo: bar3\n", "\n", + ] + private_packages_file = [ + "foo: bar\n", "\n", + "CPV: public1\n", "foo: bar1\n", "\n", + "CPV: public2\n", "foo: bar3\n", "\n", + ] + self.mox.ReplayAll() + temp_packages_file = tempfile.NamedTemporaryFile() + temp_packages_file.write("".join(full_packages_file)) + temp_packages_file.flush() + new_packages_file = prebuilt.FilterPackagesFile(temp_packages_file.name) + new_contents = open(new_packages_file.name).read() + self.assertEqual("".join(private_packages_file), new_contents) + self.assertEqual("".join(private_packages_file), new_packages_file.read()) + new_packages_file.close() + + if __name__ == '__main__': unittest.main()