|
# Copyright 2016 Google Inc.
|
#
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
# you may not use this file except in compliance with the License.
|
# You may obtain a copy of the License at
|
#
|
# http://www.apache.org/licenses/LICENSE-2.0
|
#
|
# Unless required by applicable law or agreed to in writing, software
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
# See the License for the specific language governing permissions and
|
# limitations under the License.
|
|
import base64
|
import logging
|
import os
|
import subprocess
|
import tempfile
|
|
from oauth2client import service_account
|
|
_OPENID_SCOPE = 'openid'
|
|
class DashboardRestClient(object):
|
"""Instance of the Dashboard REST client.
|
|
Attributes:
|
post_cmd: String, The command-line string to post data to the dashboard,
|
e.g. 'wget <url> --post-file '
|
service_json_path: String, The path to the service account keyfile
|
created from Google App Engine settings.
|
auth_token: ServiceAccountCredentials object or None if not
|
initialized.
|
"""
|
|
def __init__(self, post_cmd, service_json_path):
|
self.post_cmd = post_cmd
|
self.service_json_path = service_json_path
|
self.auth_token = None
|
|
def Initialize(self):
|
"""Initializes the client with an auth token and access token.
|
|
Returns:
|
True if the client is initialized successfully, False otherwise.
|
"""
|
try:
|
self.auth_token = service_account.ServiceAccountCredentials.from_json_keyfile_name(
|
self.service_json_path, [_OPENID_SCOPE])
|
self.auth_token.get_access_token()
|
except IOError as e:
|
logging.error("Error reading service json keyfile: %s", e)
|
return False
|
except (ValueError, KeyError) as e:
|
logging.error("Invalid service json keyfile: %s", e)
|
return False
|
return True
|
|
|
def _GetToken(self):
|
"""Gets an OAuth2 token using from a service account json keyfile.
|
|
Uses the service account keyfile located at 'service_json_path', provided
|
to the constructor, to request an OAuth2 token.
|
|
Returns:
|
String, an OAuth2 token using the service account credentials.
|
None if authentication fails.
|
"""
|
return str(self.auth_token.get_access_token().access_token)
|
|
def AddAuthToken(self, post_message):
|
"""Add OAuth2 token to the dashboard message.
|
|
Args:
|
post_message: DashboardPostMessage, The data to post.
|
|
Returns:
|
True if successful, False otherwise
|
"""
|
token = self._GetToken()
|
if not token:
|
return False
|
|
post_message.access_token = token
|
|
def PostData(self, post_message):
|
"""Post data to the dashboard database.
|
|
Puts data into the dashboard database using its proto REST endpoint.
|
|
Args:
|
post_message: DashboardPostMessage, The data to post.
|
|
Returns:
|
True if successful, False otherwise
|
"""
|
post_bytes = base64.b64encode(post_message.SerializeToString())
|
|
with tempfile.NamedTemporaryFile(delete=False) as file:
|
file.write(post_bytes)
|
p = subprocess.Popen(
|
self.post_cmd.format(path=file.name),
|
shell=True,
|
stdout=subprocess.PIPE,
|
stderr=subprocess.PIPE)
|
output, err = p.communicate()
|
os.remove(file.name)
|
|
if p.returncode or err:
|
logging.error("Row insertion failed: %s", err)
|
return False
|
return True
|
