Scripts to download files from google storage based on sha1 sums

common.py

+# Copyright (c) 2012 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.
+
+"""Utility functions common to the Google storage scripts."""
+
+import hashlib
+import os
+import re
+import subprocess2

[M-A Ruel, 2013/02/25 15:15:06]

It's not a stdlib, so put it in a subgroup after, e.g.

import hashlib
import os
import re
import sys

import subprocess2


class Gsutil(object):
  ...

[Ryan Tseng, 2013/02/27 02:06:56]

Done.

+import sys
+
+  
+class Gsutil(object):
+  """Call gsutil with some predefined settings.""" 
+  def __init__(self, path, boto_path=None, timeout=None):
+    if not os.path.exists(path):
+      raise OSError('GSUtil not found in %s' % path)
+    self.path = path
+    self.timeout = timeout
+    self.boto_path = boto_path
+
+  def call(self, *args):
+    env = os.environ.copy()
+    if self.boto_path is not None:
+      env['AWS_CREDENTIAL_FILE'] = self.boto_path
+    return subprocess2.call((sys.executable, self.path) + args,
+                            env=env,
+                            timeout=self.timeout)
+
+  def check_call(self, *args):
+    env = os.environ.copy()
+    if self.boto_path is not None:
+      env['AWS_CREDENTIAL_FILE'] = self.boto_path
+    ((out, err), code) = subprocess2.communicate(
+        (sys.executable, self.path) + args,
+        stdout=subprocess2.PIPE,
+        stderr=subprocess2.PIPE,
+        env=env,
+        timeout=self.timeout)
+
+    # Parse output.
+    status_code_match = re.search('status=([0-9]+)', err)
+    if status_code_match:
+      return int(status_code_match.groups(1))
+    elif ('You are attempting to access protected data with '
+          'no configured credentials.' in err):
+      return (403, out, err)
+    elif 'No such object' in err:
+      return (404, out, err)
+    else:
+      return (code, out, err)
+
+  def clone(self):
+    return Gsutil(self.path, self.boto_path, self.timeout)
+
+
+def GetSHA1(filename):
+  sha1 = hashlib.sha1()
+  with open(filename, 'rb') as f:
+    while True:
+      # Read in 1mb chunks, so it doesn't all have to be loaded into memory.
+      chunk = f.read(1024*1024)
+      if not chunk:
+        break
+      sha1.update(chunk)
+  return sha1.hexdigest()
+
+
+def GetMD5(filename, lock, use_md5):
+  # See if we can find an existing MD5 sum stored in a file.
+  if use_md5 and os.path.exists('%s.md5' % filename):
+    with open('%s.md5' % filename) as f:
+      md5_match = re.search('([a-z0-9]{32})', f.read())
+      if md5_match:
+        return md5_match.groups()[0]
+
+  # Calculate the MD5 checksum of the file.
+  md5_calculator = hashlib.md5()
+  with lock:
+    with open(filename, 'rb') as f:
+      while True:
+        chunk = f.read(1024*1024)
+        if not chunk:
+          break
+        md5_calculator.update(chunk)
+  local_md5 = md5_calculator.hexdigest()
+  if use_md5:
+    with open('%s.md5' % filename, 'w') as f:
+      f.write(local_md5)
+  return local_md5