1# -*- coding: utf-8 -*-
2# Copyright 2014 Google Inc. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15"""Shell tab completion."""
16
17import itertools
18import json
19import threading
20import time
21
22import boto
23
24from boto.gs.acl import CannedACLStrings
25from gslib.storage_url import IsFileUrlString
26from gslib.storage_url import StorageUrlFromString
27from gslib.storage_url import StripOneSlash
28from gslib.util import GetTabCompletionCacheFilename
29from gslib.util import GetTabCompletionLogFilename
30from gslib.wildcard_iterator import CreateWildcardIterator
31
32TAB_COMPLETE_CACHE_TTL = 15
33
34_TAB_COMPLETE_MAX_RESULTS = 1000
35
36_TIMEOUT_WARNING = """
37Tab completion aborted (took >%ss), you may complete the command manually.
38The timeout can be adjusted in the gsutil configuration file.
39""".rstrip()
40
41
42class CompleterType(object):
43  CLOUD_BUCKET = 'cloud_bucket'
44  CLOUD_OBJECT = 'cloud_object'
45  CLOUD_OR_LOCAL_OBJECT = 'cloud_or_local_object'
46  LOCAL_OBJECT = 'local_object'
47  LOCAL_OBJECT_OR_CANNED_ACL = 'local_object_or_canned_acl'
48  NO_OP = 'no_op'
49
50
51class LocalObjectCompleter(object):
52  """Completer object for local files."""
53
54  def __init__(self):
55    # This is only safe to import if argcomplete is present in the install
56    # (which happens for Cloud SDK installs), so import on usage, not on load.
57    # pylint: disable=g-import-not-at-top
58    from argcomplete.completers import FilesCompleter
59    self.files_completer = FilesCompleter()
60
61  def __call__(self, prefix, **kwargs):
62    return self.files_completer(prefix, **kwargs)
63
64
65class LocalObjectOrCannedACLCompleter(object):
66  """Completer object for local files and canned ACLs.
67
68  Currently, only Google Cloud Storage canned ACL names are supported.
69  """
70
71  def __init__(self):
72    self.local_object_completer = LocalObjectCompleter()
73
74  def __call__(self, prefix, **kwargs):
75    local_objects = self.local_object_completer(prefix, **kwargs)
76    canned_acls = [acl for acl in CannedACLStrings if acl.startswith(prefix)]
77    return local_objects + canned_acls
78
79
80class TabCompletionCache(object):
81  """Cache for tab completion results."""
82
83  def __init__(self, prefix, results, timestamp, partial_results):
84    self.prefix = prefix
85    self.results = results
86    self.timestamp = timestamp
87    self.partial_results = partial_results
88
89  @staticmethod
90  def LoadFromFile(filename):
91    """Instantiates the cache from a file.
92
93    Args:
94      filename: The file to load.
95    Returns:
96      TabCompletionCache instance with loaded data or an empty cache
97          if the file cannot be loaded
98    """
99    try:
100      with open(filename, 'r') as fp:
101        cache_dict = json.loads(fp.read())
102        prefix = cache_dict['prefix']
103        results = cache_dict['results']
104        timestamp = cache_dict['timestamp']
105        partial_results = cache_dict['partial-results']
106    except Exception:  # pylint: disable=broad-except
107      # Guarding against incompatible format changes in the cache file.
108      # Erring on the side of not breaking tab-completion in case of cache
109      # issues.
110      prefix = None
111      results = []
112      timestamp = 0
113      partial_results = False
114
115    return TabCompletionCache(prefix, results, timestamp, partial_results)
116
117  def GetCachedResults(self, prefix):
118    """Returns the cached results for prefix or None if not in cache."""
119    current_time = time.time()
120    if current_time - self.timestamp >= TAB_COMPLETE_CACHE_TTL:
121      return None
122
123    results = None
124
125    if prefix == self.prefix:
126      results = self.results
127    elif (not self.partial_results and prefix.startswith(self.prefix)
128          and prefix.count('/') == self.prefix.count('/')):
129      results = [x for x in self.results if x.startswith(prefix)]
130
131    if results is not None:
132      # Update cache timestamp to make sure the cache entry does not expire if
133      # the user is performing multiple completions in a single
134      # bucket/subdirectory since we can answer these requests from the cache.
135      # e.g. gs://prefix<tab> -> gs://prefix-mid<tab> -> gs://prefix-mid-suffix
136      self.timestamp = time.time()
137      return results
138
139  def UpdateCache(self, prefix, results, partial_results):
140    """Updates the in-memory cache with the results for the given prefix."""
141    self.prefix = prefix
142    self.results = results
143    self.partial_results = partial_results
144    self.timestamp = time.time()
145
146  def WriteToFile(self, filename):
147    """Writes out the cache to the given file."""
148    json_str = json.dumps({
149        'prefix': self.prefix,
150        'results': self.results,
151        'partial-results': self.partial_results,
152        'timestamp': self.timestamp,
153    })
154
155    try:
156      with open(filename, 'w') as fp:
157        fp.write(json_str)
158    except IOError:
159      pass
160
161
162class CloudListingRequestThread(threading.Thread):
163  """Thread that performs a listing request for the given URL string."""
164
165  def __init__(self, wildcard_url_str, gsutil_api):
166    """Instantiates Cloud listing request thread.
167
168    Args:
169      wildcard_url_str: The URL to list.
170      gsutil_api: gsutil Cloud API instance to use.
171    """
172    super(CloudListingRequestThread, self).__init__()
173    self.daemon = True
174    self._wildcard_url_str = wildcard_url_str
175    self._gsutil_api = gsutil_api
176    self.results = None
177
178  def run(self):
179    it = CreateWildcardIterator(
180        self._wildcard_url_str, self._gsutil_api).IterAll(
181            bucket_listing_fields=['name'])
182    self.results = [
183        str(c) for c in itertools.islice(it, _TAB_COMPLETE_MAX_RESULTS)]
184
185
186class TimeoutError(Exception):
187  pass
188
189
190class CloudObjectCompleter(object):
191  """Completer object for Cloud URLs."""
192
193  def __init__(self, gsutil_api, bucket_only=False):
194    """Instantiates completer for Cloud URLs.
195
196    Args:
197      gsutil_api: gsutil Cloud API instance to use.
198      bucket_only: Whether the completer should only match buckets.
199    """
200    self._gsutil_api = gsutil_api
201    self._bucket_only = bucket_only
202
203  def _PerformCloudListing(self, wildcard_url, timeout):
204    """Perform a remote listing request for the given wildcard URL.
205
206    Args:
207      wildcard_url: The wildcard URL to list.
208      timeout: Time limit for the request.
209    Returns:
210      Cloud resources matching the given wildcard URL.
211    Raises:
212      TimeoutError: If the listing does not finish within the timeout.
213    """
214    request_thread = CloudListingRequestThread(wildcard_url, self._gsutil_api)
215    request_thread.start()
216    request_thread.join(timeout)
217
218    if request_thread.is_alive():
219      # This is only safe to import if argcomplete is present in the install
220      # (which happens for Cloud SDK installs), so import on usage, not on load.
221      # pylint: disable=g-import-not-at-top
222      import argcomplete
223      argcomplete.warn(_TIMEOUT_WARNING % timeout)
224      raise TimeoutError()
225
226    results = request_thread.results
227
228    return results
229
230  def __call__(self, prefix, **kwargs):
231    if not prefix:
232      prefix = 'gs://'
233    elif IsFileUrlString(prefix):
234      return []
235
236    wildcard_url = prefix + '*'
237    url = StorageUrlFromString(wildcard_url)
238    if self._bucket_only and not url.IsBucket():
239      return []
240
241    timeout = boto.config.getint('GSUtil', 'tab_completion_timeout', 5)
242    if timeout == 0:
243      return []
244
245    start_time = time.time()
246
247    cache = TabCompletionCache.LoadFromFile(GetTabCompletionCacheFilename())
248    cached_results = cache.GetCachedResults(prefix)
249
250    timing_log_entry_type = ''
251    if cached_results is not None:
252      results = cached_results
253      timing_log_entry_type = ' (from cache)'
254    else:
255      try:
256        results = self._PerformCloudListing(wildcard_url, timeout)
257        if self._bucket_only and len(results) == 1:
258          results = [StripOneSlash(results[0])]
259        partial_results = (len(results) == _TAB_COMPLETE_MAX_RESULTS)
260        cache.UpdateCache(prefix, results, partial_results)
261      except TimeoutError:
262        timing_log_entry_type = ' (request timeout)'
263        results = []
264
265    cache.WriteToFile(GetTabCompletionCacheFilename())
266
267    end_time = time.time()
268    num_results = len(results)
269    elapsed_seconds = end_time - start_time
270    _WriteTimingLog(
271        '%s results%s in %.2fs, %.2f results/second for prefix: %s\n' %
272        (num_results, timing_log_entry_type, elapsed_seconds,
273         num_results / elapsed_seconds, prefix))
274
275    return results
276
277
278class CloudOrLocalObjectCompleter(object):
279  """Completer object for Cloud URLs or local files.
280
281  Invokes the Cloud object completer if the input looks like a Cloud URL and
282  falls back to local file completer otherwise.
283  """
284
285  def __init__(self, gsutil_api):
286    self.cloud_object_completer = CloudObjectCompleter(gsutil_api)
287    self.local_object_completer = LocalObjectCompleter()
288
289  def __call__(self, prefix, **kwargs):
290    if IsFileUrlString(prefix):
291      completer = self.local_object_completer
292    else:
293      completer = self.cloud_object_completer
294    return completer(prefix, **kwargs)
295
296
297class NoOpCompleter(object):
298  """Completer that always returns 0 results."""
299
300  def __call__(self, unused_prefix, **unused_kwargs):
301    return []
302
303
304def MakeCompleter(completer_type, gsutil_api):
305  """Create a completer instance of the given type.
306
307  Args:
308    completer_type: The type of completer to create.
309    gsutil_api: gsutil Cloud API instance to use.
310  Returns:
311    A completer instance.
312  Raises:
313    RuntimeError: if completer type is not supported.
314  """
315  if completer_type == CompleterType.CLOUD_OR_LOCAL_OBJECT:
316    return CloudOrLocalObjectCompleter(gsutil_api)
317  elif completer_type == CompleterType.LOCAL_OBJECT:
318    return LocalObjectCompleter()
319  elif completer_type == CompleterType.LOCAL_OBJECT_OR_CANNED_ACL:
320    return LocalObjectOrCannedACLCompleter()
321  elif completer_type == CompleterType.CLOUD_BUCKET:
322    return CloudObjectCompleter(gsutil_api, bucket_only=True)
323  elif completer_type == CompleterType.CLOUD_OBJECT:
324    return CloudObjectCompleter(gsutil_api)
325  elif completer_type == CompleterType.NO_OP:
326    return NoOpCompleter()
327  else:
328    raise RuntimeError(
329        'Unknown completer "%s"' % completer_type)
330
331
332def _WriteTimingLog(message):
333  """Write an entry to the tab completion timing log, if it's enabled."""
334  if boto.config.getbool('GSUtil', 'tab_completion_time_logs', False):
335    with open(GetTabCompletionLogFilename(), 'ab') as fp:
336      fp.write(message)
337
338