1#!/usr/bin/python
2# Copyright 2016 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
16"""A library for watching changes in the metadata server."""
17
18import functools
19import json
20import logging
21import os
22import socket
23import time
24
25from google_compute_engine.compat import httpclient
26from google_compute_engine.compat import urlerror
27from google_compute_engine.compat import urlparse
28from google_compute_engine.compat import urlrequest
29
30METADATA_SERVER = 'http://metadata.google.internal/computeMetadata/v1'
31
32
33class StatusException(urlerror.HTTPError):
34
35  def __init__(self, response):
36    url = response.geturl()
37    code = response.getcode()
38    message = httpclient.responses.get(code)
39    headers = response.headers
40    super(StatusException, self).__init__(url, code, message, headers, response)
41
42
43def RetryOnUnavailable(func):
44  """Function decorator to retry on a service unavailable exception."""
45
46  @functools.wraps(func)
47  def Wrapper(*args, **kwargs):
48    while True:
49      try:
50        response = func(*args, **kwargs)
51      except (httpclient.HTTPException, socket.error, urlerror.URLError) as e:
52        time.sleep(5)
53        if (isinstance(e, urlerror.HTTPError)
54            and e.getcode() == httpclient.SERVICE_UNAVAILABLE):
55          continue
56        elif isinstance(e, socket.timeout):
57          continue
58        raise
59      else:
60        if response.getcode() == httpclient.OK:
61          return response
62        else:
63          raise StatusException(response)
64  return Wrapper
65
66
67class MetadataWatcher(object):
68  """Watches for changes in metadata."""
69
70  def __init__(self, logger=None, timeout=60):
71    """Constructor.
72
73    Args:
74      logger: logger object, used to write to SysLog and serial port.
75      timeout: int, timeout in seconds for metadata requests.
76    """
77    self.etag = 0
78    self.logger = logger or logging
79    self.timeout = timeout
80
81  @RetryOnUnavailable
82  def _GetMetadataRequest(self, metadata_url, params=None, timeout=None):
83    """Performs a GET request with the metadata headers.
84
85    Args:
86      metadata_url: string, the URL to perform a GET request on.
87      params: dictionary, the query parameters in the GET request.
88      timeout: int, timeout in seconds for metadata requests.
89
90    Returns:
91      HTTP response from the GET request.
92
93    Raises:
94      urlerror.HTTPError: raises when the GET request fails.
95    """
96    headers = {'Metadata-Flavor': 'Google'}
97    params = urlparse.urlencode(params or {})
98    url = '%s?%s' % (metadata_url, params)
99    request = urlrequest.Request(url, headers=headers)
100    request_opener = urlrequest.build_opener(urlrequest.ProxyHandler({}))
101    timeout = timeout or self.timeout
102    return request_opener.open(request, timeout=timeout*1.1)
103
104  def _UpdateEtag(self, response):
105    """Update the etag from an API response.
106
107    Args:
108      response: HTTP response with a header field.
109
110    Returns:
111      bool, True if the etag in the response header updated.
112    """
113    etag = response.headers.get('etag', self.etag)
114    etag_updated = self.etag != etag
115    self.etag = etag
116    return etag_updated
117
118  def _GetMetadataUpdate(
119      self, metadata_key='', recursive=True, wait=True, timeout=None):
120    """Request the contents of metadata server and deserialize the response.
121
122    Args:
123      metadata_key: string, the metadata key to watch for changes.
124      recursive: bool, True if we should recursively watch for metadata changes.
125      wait: bool, True if we should wait for a metadata change.
126      timeout: int, timeout in seconds for returning metadata output.
127
128    Returns:
129      json, the deserialized contents of the metadata server.
130    """
131    metadata_key = os.path.join(metadata_key, '') if recursive else metadata_key
132    metadata_url = os.path.join(METADATA_SERVER, metadata_key)
133    params = {
134        'alt': 'json',
135        'last_etag': self.etag,
136        'recursive': recursive,
137        'timeout_sec': timeout or self.timeout,
138        'wait_for_change': wait,
139    }
140    while True:
141      response = self._GetMetadataRequest(
142          metadata_url, params=params, timeout=timeout)
143      etag_updated = self._UpdateEtag(response)
144      if wait and not etag_updated and not timeout:
145        # Retry until the etag is updated.
146        continue
147      else:
148        # One of the following are true:
149        # - Waiting for change is not required.
150        # - The etag is updated.
151        # - The user specified a request timeout.
152        break
153    return json.loads(response.read().decode('utf-8'))
154
155  def _HandleMetadataUpdate(
156      self, metadata_key='', recursive=True, wait=True, timeout=None,
157      retry=True):
158    """Wait for a successful metadata response.
159
160    Args:
161      metadata_key: string, the metadata key to watch for changes.
162      recursive: bool, True if we should recursively watch for metadata changes.
163      wait: bool, True if we should wait for a metadata change.
164      timeout: int, timeout in seconds for returning metadata output.
165      retry: bool, True if we should retry on failure.
166
167    Returns:
168      json, the deserialized contents of the metadata server.
169    """
170    exception = None
171    while True:
172      try:
173        return self._GetMetadataUpdate(
174            metadata_key=metadata_key, recursive=recursive, wait=wait,
175            timeout=timeout)
176      except (httpclient.HTTPException, socket.error, urlerror.URLError) as e:
177        if not isinstance(e, type(exception)):
178          exception = e
179          self.logger.error('GET request error retrieving metadata. %s.', e)
180        if retry:
181          continue
182        else:
183          break
184
185  def WatchMetadata(
186      self, handler, metadata_key='', recursive=True, timeout=None):
187    """Watch for changes to the contents of the metadata server.
188
189    Args:
190      handler: callable, a function to call with the updated metadata contents.
191      metadata_key: string, the metadata key to watch for changes.
192      recursive: bool, True if we should recursively watch for metadata changes.
193      timeout: int, timeout in seconds for returning metadata output.
194    """
195    while True:
196      response = self._HandleMetadataUpdate(
197          metadata_key=metadata_key, recursive=recursive, wait=True,
198          timeout=timeout)
199      try:
200        handler(response)
201      except Exception as e:
202        self.logger.exception('Exception calling the response handler. %s.', e)
203
204  def GetMetadata(
205      self, metadata_key='', recursive=True, timeout=None, retry=True):
206    """Retrieve the contents of metadata server for a metadata key.
207
208    Args:
209      metadata_key: string, the metadata key to watch for changes.
210      recursive: bool, True if we should recursively watch for metadata changes.
211      timeout: int, timeout in seconds for returning metadata output.
212      retry: bool, True if we should retry on failure.
213
214    Returns:
215      json, the deserialized contents of the metadata server or None if error.
216    """
217    return self._HandleMetadataUpdate(
218        metadata_key=metadata_key, recursive=recursive, wait=False,
219        timeout=timeout, retry=retry)
220