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