OLD | NEW |
---|---|
(Empty) | |
1 #!/usr/bin/env python | |
2 # Copyright (c) 2014 The Chromium Authors. All rights reserved. | |
3 # Use of this source code is governed by a BSD-style license that can be | |
4 # found in the LICENSE file. | |
5 | |
6 """Run a pinned gsutil.""" | |
7 | |
8 | |
9 import argparse | |
10 import shutil | |
11 import zipfile | |
12 import hashlib | |
13 import base64 | |
14 import os | |
15 import sys | |
16 import json | |
17 import urllib | |
18 import subprocess | |
19 | |
20 | |
21 GSUTIL_URL = 'https://storage.googleapis.com/pub/' | |
22 API_URL = 'https://www.googleapis.com/storage/v1/b/pub/o/' | |
23 | |
24 THIS_DIR = os.path.dirname(os.path.abspath(__file__)) | |
25 DEFAULT_BIN_DIR = os.path.join(THIS_DIR, 'external_bin', 'gsutil') | |
26 DEFAULT_FALLBACK_GSUTIL = os.path.join( | |
27 THIS_DIR, 'third_party', 'gsutil', 'gsutil') | |
28 | |
29 | |
30 class SubprocessError(Exception): | |
31 pass | |
32 class InvalidGsutilError(Exception): | |
pgervais
2014/12/10 00:19:35
Missing blank line
hinoka
2014/12/10 01:46:46
Done.
| |
33 pass | |
34 | |
35 | |
36 def call(args, verbose=True, **kwargs): | |
37 kwargs['stdout'] = subprocess.PIPE | |
38 kwargs['stderr'] = subprocess.STDOUT | |
39 proc = subprocess.Popen(args, **kwargs) | |
40 out = [] | |
41 for line in proc.stdout: | |
42 out.append(line) | |
43 if verbose: | |
44 sys.stdout.write(line) | |
45 code = proc.wait() | |
46 if code: | |
47 raise SubprocessError('%s failed with %s' % (args, code)) | |
48 return ''.join(out) | |
49 | |
50 | |
51 def download_gsutil(version, target_dir): | |
52 """Downloads gsutil into the target_dir.""" | |
53 filename = 'gsutil_%s.zip' % version | |
54 target_filename = os.path.join(target_dir, filename) | |
55 | |
56 # Check if the target exists already. | |
57 if os.path.exists(target_filename): | |
58 md5_calc = hashlib.md5() | |
59 with open(target_filename, 'rb') as f: | |
60 while True: | |
61 buf = f.read(4096) | |
62 if not buf: | |
63 break | |
64 md5_calc.update(buf) | |
65 local_md5 = md5_calc.hexdigest() | |
66 | |
67 metadata_url = '%s%s' % (API_URL, filename) | |
68 metadata = json.load(urllib.urlopen(metadata_url)) | |
69 remote_md5 = base64.b64decode(metadata['md5Hash']) | |
70 | |
71 if local_md5 == remote_md5: | |
72 return target_filename | |
73 os.remove(target_filename) | |
74 | |
75 # Do the download. | |
76 url = '%s%s' % (GSUTIL_URL, filename) | |
77 u = urllib.urlopen(url) | |
78 with open(target_filename, 'wb') as f: | |
79 while True: | |
80 buf = u.read(4096) | |
81 if not buf: | |
82 break | |
83 f.write(buf) | |
84 return target_filename | |
85 | |
86 | |
87 def check_gsutil(gsutil_bin): | |
88 """Run gsutil version and make sure it runs.""" | |
89 try: | |
90 call([sys.executable, gsutil_bin, 'version'], verbose=False) | |
91 return True | |
92 except SubprocessError: | |
93 return False | |
94 | |
95 | |
96 def ensure_gsutil(version, target): | |
97 bin_dir = os.path.join(target, 'gsutil_%s' % version) | |
98 gsutil_bin = os.path.join(bin_dir, 'gsutil', 'gsutil') | |
99 if os.path.isfile(gsutil_bin) and check_gsutil(gsutil_bin): | |
100 # Everything is awesome! we're all done here. | |
101 return gsutil_bin | |
102 | |
103 if os.path.isdir(bin_dir): | |
104 # Clean up if we're redownloading a corrupted gsutil. | |
105 shutil.rmtree(bin_dir) | |
106 cache_dir = os.path.join(target, '.cache_dir') | |
107 if not os.path.isdir(cache_dir): | |
108 os.makedirs(cache_dir) | |
109 target_zip_filename = download_gsutil(version, cache_dir) | |
110 with zipfile.ZipFile(target_zip_filename, 'r') as target_zip: | |
111 target_zip.extractall(bin_dir) | |
112 | |
113 # Final check that the gsutil bin is okay. This should never fail. | |
114 if not check_gsutil(gsutil_bin): | |
115 raise InvalidGsutilError() | |
116 | |
117 return gsutil_bin | |
118 | |
119 | |
120 def run_gsutil(force_version, fallback, target, args): | |
121 if force_version: | |
122 gsutil_bin = ensure_gsutil(force_version, target) | |
123 else: | |
124 gsutil_bin = fallback | |
125 cmd = [sys.executable, gsutil_bin] + args | |
126 call(cmd) | |
127 | |
128 | |
129 def parse_args(): | |
130 parser = argparse.ArgumentParser() | |
131 parser.add_argument('--force_version') | |
pgervais
2014/12/10 00:19:35
Please use --force-version instead.
hinoka
2014/12/10 01:46:46
Done.
| |
132 parser.add_argument('--fallback', default=DEFAULT_FALLBACK_GSUTIL) | |
133 parser.add_argument('--target', default=DEFAULT_BIN_DIR) | |
134 parser.add_argument('args', nargs=argparse.REMAINDER) | |
135 | |
136 args = parser.parse_args() | |
137 return args.force_version, args.fallback, args.target, args.args | |
138 | |
139 | |
140 def main(): | |
141 force_version, fallback, target, args = parse_args() | |
142 run_gsutil(force_version, fallback, target, args) | |
143 | |
144 if __name__ == '__main__': | |
145 sys.exit(main()) | |
OLD | NEW |