Bug 1268088 - Set Content-Type for uploaded files r=rail a=release DONTBUILD
authorConnor Sheehan <csheehan@mozilla.com>
Mon, 16 May 2016 15:11:24 -0400
changeset 332896 0b344441f45b2e88ee05787ff5e8f2a4b987b802
parent 332895 049fb75980cf45695110c8b6ef105f06d10cbe0c
child 332897 73cd8f727097aa78f01cc0697a55aeae0f1337b8
push id6048
push userkmoir@mozilla.com
push dateMon, 06 Jun 2016 19:02:08 +0000
treeherdermozilla-beta@46d72a56c57d [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersrail, release
bugs1268088
milestone48.0a2
Bug 1268088 - Set Content-Type for uploaded files r=rail a=release DONTBUILD MozReview-Commit-ID: LfS1VSsPwAY
testing/mozharness/scripts/release/beet_mover.py
--- a/testing/mozharness/scripts/release/beet_mover.py
+++ b/testing/mozharness/scripts/release/beet_mover.py
@@ -19,16 +19,17 @@ import sh
 import redo
 
 sys.path.insert(1, os.path.dirname(os.path.dirname(sys.path[0])))
 from mozharness.base.log import FATAL
 from mozharness.base.python import VirtualenvMixin
 from mozharness.base.script import BaseScript
 from mozharness.mozilla.aws import pop_aws_auth_from_env
 import mozharness
+import mimetypes
 
 
 def get_hash(content, hash_type="md5"):
     h = hashlib.new(hash_type)
     h.update(content)
     return h.hexdigest()
 
 
@@ -102,16 +103,28 @@ DEFAULT_EXCLUDES = [
     r"^.*json$",
     r"^.*/host.*$",
     r"^.*/mar-tools/.*$",
     r"^.*robocop.apk$",
     r"^.*contrib.*"
 ]
 CACHE_DIR = 'cache'
 
+MIME_MAP = {
+    '': 'text/plain',
+    '.asc': 'text/plain',
+    '.beet': 'text/plain',
+    '.bundle': 'application/octet-stream',
+    '.bz2': 'application/octet-stream',
+    '.checksums': 'text/plain',
+    '.dmg': 'application/x-iso9660-image',
+    '.mar': 'application/octet-stream',
+    '.xpi': 'application/zip'
+}
+
 
 class BeetMover(BaseScript, VirtualenvMixin, object):
     def __init__(self, aws_creds):
         beetmover_kwargs = {
             'config_options': CONFIG_OPTIONS,
             'all_actions': [
                 # 'clobber',
                 'create-virtualenv',
@@ -149,16 +162,17 @@ class BeetMover(BaseScript, VirtualenvMi
         self.bucket = c['bucket']
         if not all(aws_creds):
             self.fatal('credentials must be passed in env: "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"')
         self.aws_key_id, self.aws_secret_key = aws_creds
         # if excludes is set from command line, use it otherwise use defaults
         self.excludes = self.config.get('excludes', DEFAULT_EXCLUDES)
         dirs = self.query_abs_dirs()
         self.dest_dir = os.path.join(dirs['abs_work_dir'], CACHE_DIR)
+        self.mime_fix()
 
     def activate_virtualenv(self):
         """
         activates virtualenv and adds module imports to a instance wide namespace.
 
         creating and activating a virtualenv onto the currently executing python interpreter is a
         bit black magic. Rather than having import statements added in various places within the
         script, we import them here immediately after we activate the newly created virtualenv
@@ -296,17 +310,19 @@ class BeetMover(BaseScript, VirtualenvMi
         key.key = s3_key  # set key name
 
         self.info("Checking if `{}` already exists".format(s3_key))
         key = bucket.get_key(s3_key)
         if not key:
             self.info("Uploading to `{}`".format(s3_key))
             key = bucket.new_key(s3_key)
             # set key value
-            self.retry(key.set_contents_from_filename, args=[source], error_level=FATAL),
+            mime_type, _ = mimetypes.guess_type(source)
+            self.retry(lambda: key.set_contents_from_filename(source, headers={'Content-Type': mime_type}),
+                       error_level=FATAL),
         else:
             if not get_hash(key.get_contents_as_string()) == get_hash(open(source).read()):
                 # for now, let's halt. If necessary, we can revisit this and allow for overwrites
                 #  to the same buildnum release with different bits
                 self.fatal("`{}` already exists with different checksum.".format(s3_key))
             self.log("`{}` has the same MD5 checksum, not uploading".format(s3_key))
 
     def scan_bits(self):
@@ -332,11 +348,16 @@ class BeetMover(BaseScript, VirtualenvMi
                               os.path.abspath(os.path.dirname(os.path.dirname(mozharness.__file__))), 'external_tools')
         self.run_command([self.query_python_path(), os.path.join(external_tools_path,'extract_and_run_command.py'),
                          '-j{}'.format(self.config['scan_parallelization']),
                          'clamscan', '--no-summary', '--', self.dest_dir])
 
     def _matches_exclude(self, keyname):
          return any(re.search(exclude, keyname) for exclude in self.excludes)
 
+    def mime_fix(self):
+        """ Add mimetypes for custom extensions """
+        mimetypes.init()
+        map(lambda (ext, mime_type,): mimetypes.add_type(mime_type, ext), MIME_MAP.items())
+
 if __name__ == '__main__':
     beet_mover = BeetMover(pop_aws_auth_from_env())
     beet_mover.run_and_exit()