• 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_test_job_offloaded_message(self, gcs_uri):
88        """Sends a test job offloaded message to the remote console.
89
90        @param gcs_uri: The test result Google Cloud Storage URI.
91
92        @returns True if the notification is successfully sent.
93            Otherwise, False.
94        """
95        pass
96
97
98# Make it easy to mock out
99def _create_pubsub_client(credential):
100    return pubsub_utils.PubSubClient(credential)
101
102
103class PubSubBasedClient(CloudConsoleClient):
104    """A Cloud PubSub based implementation of the CloudConsoleClient interface.
105    """
106    def __init__(
107            self,
108            credential=moblab_host.MOBLAB_SERVICE_ACCOUNT_LOCATION,
109            pubsub_topic=_PUBSUB_TOPIC):
110        """Constructor.
111
112        @param credential: The service account credential filename. Default to
113            '/home/moblab/.service_account.json'.
114        @param pubsub_topic: The cloud pubsub topic name to use.
115        """
116        super(PubSubBasedClient, self).__init__()
117        self._pubsub_client = _create_pubsub_client(credential)
118        self._pubsub_topic = pubsub_topic
119
120
121    def _create_message(self, data, msg_attributes):
122        """Creates a cloud pubsub notification object.
123
124        @param data: The message data as a string.
125        @param msg_attributes: The message attribute map.
126
127        @returns: A pubsub message object with data and attributes.
128        """
129        message = {}
130        if data:
131            message['data'] = data
132        if msg_attributes:
133            message['attributes'] = msg_attributes
134        return message
135
136    def _create_message_attributes(self, message_type_enum):
137        """Creates a cloud pubsub notification message attribute map.
138
139        Fills in the version, moblab mac address, and moblab id information
140        as attributes.
141
142        @param message_type_enum The message type enum.
143
144        @returns: A pubsub messsage attribute map.
145        """
146        msg_attributes = {}
147        msg_attributes[_get_attribute_name(cpcon.ATTR_MESSAGE_TYPE)] = (
148                _get_message_type_name(message_type_enum))
149        msg_attributes[_get_attribute_name(cpcon.ATTR_MESSAGE_VERSION)] = (
150                CURRENT_MESSAGE_VERSION)
151        msg_attributes[_get_attribute_name(cpcon.ATTR_MOBLAB_MAC_ADDRESS)] = (
152                utils.get_moblab_serial_number())
153        msg_attributes[_get_attribute_name(cpcon.ATTR_MOBLAB_ID)] = (
154                utils.get_moblab_id())
155        return msg_attributes
156
157    def _create_test_job_offloaded_message(self, gcs_uri):
158        """Construct a test result notification.
159
160        TODO(ntang): switch LEGACY to new message format.
161        @param gcs_uri: The test result Google Cloud Storage URI.
162
163        @returns The notification message.
164        """
165        data = base64.b64encode(LEGACY_TEST_OFFLOAD_MESSAGE)
166        msg_attributes = {}
167        msg_attributes[LEGACY_ATTR_VERSION] = CURRENT_MESSAGE_VERSION
168        msg_attributes[LEGACY_ATTR_MOBLAB_MAC] = (
169                utils.get_moblab_serial_number())
170        msg_attributes[LEGACY_ATTR_MOBLAB_ID] = utils.get_moblab_id()
171        msg_attributes[LEGACY_ATTR_GCS_URI] = gcs_uri
172
173        return self._create_message(data, msg_attributes)
174
175
176    def send_test_job_offloaded_message(self, gcs_uri):
177        """Notify the cloud console a test job is offloaded.
178
179        @param gcs_uri: The test result Google Cloud Storage URI.
180
181        @returns True if the notification is successfully sent.
182            Otherwise, False.
183        """
184        logging.info('Notification on gcs_uri %s', gcs_uri)
185        message = self._create_test_job_offloaded_message(gcs_uri)
186        return self._publish_notification(message)
187
188
189    def _publish_notification(self, message):
190        msg_ids = self._pubsub_client.publish_notifications(
191                self._pubsub_topic, [message])
192
193        if msg_ids:
194            logging.debug('Successfully sent out a notification')
195            return True
196        logging.warning('Failed to send notification %s', str(message))
197        return False
198
199    def send_heartbeat(self):
200        """Sends a heartbeat.
201
202        @returns True if the heartbeat notification is successfully sent.
203            Otherwise, False.
204        """
205        logging.info('Sending a heartbeat')
206
207        event = cpcon.Heartbeat()
208        # Don't sent local timestamp for now.
209        data = event.SerializeToString()
210        try:
211            attributes = self._create_message_attributes(
212                    cpcon.MSG_MOBLAB_HEARTBEAT)
213            message = self._create_message(data, attributes)
214        except ValueError:
215            logging.exception('Failed to create message.')
216            return False
217        return self._publish_notification(message)
218
219    def send_event(self, event_type=None, event_data=None):
220        """Sends an event notification to the remote console.
221
222        @param event_type: The event type that is defined in the protobuffer
223            file 'cloud_console.proto'.
224        @param event_data: The event data.
225
226        @returns True if the notification is successfully sent.
227            Otherwise, False.
228        """
229        logging.info('Send an event.')
230        if not event_type:
231            logging.info('Failed to send event without a type.')
232            return False
233
234        event = cpcon.RemoteEventMessage()
235        if event_data:
236            event.data = event_data
237        else:
238            event.data = ''
239        event.type = event_type
240        data = event.SerializeToString()
241        try:
242            attributes = self._create_message_attributes(
243                    cpcon.MSG_MOBLAB_REMOTE_EVENT)
244            message = self._create_message(data, attributes)
245        except ValueError:
246            logging.exception('Failed to create message.')
247            return False
248        return self._publish_notification(message)
249