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 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