• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
16import base64
17import logging
18import os
19import subprocess
20import tempfile
21
22from oauth2client import service_account
23
24_OPENID_SCOPE = 'openid'
25
26class 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 PostData(self, post_message):
75        """Post data to the dashboard database.
76
77        Puts data into the dashboard database using its proto REST endpoint.
78
79        Args:
80            post_message: DashboardPostMessage, The data to post.
81
82        Returns:
83            True if successful, False otherwise
84        """
85        token = self._GetToken()
86        if not token:
87            return False
88
89        post_message.access_token = token
90        post_bytes = base64.b64encode(post_message.SerializeToString())
91
92        with tempfile.NamedTemporaryFile(delete=False) as file:
93            file.write(post_bytes)
94        p = subprocess.Popen(
95            self.post_cmd.format(path=file.name),
96            shell=True,
97            stdout=subprocess.PIPE,
98            stderr=subprocess.PIPE)
99        output, err = p.communicate()
100        os.remove(file.name)
101
102        if p.returncode or err:
103            logging.error("Row insertion failed: %s", err)
104            return False
105        return True
106