Home | History | Annotate | Download | only in web
      1 
      2 # Copyright 2016 Google Inc.
      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 import base64
     17 import logging
     18 import os
     19 import subprocess
     20 import tempfile
     21 
     22 from oauth2client import service_account
     23 
     24 _OPENID_SCOPE = 'openid'
     25 
     26 class DashboardRestClient(object):
     27     """Instance of the Dashboard REST client.
     28 
     29     Attributes:
     30         post_cmd: String, The command-line string to post data to the dashboard,
     31                   e.g. 'wget <url> --post-file '
     32         service_json_path: String, The path to the service account keyfile
     33                            created from Google App Engine settings.
     34         auth_token: ServiceAccountCredentials object or None if not
     35                     initialized.
     36     """
     37 
     38     def __init__(self, post_cmd, service_json_path):
     39         self.post_cmd = post_cmd
     40         self.service_json_path = service_json_path
     41         self.auth_token = None
     42 
     43     def Initialize(self):
     44         """Initializes the client with an auth token and access token.
     45 
     46         Returns:
     47             True if the client is initialized successfully, False otherwise.
     48         """
     49         try:
     50             self.auth_token = service_account.ServiceAccountCredentials.from_json_keyfile_name(
     51                 self.service_json_path, [_OPENID_SCOPE])
     52             self.auth_token.get_access_token()
     53         except IOError as e:
     54             logging.error("Error reading service json keyfile: %s", e)
     55             return False
     56         except (ValueError, KeyError) as e:
     57             logging.error("Invalid service json keyfile: %s", e)
     58             return False
     59         return True
     60 
     61 
     62     def _GetToken(self):
     63         """Gets an OAuth2 token using from a service account json keyfile.
     64 
     65         Uses the service account keyfile located at 'service_json_path', provided
     66         to the constructor, to request an OAuth2 token.
     67 
     68         Returns:
     69             String, an OAuth2 token using the service account credentials.
     70             None if authentication fails.
     71         """
     72         return str(self.auth_token.get_access_token().access_token)
     73 
     74     def AddAuthToken(self, post_message):
     75         """Add OAuth2 token to the dashboard message.
     76 
     77         Args:
     78             post_message: DashboardPostMessage, The data to post.
     79 
     80         Returns:
     81             True if successful, False otherwise
     82         """
     83         token = self._GetToken()
     84         if not token:
     85             return False
     86 
     87         post_message.access_token = token
     88 
     89     def PostData(self, post_message):
     90         """Post data to the dashboard database.
     91 
     92         Puts data into the dashboard database using its proto REST endpoint.
     93 
     94         Args:
     95             post_message: DashboardPostMessage, The data to post.
     96 
     97         Returns:
     98             True if successful, False otherwise
     99         """
    100         post_bytes = base64.b64encode(post_message.SerializeToString())
    101 
    102         with tempfile.NamedTemporaryFile(delete=False) as file:
    103             file.write(post_bytes)
    104         p = subprocess.Popen(
    105             self.post_cmd.format(path=file.name),
    106             shell=True,
    107             stdout=subprocess.PIPE,
    108             stderr=subprocess.PIPE)
    109         output, err = p.communicate()
    110         os.remove(file.name)
    111 
    112         if p.returncode or err:
    113             logging.error("Row insertion failed: %s", err)
    114             return False
    115         return True
    116