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, 'bin', 'gsutil') | |
dnj
2014/11/20 23:03:02
My only concern with 'bin' is that it's part of th
Ryan Tseng
2014/11/21 02:14:55
Good point, I'll use external_bin
| |
26 | |
27 | |
28 class SubprocessError(Exception): | |
29 pass | |
30 class InvalidGsutilError(Exception): | |
31 pass | |
32 | |
33 | |
34 def call(args, verbose=True, **kwargs): | |
35 kwargs['stdout'] = subprocess.PIPE | |
36 kwargs['stderr'] = subprocess.STDOUT | |
37 proc = subprocess.Popen(args, **kwargs) | |
38 out = [] | |
39 for line in proc.stdout: | |
40 out.append(line) | |
41 if verbose: | |
42 sys.stdout.write(line) | |
43 code = proc.wait() | |
44 if code: | |
45 raise SubprocessError('%s failed with %s' % (args, code)) | |
46 return ''.join(out) | |
47 | |
48 | |
49 def download_gsutil(version, target_dir): | |
50 """Downloads gsutil into the target_dir.""" | |
51 filename = 'gsutil_%s.zip' % version | |
52 target_filename = os.path.join(target_dir, filename) | |
53 | |
54 # Check if the target exists already. | |
55 if os.path.exists(target_filename): | |
56 md5_calc = hashlib.md5() | |
57 with open(target_filename, 'rb') as f: | |
58 while True: | |
59 buf = f.read(4096) | |
60 if not buf: | |
61 break | |
62 md5_calc.update(buf) | |
63 local_md5 = md5_calc.hexdigest() | |
64 | |
65 metadata_url = '%s%s' % (API_URL, filename) | |
66 metadata = json.load(urllib.urlopen(metadata_url)) | |
67 remote_md5 = base64.b64decode(metadata['md5Hash']) | |
68 | |
69 if local_md5 == remote_md5: | |
70 return target_filename | |
71 os.remove(target_filename) | |
72 | |
73 # Do the download. | |
74 url = '%s%s' % (GSUTIL_URL, filename) | |
75 u = urllib.urlopen(url) | |
76 with open(target_filename, 'wb') as f: | |
77 while True: | |
78 buf = u.read(4096) | |
79 if not buf: | |
80 break | |
81 f.write(buf) | |
82 return target_filename | |
83 | |
84 | |
85 def check_gsutil(gsutil_bin): | |
86 """Run gsutil version and make sure it runs.""" | |
87 try: | |
88 call([gsutil_bin, 'version'], verbose=False) | |
89 return True | |
90 except SubprocessError: | |
91 return False | |
92 | |
93 | |
94 def ensure_gsutil(version, target): | |
95 bin_dir = os.path.join(target, 'gsutil_%s' % version) | |
96 gsutil_bin = os.path.join(bin_dir, 'gsutil', 'gsutil') | |
97 if os.path.isfile(gsutil_bin) and check_gsutil(gsutil_bin): | |
98 # Everything is awesome! we're all done here. | |
99 return gsutil_bin | |
100 | |
101 if os.path.isdir(bin_dir): | |
102 # Clean up if we're redownloading a corrupted gsutil. | |
103 shutil.rmtree(bin_dir) | |
104 cache_dir = os.path.join(target, '.cache_dir') | |
105 if not os.path.isdir(cache_dir): | |
106 os.makedirs(cache_dir) | |
107 target_zip_filename = download_gsutil(version, cache_dir) | |
108 with zipfile.ZipFile(target_zip_filename, 'r') as target_zip: | |
109 target_zip.extractall(bin_dir) | |
110 os.chmod(gsutil_bin, 0755) # python zipfile doesn't set exec bit. | |
dnj
2014/11/20 23:03:02
Windows will ignore this. Does anything special ne
Ryan Tseng
2014/11/21 02:14:55
Actually if i used [sys.executable, <path to gsuti
| |
111 | |
112 # Final check that the gsutil bin is okay. This should never fail. | |
113 if not check_gsutil(gsutil_bin): | |
114 raise InvalidGsutilError() | |
115 | |
116 return gsutil_bin | |
117 | |
118 | |
119 def run_gsutil(_force_version, fallback, _target, args): | |
120 # TODO(hinoka): Actually use this | |
121 # gsutil_bin = ensure_gsutil(force_version, target) | |
122 cmd = [sys.executable, fallback] + args | |
123 call(cmd) | |
124 | |
125 | |
126 def parse_args(): | |
127 parser = argparse.ArgumentParser() | |
128 parser.add_argument('--force_version', default='4.6') | |
129 parser.add_argument('--fallback') | |
130 parser.add_argument('--target', default=DEFAULT_BIN_DIR) | |
131 parser.add_argument('args', nargs='*') | |
dnj
2014/11/20 23:03:02
Replace '*' with 'argparse.REMAINDER' to make argp
Ryan Tseng
2014/11/21 02:14:55
Good catch, done.
| |
132 | |
133 args = parser.parse_args() | |
134 return args.force_version, args.fallback, args.target, args.args | |
135 | |
136 | |
137 def main(): | |
138 force_version, fallback, target, args = parse_args() | |
139 run_gsutil(force_version, fallback, target, args) | |
140 | |
141 if __name__ == '__main__': | |
142 sys.exit(main()) | |
OLD | NEW |