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

chrome/browser/resources/safe_browsing/gen_file_type_proto.py

 #!/usr/bin/python
 # Copyright 2016 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """
  Convert the ASCII download_file_types.asciipb proto into a binary resource.
 
  We generate a separate variant of the binary proto for each platform,
  each which contains only the values that platform needs.
 """
 
 import optparse
+import os
 import re
 import subprocess
 import sys
+import traceback
 
 
 def ImportProtoModules(paths):
   """Import the protobuf modiles we need. |paths| is list of import paths"""
   for path in paths:
     # Put the path to our proto libraries in front, so that we don't use system
     # protobuf.
     sys.path.insert(1, path)
 
   import download_file_types_pb2 as config_pb2
@@ skipping 59 matching lines @@
 
   # Now clear out the full list and replace it with 1 entry.
   del file_type.platform_settings[:]
   new_setting = file_type.platform_settings.add()
   new_setting.CopyFrom(setting_match)
   new_setting.ClearField('platform')
 
 
 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)
 
   # Ensure there's only one platform_settings for the default.
   PrunePlatformSettings(new_pb.default_file_type, None, platform_type)
 
   # This can be extended if we want to match weird extensions.
   # Just no dots, non-UTF8, or uppercase chars.
   invalid_char_re = re.compile('[^a-z0-9_-]')
@@ skipping 14 matching lines @@
             file_type.extension, file_type.uma_value))
     uma_values_used.add(file_type.uma_value)
 
     # Modify file_type to include only the best match platform_setting.
     PrunePlatformSettings(
         file_type, new_pb.default_file_type.platform_settings[0], 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():
   parser = optparse.OptionParser()
   # TODO(nparker): Remove this once the bug is fixed.
   parser.add_option('-w', '--wrap', action="store_true", default=False,
                      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')
   parser.add_option('-i', '--infile',
                     help='The ASCII DownloadFileType-proto file to read.')
   parser.add_option('-o', '--outfile',
                     help='The binary file to write to.')
   parser.add_option('-p', '--path', action="append",
                     help='Repeat this as needed.  Directory(s) containing ' +
                     'the download_file_types_pb2.py and ' +
                     'google.protobuf.text_format modules')
   (opts, args) = parser.parse_args()
   if opts.infile is None or opts.outfile is None:
     parser.print_help()
     return 1
 
   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:
       command += ['-p', path]
     sys.exit(subprocess.call(command))
 
   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
 
 
 if __name__ == '__main__':
   sys.exit(main())