1# Copyright (c) 2012 The Chromium 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
5import binascii
6import copy
7import logging
8import os
9import pprint
10import re
11import time
12import xmlrpclib
13import json
14import urllib2
15import time
16
17import ap_spec
18import web_driver_core_helpers
19
20from autotest_lib.client.common_lib import error
21from autotest_lib.client.common_lib import global_config
22from autotest_lib.client.common_lib.cros.network import ap_constants
23from autotest_lib.client.common_lib.cros.network import xmlrpc_datatypes
24from autotest_lib.client.common_lib.cros.network import xmlrpc_security_types
25from autotest_lib.server.cros.ap_configurators import ap_configurator
26
27try:
28  from selenium import webdriver
29except ImportError:
30  raise ImportError('Could not locate the webdriver package. '
31                    'Did you emerge it into your chroot?')
32
33
34class DynamicAPConfigurator(web_driver_core_helpers.WebDriverCoreHelpers,
35                            ap_configurator.APConfiguratorAbstract):
36    """Base class for objects to configure access points using webdriver."""
37
38
39    def __init__(self, ap_config):
40        """Construct a DynamicAPConfigurator.
41
42        @param ap_config: information from the configuration file
43        @param set_ap_spec: APSpec object that when passed will set all
44                            of the configuration options
45
46        """
47        super(DynamicAPConfigurator, self).__init__()
48        rpm_frontend_server = global_config.global_config.get_config_value(
49                'CROS', 'rpm_frontend_uri')
50        self.rpm_client = xmlrpclib.ServerProxy(
51                rpm_frontend_server, verbose=False)
52
53        # Load the data for the config file
54        self.admin_interface_url = ap_config.get_admin()
55        self.class_name = ap_config.get_class()
56        self._short_name = ap_config.get_model()
57        self.mac_address = ap_config.get_wan_mac()
58        self.host_name = ap_config.get_wan_host()
59        # Get corresponding PDU from host name.
60        self.pdu = re.sub('host\d+', 'rpm1', self.host_name) + '.cros'
61        self.config_data = ap_config
62
63        name_dict = {'Router name': self._short_name,
64                     'Controller class': self.class_name,
65                     '2.4 GHz MAC Address': ap_config.get_bss(),
66                     '5 GHz MAC Address': ap_config.get_bss5(),
67                     'Hostname': ap_config.get_wan_host()}
68
69        self._name = str('%s' % pprint.pformat(name_dict))
70
71        # Set a default band, this can be overriden by the subclasses
72        self.current_band = ap_spec.BAND_2GHZ
73        self._ssid = None
74
75        # Diagnostic members
76        self._command_list = []
77        self._screenshot_list = []
78        self._traceback = None
79
80        self.driver_connection_established = False
81        self.router_on = False
82        self._configuration_success = ap_constants.CONFIG_SUCCESS
83        self._webdriver_port = 9515
84
85        self.ap_spec = None
86        self.webdriver_hostname = None
87
88    def __del__(self):
89        """Cleanup webdriver connections"""
90        try:
91            self.driver.close()
92        except:
93            pass
94
95
96    def __str__(self):
97        """Prettier display of the object"""
98        return('AP Name: %s\n'
99               'BSS: %s\n'
100               'SSID: %s\n'
101               'Short name: %s' % (self.name, self.get_bss(),
102               self._ssid, self.short_name))
103
104
105    @property
106    def configurator_type(self):
107        """Returns the configurator type."""
108        return ap_spec.CONFIGURATOR_DYNAMIC
109
110
111    @property
112    def ssid(self):
113        """Returns the SSID."""
114        return self._ssid
115
116
117    def add_item_to_command_list(self, method, args, page, priority):
118        """
119        Adds commands to be executed against the AP web UI.
120
121        @param method: the method to run
122        @param args: the arguments for the method you want executed
123        @param page: the page on the web ui where to run the method against
124        @param priority: the priority of the method
125
126        """
127        self._command_list.append({'method': method,
128                                   'args': copy.copy(args),
129                                   'page': page,
130                                   'priority': priority})
131
132
133    def reset_command_list(self):
134        """Resets all internal command state."""
135        logging.error('Dumping command list %s', self._command_list)
136        self._command_list = []
137        self.destroy_driver_connection()
138
139
140    def save_screenshot(self):
141        """
142        Stores and returns the screenshot as a base 64 encoded string.
143
144        @returns the screenshot as a base 64 encoded string; if there was
145        an error saving the screenshot None is returned.
146
147        """
148        screenshot = None
149        if self.driver_connection_established:
150            try:
151                # driver.get_screenshot_as_base64 takes a screenshot that is
152                # whatever the size of the window is.  That can be anything,
153                # forcing a size that will get everything we care about.
154                window_size = self.driver.get_window_size()
155                self.driver.set_window_size(2000, 5000)
156                screenshot = self.driver.get_screenshot_as_base64()
157                self.driver.set_window_size(window_size['width'],
158                                            window_size['height'])
159            except Exception as e:
160                # The messages differ based on the webdriver version
161                logging.error('Getting the screenshot failed. %s', e)
162                # TODO (krisr) this too can fail with an exception.
163                self._check_for_alert_in_message(str(e),
164                                                 self._handler(None))
165                logging.error('Alert was handled.')
166                screenshot = None
167            if screenshot:
168                self._screenshot_list.append(screenshot)
169        return screenshot
170
171
172    def get_all_screenshots(self):
173        """Returns a list of screenshots."""
174        return self._screenshot_list
175
176
177    def clear_screenshot_list(self):
178        """Clear the list of currently stored screenshots."""
179        self._screenshot_list = []
180
181
182    def _save_all_pages(self):
183        """Iterate through AP pages, saving screenshots"""
184        self.establish_driver_connection()
185        if not self.driver_connection_established:
186            logging.error('Unable to establish webdriver connection to '
187                          'retrieve screenshots.')
188            return
189        for page in range(1, self.get_number_of_pages() + 1):
190            self.navigate_to_page(page)
191            self.save_screenshot()
192
193
194    def _write_screenshots(self, filename, outputdir):
195        """
196        Writes screenshots to filename in outputdir
197
198        @param filename: a string prefix for screenshot filenames
199        @param outputdir: a string directory name to save screenshots
200
201        """
202        for (i, image) in enumerate(self.get_all_screenshots()):
203            path = os.path.join(outputdir,
204                                str('%s_%d.png' % (filename, (i + 1))))
205            with open(path, 'wb') as f:
206                f.write(image.decode('base64'))
207
208
209    @property
210    def traceback(self):
211        """
212        Returns the traceback of a configuration error as a string.
213
214        Note that if configuration_success returns CONFIG_SUCCESS this will
215        be none.
216
217        """
218        return self._traceback
219
220
221    @traceback.setter
222    def traceback(self, value):
223        """
224        Set the traceback.
225
226        If the APConfigurator crashes use this to store what the traceback
227        was as a string.  It can be used later to debug configurator errors.
228
229        @param value: a string representation of the exception traceback
230
231        """
232        self._traceback = value
233
234
235    def check_webdriver_ready(self, webdriver_hostname, webdriver_port):
236        """Checks if webdriver binary is installed and running.
237
238        @param webdriver_hostname: locked webdriver instance
239        @param webdriver_port: port of the webdriver server
240
241        @returns a string: the address of webdriver running on port.
242
243        @raises TestError: Webdriver is not running.
244        """
245        address = webdriver_hostname + '.cros'
246        url = 'http://%s:%d/session' % (address, webdriver_port)
247        req = urllib2.Request(url, '{"desiredCapabilities":{}}')
248        try:
249            time.sleep(20)
250            response = urllib2.urlopen(req)
251            json_dict = json.loads(response.read())
252            if json_dict['status'] == 0:
253                # Connection was successful, close the session
254                session_url = os.path.join(url, json_dict['sessionId'])
255                req = urllib2.Request(session_url)
256                req.get_method = lambda: 'DELETE'
257                response = urllib2.urlopen(req)
258                logging.info('Webdriver connection established to server %s',
259                            address)
260                return webdriver_hostname
261        except:
262            err = 'Could not establish connection: %s', webdriver_hostname
263            raise error.TestError(err)
264
265
266    @property
267    def webdriver_port(self):
268        """Returns the webdriver port."""
269        return self._webdriver_port
270
271
272    @webdriver_port.setter
273    def webdriver_port(self, value):
274        """
275        Set the webdriver server port.
276
277        @param value: the port number of the webdriver server
278
279        """
280        self._webdriver_port = value
281
282
283    @property
284    def name(self):
285        """Returns a string to describe the router."""
286        return self._name
287
288
289    @property
290    def short_name(self):
291        """Returns a short string to describe the router."""
292        return self._short_name
293
294
295    def get_number_of_pages(self):
296        """Returns the number of web pages used to configure the router.
297
298        Note: This is used internally by apply_settings, and this method must be
299              implemented by the derived class.
300
301        Note: The derived class must implement this method.
302
303        """
304        raise NotImplementedError
305
306
307    def get_supported_bands(self):
308        """Returns a list of dictionaries describing the supported bands.
309
310        Example: returned is a dictionary of band and a list of channels. The
311                 band object returned must be one of those defined in the
312                 __init___ of this class.
313
314        supported_bands = [{'band' : self.band_2GHz,
315                            'channels' : [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]},
316                           {'band' : ap_spec.BAND_5GHZ,
317                            'channels' : [26, 40, 44, 48, 149, 153, 165]}]
318
319        Note: The derived class must implement this method.
320
321        @return a list of dictionaries as described above
322
323        """
324        raise NotImplementedError
325
326
327    def get_bss(self):
328        """Returns the bss of the AP."""
329        if self.current_band == ap_spec.BAND_2GHZ:
330            return self.config_data.get_bss()
331        else:
332            return self.config_data.get_bss5()
333
334
335    def _get_channel_popup_position(self, channel):
336        """Internal method that converts a channel value to a popup position."""
337        supported_bands = self.get_supported_bands()
338        for band in supported_bands:
339            if band['band'] == self.current_band:
340                return band['channels'].index(channel)
341        raise RuntimeError('The channel passed %d to the band %s is not '
342                           'supported.' % (channel, band))
343
344
345    def get_supported_modes(self):
346        """
347        Returns a list of dictionaries describing the supported modes.
348
349        Example: returned is a dictionary of band and a list of modes. The band
350                 and modes objects returned must be one of those defined in the
351                 __init___ of this class.
352
353        supported_modes = [{'band' : ap_spec.BAND_2GHZ,
354                            'modes' : [mode_b, mode_b | mode_g]},
355                           {'band' : ap_spec.BAND_5GHZ,
356                            'modes' : [mode_a, mode_n, mode_a | mode_n]}]
357
358        Note: The derived class must implement this method.
359
360        @return a list of dictionaries as described above
361
362        """
363        raise NotImplementedError
364
365
366    def is_visibility_supported(self):
367        """
368        Returns if AP supports setting the visibility (SSID broadcast).
369
370        @return True if supported; False otherwise.
371
372        """
373        return True
374
375
376    def is_band_and_channel_supported(self, band, channel):
377        """
378        Returns if a given band and channel are supported.
379
380        @param band: the band to check if supported
381        @param channel: the channel to check if supported
382
383        @return True if combination is supported; False otherwise.
384
385        """
386        bands = self.get_supported_bands()
387        for current_band in bands:
388            if (current_band['band'] == band and
389                channel in current_band['channels']):
390                return True
391        return False
392
393
394    def is_security_mode_supported(self, security_mode):
395        """
396        Returns if a given security_type is supported.
397
398        Note: The derived class must implement this method.
399
400        @param security_mode: one of the following modes:
401                         self.security_disabled,
402                         self.security_wep,
403                         self.security_wpapsk,
404                         self.security_wpa2psk
405
406        @return True if the security mode is supported; False otherwise.
407
408        """
409        raise NotImplementedError
410
411
412    def navigate_to_page(self, page_number):
413        """
414        Navigates to the page corresponding to the given page number.
415
416        This method performs the translation between a page number and a url to
417        load. This is used internally by apply_settings.
418
419        Note: The derived class must implement this method.
420
421        @param page_number: page number of the page to load
422
423        """
424        raise NotImplementedError
425
426
427    def power_cycle_router_up(self):
428        """Queues the power cycle up command."""
429        self.add_item_to_command_list(self._power_cycle_router_up, (), 1, 0)
430
431
432    def _power_cycle_router_up(self):
433        """Turns the ap off and then back on again."""
434        self.rpm_client.queue_request(self.host_name, 'OFF')
435        self.router_on = False
436        self._power_up_router()
437
438
439    def power_down_router(self):
440        """Queues up the power down command."""
441        self.add_item_to_command_list(self._power_down_router, (), 1, 999)
442
443
444    def _power_down_router(self):
445        """Turns off the power to the ap via the power strip."""
446        self.check_pdu_status()
447        self.rpm_client.queue_request(self.host_name, 'OFF')
448        self.router_on = False
449
450
451    def power_up_router(self):
452        """Queues up the power up command."""
453        self.add_item_to_command_list(self._power_up_router, (), 1, 0)
454
455
456    def _power_up_router(self):
457        """
458        Turns on the power to the ap via the power strip.
459
460        This method returns once it can navigate to a web page of the ap UI.
461
462        """
463        if self.router_on:
464            return
465        self.check_pdu_status()
466        self.rpm_client.queue_request(self.host_name, 'ON')
467        self.establish_driver_connection()
468        # Depending on the response of the webserver for the AP, or lack
469        # there of, the amount of time navigate_to_page and refresh take
470        # is indeterminate.  Give the APs 5 minutes of real time and then
471        # give up.
472        timeout = time.time() + (5 * 60)
473        half_way = time.time() + (2.5 * 60)
474        performed_power_cycle = False
475        while time.time() < timeout:
476            try:
477                logging.info('Attempting to load page')
478                self.navigate_to_page(1)
479                logging.debug('Page navigation complete')
480                self.router_on = True
481                return
482            # Navigate to page may throw a Selemium error or its own
483            # RuntimeError depending on the implementation.  Either way we are
484            # bringing a router back from power off, we need to be patient.
485            except:
486                logging.info('Forcing a page refresh')
487                self.driver.refresh()
488                logging.info('Waiting for router %s to come back up.',
489                             self.name)
490                # Sometime the APs just don't come up right.
491                if not performed_power_cycle and time.time() > half_way:
492                    logging.info('Cannot connect to AP, forcing cycle')
493                    self.rpm_client.queue_request(self.host_name, 'CYCLE')
494                    performed_power_cycle = True
495                    logging.info('Power cycle complete')
496        raise RuntimeError('Unable to load admin page after powering on the '
497                           'router: %s' % self.name)
498
499
500    def save_page(self, page_number):
501        """
502        Saves the given page.
503
504        Note: The derived class must implement this method.
505
506        @param page_number: Page number of the page to save.
507
508        """
509        raise NotImplementedError
510
511
512    def set_using_ap_spec(self, set_ap_spec, power_up=True):
513        """
514        Sets all configurator options.
515
516        @param set_ap_spec: APSpec object
517
518        """
519        if power_up:
520            self.power_up_router()
521        if self.is_visibility_supported():
522            self.set_visibility(set_ap_spec.visible)
523        if (set_ap_spec.security == ap_spec.SECURITY_TYPE_WPAPSK or
524            set_ap_spec.security == ap_spec.SECURITY_TYPE_WPA2PSK):
525            self.set_security_wpapsk(set_ap_spec.security, set_ap_spec.password)
526        else:
527            self.set_security_disabled()
528        self.set_band(set_ap_spec.band)
529        self.set_mode(set_ap_spec.mode)
530        self.set_channel(set_ap_spec.channel)
531
532        # Update ssid
533        raw_ssid = '%s_%s_ch%d_%s' % (
534                self.short_name,
535                ap_spec.mode_string_for_mode(set_ap_spec.mode),
536                set_ap_spec.channel,
537                set_ap_spec.security)
538        self._ssid = raw_ssid.replace(' ', '_').replace('.', '_')[:32]
539        self.set_ssid(self._ssid)
540        self.ap_spec = set_ap_spec
541        self.webdriver_hostname = set_ap_spec.webdriver_hostname
542
543    def set_mode(self, mode, band=None):
544        """
545        Sets the mode.
546
547        Note: The derived class must implement this method.
548
549        @param mode: must be one of the modes listed in __init__()
550        @param band: the band to select
551
552        """
553        raise NotImplementedError
554
555
556    def set_radio(self, enabled=True):
557        """
558        Turns the radio on and off.
559
560        Note: The derived class must implement this method.
561
562        @param enabled: True to turn on the radio; False otherwise
563
564        """
565        raise NotImplementedError
566
567
568    def set_ssid(self, ssid):
569        """
570        Sets the SSID of the wireless network.
571
572        Note: The derived class must implement this method.
573
574        @param ssid: name of the wireless network
575
576        """
577        raise NotImplementedError
578
579
580    def set_channel(self, channel):
581        """
582        Sets the channel of the wireless network.
583
584        Note: The derived class must implement this method.
585
586        @param channel: integer value of the channel
587
588        """
589        raise NotImplementedError
590
591
592    def set_band(self, band):
593        """
594        Sets the band of the wireless network.
595
596        Currently there are only two possible values for band: 2kGHz and 5kGHz.
597        Note: The derived class must implement this method.
598
599        @param band: Constant describing the band type
600
601        """
602        raise NotImplementedError
603
604
605    def set_security_disabled(self):
606        """
607        Disables the security of the wireless network.
608
609        Note: The derived class must implement this method.
610
611        """
612        raise NotImplementedError
613
614
615    def set_security_wep(self, key_value, authentication):
616        """
617        Enabled WEP security for the wireless network.
618
619        Note: The derived class must implement this method.
620
621        @param key_value: encryption key to use
622        @param authentication: one of two supported WEP authentication types:
623                               open or shared.
624        """
625        raise NotImplementedError
626
627
628    def set_security_wpapsk(self, security, shared_key, update_interval=1800):
629        """Enabled WPA using a private security key for the wireless network.
630
631        Note: The derived class must implement this method.
632
633        @param security: Required security for AP configuration
634        @param shared_key: shared encryption key to use
635        @param update_interval: number of seconds to wait before updating
636
637        """
638        raise NotImplementedError
639
640    def set_visibility(self, visible=True):
641        """Set the visibility of the wireless network.
642
643        Note: The derived class must implement this method.
644
645        @param visible: True for visible; False otherwise
646
647        """
648        raise NotImplementedError
649
650
651    def establish_driver_connection(self):
652        """Makes a connection to the webdriver service."""
653        if self.driver_connection_established:
654            return
655        # Load the Auth extension
656
657        webdriver_hostname = self.ap_spec.webdriver_hostname
658        webdriver_ready = self.check_webdriver_ready(webdriver_hostname,
659                                                     self._webdriver_port)
660        webdriver_server = webdriver_ready + '.cros'
661        if webdriver_server is None:
662            raise RuntimeError('Unable to connect to webdriver locally or '
663                               'via the lab service.')
664        extension_path = os.path.join(os.path.dirname(__file__),
665                                      'basic_auth_extension.crx')
666        f = open(extension_path, 'rb')
667        base64_extensions = []
668        base64_ext = (binascii.b2a_base64(f.read()).strip())
669        base64_extensions.append(base64_ext)
670        f.close()
671        webdriver_url = ('http://%s:%d' % (webdriver_server,
672                                           self._webdriver_port))
673        capabilities = {'chromeOptions' : {'extensions' : base64_extensions}}
674        self.driver = webdriver.Remote(webdriver_url, capabilities)
675        self.driver_connection_established = True
676
677
678    def destroy_driver_connection(self):
679        """Breaks the connection to the webdriver service."""
680        try:
681            self.driver.close()
682        except Exception, e:
683            logging.debug('Webdriver is crashed, should be respawned %d',
684                          time.time())
685        finally:
686            self.driver_connection_established = False
687
688
689    def apply_settings(self):
690        """Apply all settings to the access point.
691
692        @param skip_success_validation: Boolean to track if method was
693                                        executed successfully.
694
695        """
696        self.configuration_success = ap_constants.CONFIG_FAIL
697        if len(self._command_list) == 0:
698            return
699
700        # If all we are doing is powering down the router, don't mess with
701        # starting up webdriver.
702        if (len(self._command_list) == 1 and
703            self._command_list[0]['method'] == self._power_down_router):
704            self._command_list[0]['method'](*self._command_list[0]['args'])
705            self._command_list.pop()
706            self.destroy_driver_connection()
707            return
708        self.establish_driver_connection()
709        # Pull items by page and then sort
710        if self.get_number_of_pages() == -1:
711            self.fail(msg='Number of pages is not set.')
712        page_range = range(1, self.get_number_of_pages() + 1)
713        for i in page_range:
714            page_commands = [x for x in self._command_list if x['page'] == i]
715            sorted_page_commands = sorted(page_commands,
716                                          key=lambda k: k['priority'])
717            if sorted_page_commands:
718                first_command = sorted_page_commands[0]['method']
719                # If the first command is bringing the router up or down,
720                # do that before navigating to a URL.
721                if (first_command == self._power_up_router or
722                    first_command == self._power_cycle_router_up or
723                    first_command == self._power_down_router):
724                    direction = 'up'
725                    if first_command == self._power_down_router:
726                        direction = 'down'
727                    logging.info('Powering %s %s', direction, self.name)
728                    first_command(*sorted_page_commands[0]['args'])
729                    sorted_page_commands.pop(0)
730
731                # If the router is off, no point in navigating
732                if not self.router_on:
733                    if len(sorted_page_commands) == 0:
734                        # If all that was requested was to power off
735                        # the router then abort here and do not set the
736                        # configuration_success bit.  The reason is
737                        # because if we failed on the configuration that
738                        # failure should remain since all tests power
739                        # down the AP when they are done.
740                        return
741                    break
742
743                self.navigate_to_page(i)
744                for command in sorted_page_commands:
745                    command['method'](*command['args'])
746                self.save_page(i)
747        self._command_list = []
748        self.configuration_success = ap_constants.CONFIG_SUCCESS
749        self._traceback = None
750        self.destroy_driver_connection()
751
752
753    def get_association_parameters(self):
754        """
755        Creates an AssociationParameters from the configured AP.
756
757        @returns AssociationParameters for the configured AP.
758
759        """
760        security_config = None
761        if self.ap_spec.security in [ap_spec.SECURITY_TYPE_WPAPSK,
762                                     ap_spec.SECURITY_TYPE_WPA2PSK]:
763            # Not all of this is required but doing it just in case.
764            security_config = xmlrpc_security_types.WPAConfig(
765                    psk=self.ap_spec.password,
766                    wpa_mode=xmlrpc_security_types.WPAConfig.MODE_MIXED_WPA,
767                    wpa_ciphers=[xmlrpc_security_types.WPAConfig.CIPHER_CCMP,
768                                 xmlrpc_security_types.WPAConfig.CIPHER_TKIP],
769                    wpa2_ciphers=[xmlrpc_security_types.WPAConfig.CIPHER_CCMP])
770        return xmlrpc_datatypes.AssociationParameters(
771                ssid=self._ssid, security_config=security_config,
772                discovery_timeout=45, association_timeout=30,
773                configuration_timeout=30, is_hidden=not self.ap_spec.visible)
774
775
776    def debug_last_failure(self, outputdir):
777        """
778        Write debug information for last AP_CONFIG_FAIL
779
780        @param outputdir: a string directory path for debug files
781        """
782        logging.error('Traceback:\n %s', self.traceback)
783        self._write_screenshots('config_failure', outputdir)
784        self.clear_screenshot_list()
785
786
787    def debug_full_state(self, outputdir):
788        """
789        Write debug information for full AP state
790
791        @param outputdir: a string directory path for debug files
792        """
793        if self.configuration_success != ap_constants.PDU_FAIL:
794            self._save_all_pages()
795            self._write_screenshots('final_configuration', outputdir)
796            self.clear_screenshot_list()
797        self.reset_command_list()
798
799
800    def store_config_failure(self, trace):
801        """
802        Store configuration failure for latter logging
803
804        @param trace: a string traceback of config exception
805        """
806        self.save_screenshot()
807        self._traceback = trace
808