From 4f7f855728824928e307415b0293f4b139545413 Mon Sep 17 00:00:00 2001
From: David James <davidjames@google.com>
Date: Wed, 20 Oct 2010 14:48:04 -0700
Subject: [PATCH] Update prebuilt.py to filter Packages database.

This ensures that the Packages database doesn't contain private packages.
This is necessary so that the database matches up with the actual files
we upload.

BUG=chromium-os:4843
TEST=Ran unit test

Change-Id: I0355b4ab867a0d08396e45b04523a487951475f4

Review URL: http://codereview.chromium.org/3930001
---
 prebuilt.py          | 62 ++++++++++++++++++++++++++++++++++++++++++++
 prebuilt_unittest.py | 36 +++++++++++++++++++++++++
 2 files changed, 98 insertions(+)

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()