• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2016 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Chrome OS Parnter Concole remote actions."""
6
7from __future__ import print_function
8
9import base64
10import logging
11
12import common
13
14from autotest_lib.client.common_lib import global_config
15from autotest_lib.client.common_lib import utils
16from autotest_lib.server.hosts import moblab_host
17from autotest_lib.site_utils import pubsub_utils
18from autotest_lib.site_utils import cloud_console_pb2 as cpcon
19
20
21_PUBSUB_TOPIC = global_config.global_config.get_config_value(
22        'CROS', 'cloud_notification_topic', default=None)
23
24# Current notification version.
25CURRENT_MESSAGE_VERSION = '1'
26
27# Test upload pubsub notification attributes
28LEGACY_ATTR_VERSION = 'version'
29LEGACY_ATTR_GCS_URI = 'gcs_uri'
30LEGACY_ATTR_MOBLAB_MAC = 'moblab_mac_address'
31LEGACY_ATTR_MOBLAB_ID = 'moblab_id'
32# the message data for new test result notification.
33LEGACY_TEST_OFFLOAD_MESSAGE = 'NEW_TEST_RESULT'
34
35
36def is_cloud_notification_enabled():
37    """Checks if cloud pubsub notification is enabled.
38
39    @returns: True if cloud pubsub notification is enabled. Otherwise, False.
40    """
41    return  global_config.global_config.get_config_value(
42        'CROS', 'cloud_notification_enabled', type=bool, default=False)
43
44
45def _get_message_type_name(message_type_enum):
46    """Gets the message type name from message type enum.
47
48    @param message_type_enum: The message type enum.
49
50    @return The corresponding message type name as string, or 'MSG_UNKNOWN'.
51    """
52    return cpcon.MessageType.Name(message_type_enum)
53
54
55def _get_attribute_name(attribute_enum):
56    """Gets the message attribute name from attribte enum.
57
58    @param attribute_enum: The attribute enum.
59
60    @return The corresponding attribute name as string, or 'ATTR_INVALID'.
61    """
62    return cpcon.MessageAttribute.Name(attribute_enum)
63
64
65class CloudConsoleClient(object):
66    """The remote interface to the Cloud Console."""
67    def send_heartbeat(self):
68        """Sends a heartbeat.
69
70        @returns True if the notification is successfully sent.
71            Otherwise, False.
72        """
73        pass
74
75    def send_event(self, event_type=None, event_data=None):
76        """Sends an event notification to the remote console.
77
78        @param event_type: The event type that is defined in the protobuffer
79            file 'cloud_console.proto'.
80        @param event_data: The event data.
81
82        @returns True if the notification is successfully sent.
83            Otherwise, False.
84        """
85        pass
86
87    def send_log(self, msg, level=None, session_id=None):
88        """Sends a log message to the remote console.
89
90        @param msg: The log message.
91        @param level: The logging level.
92        @param session_id: The current session id.
93
94        @returns True if the notification is successfully sent.
95            Otherwise, False.
96        """
97        pass
98
99    def send_alert(self, msg, level=None, session_id=None):
100        """Sends an alert to the remote console.
101
102        @param msg: The alert message.
103        @param level: The logging level.
104        @param session_id: The current session id.
105
106        @returns True if the notification is successfully sent.
107            Otherwise, False.
108        """
109        pass
110
111    def send_test_job_offloaded_message(self, gcs_uri):
112        """Sends a test job offloaded message to the remote console.
113
114        @param gcs_uri: The test result Google Cloud Storage URI.
115
116        @returns True if the notification is successfully sent.
117            Otherwise, False.
118        """
119        pass
120
121
122# Make it easy to mock out
123def _create_pubsub_client(credential):
124    return pubsub_utils.PubSubClient(credential)
125
126
127class PubSubBasedClient(CloudConsoleClient):
128    """A Cloud PubSub based implementation of the CloudConsoleClient interface.
129    """
130    def __init__(
131            self,
132            credential=moblab_host.MOBLAB_SERVICE_ACCOUNT_LOCATION,
133            pubsub_topic=_PUBSUB_TOPIC):
134        """Constructor.
135
136        @param credential: The service account credential filename. Default to
137            '/home/moblab/.service_account.json'.
138        @param pubsub_topic: The cloud pubsub topic name to use.
139        """
140        super(PubSubBasedClient, self).__init__()
141        self._pubsub_client = _create_pubsub_client(credential)
142        self._pubsub_topic = pubsub_topic
143
144
145    def _create_message(self, data, msg_attributes):
146        """Creates a cloud pubsub notification object.
147
148        @param data: The message data as a string.
149        @param msg_attributes: The message attribute map.
150
151        @returns: A pubsub message object with data and attributes.
152        """
153        message = {}
154        if data:
155            message['data'] = data
156        if msg_attributes:
157            message['attributes'] = msg_attributes
158        return message
159
160    def _create_message_attributes(self, message_type_enum):
161        """Creates a cloud pubsub notification message attribute map.
162
163        Fills in the version, moblab mac address, and moblab id information
164        as attributes.
165
166        @param message_type_enum The message type enum.
167
168        @returns: A pubsub messsage attribute map.
169        """
170        msg_attributes = {}
171        msg_attributes[_get_attribute_name(cpcon.ATTR_MESSAGE_TYPE)] = (
172                _get_message_type_name(message_type_enum))
173        msg_attributes[_get_attribute_name(cpcon.ATTR_MESSAGE_VERSION)] = (
174                CURRENT_MESSAGE_VERSION)
175        msg_attributes[_get_attribute_name(cpcon.ATTR_MOBLAB_MAC_ADDRESS)] = (
176                utils.get_moblab_serial_number())
177        msg_attributes[_get_attribute_name(cpcon.ATTR_MOBLAB_ID)] = (
178                utils.get_moblab_id())
179        return msg_attributes
180
181    def _create_test_job_offloaded_message(self, gcs_uri):
182        """Construct a test result notification.
183
184        TODO(ntang): switch LEGACY to new message format.
185        @param gcs_uri: The test result Google Cloud Storage URI.
186
187        @returns The notification message.
188        """
189        data = base64.b64encode(LEGACY_TEST_OFFLOAD_MESSAGE)
190        msg_attributes = {}
191        msg_attributes[LEGACY_ATTR_VERSION] = CURRENT_MESSAGE_VERSION
192        msg_attributes[LEGACY_ATTR_MOBLAB_MAC] = (
193                utils.get_moblab_serial_number())
194        msg_attributes[LEGACY_ATTR_MOBLAB_ID] = utils.get_moblab_id()
195        msg_attributes[LEGACY_ATTR_GCS_URI] = gcs_uri
196
197        return self._create_message(data, msg_attributes)
198
199
200    def send_test_job_offloaded_message(self, gcs_uri):
201        """Notify the cloud console a test job is offloaded.
202
203        @param gcs_uri: The test result Google Cloud Storage URI.
204
205        @returns True if the notification is successfully sent.
206            Otherwise, False.
207        """
208        logging.info('Notification on gcs_uri %s', gcs_uri)
209        message = self._create_test_job_offloaded_message(gcs_uri)
210        return self._publish_notification(message)
211
212
213    def _publish_notification(self, message):
214        msg_ids = self._pubsub_client.publish_notifications(
215                self._pubsub_topic, [message])
216
217        if msg_ids:
218            logging.debug('Successfully sent out a notification')
219            return True
220        logging.warning('Failed to send notification %s', str(message))
221        return False
222
223    def send_heartbeat(self):
224        """Sends a heartbeat.
225
226        @returns True if the heartbeat notification is successfully sent.
227            Otherwise, False.
228        """
229        logging.info('Sending a heartbeat')
230
231        event = cpcon.Heartbeat()
232        # Don't sent local timestamp for now.
233        data = event.SerializeToString()
234        try:
235            attributes = self._create_message_attributes(
236                    cpcon.MSG_MOBLAB_HEARTBEAT)
237            message = self._create_message(data, attributes)
238        except ValueError:
239            logging.exception('Failed to create message.')
240            return False
241        return self._publish_notification(message)
242
243    def send_event(self, event_type=None, event_data=None):
244        """Sends an event notification to the remote console.
245
246        @param event_type: The event type that is defined in the protobuffer
247            file 'cloud_console.proto'.
248        @param event_data: The event data.
249
250        @returns True if the notification is successfully sent.
251            Otherwise, False.
252        """
253        logging.info('Send an event.')
254        if not event_type:
255            logging.info('Failed to send event without a type.')
256            return False
257
258        event = cpcon.RemoteEventMessage()
259        if event_data:
260            event.data = event_data
261        else:
262            event.data = ''
263        event.type = event_type
264        data = event.SerializeToString()
265        try:
266            attributes = self._create_message_attributes(
267                    cpcon.MSG_MOBLAB_REMOTE_EVENT)
268            message = self._create_message(data, attributes)
269        except ValueError:
270            logging.exception('Failed to create message.')
271            return False
272        return self._publish_notification(message)
273