Script to push download_file_types.pb to all platforms, via Component Updater

chrome/browser/resources/safe_browsing/gen_file_type_proto.py

Index: chrome/browser/resources/safe_browsing/gen_file_type_proto.py
diff --git a/chrome/browser/resources/safe_browsing/gen_file_type_proto.py b/chrome/browser/resources/safe_browsing/gen_file_type_proto.py
index 1e55f9a10c140b543f45f5c6214492ca03ab8884..4012b7dbfb5296a6a82853c9a471bac79714cc53 100755
--- a/chrome/browser/resources/safe_browsing/gen_file_type_proto.py
+++ b/chrome/browser/resources/safe_browsing/gen_file_type_proto.py
@@ -11,9 +11,11 @@
 """
 
 import optparse
+import os
 import re
 import subprocess
 import sys
+import traceback
 
 
 def ImportProtoModules(paths):
@@ -93,6 +95,7 @@ def PrunePlatformSettings(file_type, default_settings, platform_type):
 
 def FilterPbForPlatform(full_pb, platform_type):
   """ Return a filtered protobuf for this platform_type """
+  assert type(platform_type) is int, "Bad platform_type type"
 
   new_pb = config_pb2.DownloadFileTypeConfig();
   new_pb.CopyFrom(full_pb)
@@ -127,25 +130,60 @@ def FilterPbForPlatform(full_pb, platform_type):
   return new_pb
 
 
-def GeneratBinaryProtos(opts):
+def FilterForPlatformAndWrite(full_pb, platform_type, outfile):
+  """ Filter and write out a file for this platform """
+  # Filter it
+  filtered_pb = FilterPbForPlatform(full_pb, platform_type);
+
+  # Serialize it
+  binary_pb_str = filtered_pb.SerializeToString()
+
+  # Write it to disk
+  open(outfile, 'wb').write(binary_pb_str)
+
+
+def MakeSubDirs(outfile):
+  """ Make the subdirectories needed to create file |outfile| """
+  dirname = os.path.dirname(outfile)
+  if not os.path.exists(dirname):
+    os.makedirs(dirname)
+
+
+def GenerateBinaryProtos(opts):
+  """ Read the ASCII proto and generate one or more binary protos. """
   # Read the ASCII
   ifile = open(opts.infile, 'r')
   ascii_pb_str = ifile.read()
   ifile.close()
 
-  # Parse it into a structure PB
-  pb = config_pb2.DownloadFileTypeConfig()
-  text_format.Merge(ascii_pb_str, pb)
+  # Parse it into a structured PB
+  full_pb = config_pb2.DownloadFileTypeConfig()
+  text_format.Merge(ascii_pb_str, full_pb)
 
-  ValidatePb(pb);
-  platform_type = PlatformTypes()[opts.type]
-  filtered_pb = FilterPbForPlatform(pb, platform_type);
+  ValidatePb(full_pb);
 
-  # Serialize it
-  binary_pb_str = filtered_pb.SerializeToString()
+  if opts.type is not None:
+    # Just one platform type
+    platform_enum = PlatformTypes()[opts.type]
+    FilterForPlatformAndWrite(full_pb, platform_enum, opts.outfile)
+  else:
+    # Make a separate file for each platform.  Assume the outfile is of the form
+    #    "..../%d/%s/some_file.pb"
+    # This is intended to be run only manually.
+    assert "%s" in opts.outfile and "%d" in opts.outfile, (
+        "--outfile must have %s and %d in it when using --all");
 
-  # Write it to disk
-  open(opts.outfile, 'wb').write(binary_pb_str)
+    for platform_type, platform_enum in PlatformTypes().iteritems():
+      outfile = opts.outfile % (full_pb.version_id, platform_type)
+      MakeSubDirs(outfile)
+      FilterForPlatformAndWrite(full_pb, platform_enum, outfile)
+
+    print "\n\nTo push these files, run the following:"
+    all_dir = os.path.abspath(opts.outfile[:opts.outfile.find("%d") - 1])

[vakh (use Gerrit instead), 2016/05/24 23:31:57]

This assumes a particular format for opts.outfile, which may not be the case
since it is a user input path. That may end up confusing the user.
A more user-friendly way to deal with this may be to split opts.outfile into
opts.outpath and opts.outbasename and always enforce the path for -all operation
to be: opts.outpath + "/%d/%s/" + opts.outbasename

[Nathan Parker, 2016/05/24 23:55:50]

Yup, that's better.  Done.

+    print ("python " +
+           "chrome/browser/resources/safe_browsing/push_file_type_proto.py " +
+           "-d " + all_dir)
+    print "\n\n"
 
 
 def main():
@@ -155,6 +193,10 @@ def main():
                      help='Wrap this script in another python '
                      'execution to disable site-packages.  This is a '
                      'fix for http://crbug.com/605592')
+
+  parser.add_option('-a', '--all', action="store_true", default=False,
+                     help='Write a separate file for every platform. '
+                    'Outfile must have a %d for version and %s for platform.')
   parser.add_option('-t', '--type',
                     help='The platform type. One of android, chromeos, ' +
                     'linux, mac, win')
@@ -174,7 +216,10 @@ def main():
   if opts.wrap:
     # Run this script again with different args to the interpreter.
     command = [sys.executable, '-S', '-s', sys.argv[0]]
-    command += ['-t', opts.type]
+    if opts.type is not None:
+      command += ['-t', opts.type]
+    if opts.all:
+      command += ['-a']
     command += ['-i', opts.infile]
     command += ['-o', opts.outfile]
     for path in opts.path:
@@ -183,16 +228,21 @@ def main():
 
   ImportProtoModules(opts.path)
 
-  if (opts.type not in PlatformTypes()):
+  if (not opts.all and opts.type not in PlatformTypes()):
     print "ERROR: Unknown platform type '%s'" % opts.type
     parser.print_help()
     return 1
 
+  if (opts.all and opts.type):
+    print "ERROR: Can't use both --type and --all"
+    parser.print_help()
+    return 1
+
   try:
-    GeneratBinaryProtos(opts)
+    GenerateBinaryProtos(opts)
   except Exception as e:
-    print "ERROR: Failed to render binary version of %s:\n  %s\n" % (
-        opts.infile, str(e))
+    print "ERROR: Failed to render binary version of %s:\n  %s\n%s" % (
+        opts.infile, str(e), traceback.format_exc())
     return 1