Send authorization headers when calling Swarming backend.

appengine/swarming/swarming_bot/bot_code/remote_client.py

+# Copyright 2016 The LUCI Authors. All rights reserved.
+# Use of this source code is governed under the Apache License, Version 2.0
+# that can be found in the LICENSE file.
+
+import threading
+
+from utils import net
+
+
+# RemoteClient will attempt to refresh the authentication headers once they are
+# this close to the expiration.
+AUTH_HEADERS_EXPIRATION_SEC = 6*60
+
+
+class InitializationError(Exception):
+  """Raised by RemoteClient.initialize on fatal errors."""
+  def __init__(self, last_error):
+    super(InitializationError, self).__init__('Failed to grab auth headers')
+    self.last_error = last_error
+
+
+class RemoteClient(object):
+  """RemoteClient knows how to make authenticated calls to the backend.
+
+  It also holds in-memory cache of authentication headers and periodically
+  refreshes them (by calling supplied callback, that usually is implemented in
+  terms of bot_config.get_authentication_headers() function).
+
+  If the callback is None, skips authentication (this is used during initial
+  stages of the bot bootstrap).
+  """
+
+  def __init__(self, server, auth_headers_callback):
+    self._server = server
+    self._auth_headers_callback = auth_headers_callback
+    self._uses_auth = False
+    self._lock = threading.Lock()
+
+  def initialize(self, quit_bit):
+    """Grabs initial auth headers, retrying on errors a bunch of times.
+
+    Raises InitializationError if all attempts fail. Aborts attempts and returns
+    if quit_bit is set to signaled state.
+    """
+    # TODO(vadimsh): Implement.
+
+  @property
+  def uses_auth(self):
+    """Returns True if get_authentication_headers() returns some headers.
+
+    If bot_config.get_authentication_headers() is not implement it will return
+    False.
+    """
+    return self._uses_auth
+
+  def get_authentication_headers(self):
+    """Returns a dict with the headers, refreshing them if necessary.
+
+    Will always return a dict (perhaps empty if no auth headers are provided by
+    the callback).
+    """
+    with self._lock:
+      # TODO(vadimsh): Implement.
+      return {}
+
+  def url_read_json(self, url_path, data=None):
+    """Does POST (if data is not None) or GET request to a JSON endpoint."""
+    return net.url_read_json(
+        self._server + url_path,
+        data=data,
+        headers=self.get_authentication_headers(),
+        timeout=AUTH_HEADERS_EXPIRATION_SEC,
+        follow_redirects=False)
+
+  def url_retrieve(self, filepath, url_path):
+    """Fetches the file from the given URL path on the server."""
+    return net.url_retrieve(
+        filepath,
+        self._server + url_path,
+        headers=self.get_authentication_headers(),
+        timeout=AUTH_HEADERS_EXPIRATION_SEC)