1# Copyright (c) 2012 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
5import ctypes
6import datetime
7import logging
8import multiprocessing
9import os
10import pexpect
11import Queue
12import re
13import threading
14import time
15
16from config import rpm_config
17import dli_urllib
18import rpm_logging_config
19
20import common
21from autotest_lib.client.common_lib import error
22from autotest_lib.client.common_lib.cros import retry
23
24RPM_CALL_TIMEOUT_MINS = rpm_config.getint('RPM_INFRASTRUCTURE',
25                                          'call_timeout_mins')
26SET_POWER_STATE_TIMEOUT_SECONDS = rpm_config.getint(
27        'RPM_INFRASTRUCTURE', 'set_power_state_timeout_seconds')
28PROCESS_TIMEOUT_BUFFER = 30
29
30
31class RPMController(object):
32    """
33    This abstract class implements RPM request queueing and
34    processes queued requests.
35
36    The actual interaction with the RPM device will be implemented
37    by the RPM specific subclasses.
38
39    It assumes that you know the RPM hostname and that the device is on
40    the specified RPM.
41
42    This class also allows support for RPM devices that can be accessed
43    directly or through a hydra serial concentrator device.
44
45    Implementation details:
46    This is an abstract class, subclasses must implement the methods
47    listed here. You must not instantiate this class but should
48    instantiate one of those leaf subclasses. Subclasses should
49    also set TYPE class attribute to indicate device type.
50
51    @var behind_hydra: boolean value to represent whether or not this RPM is
52                        behind a hydra device.
53    @var hostname: hostname for this rpm device.
54    @var is_running_lock: lock used to control access to _running.
55    @var request_queue: queue used to store requested outlet state changes.
56    @var queue_lock: lock used to control access to request_queue.
57    @var _running: boolean value to represent if this controller is currently
58                   looping over queued requests.
59    """
60
61
62    SSH_LOGIN_CMD = ('ssh -l %s -o StrictHostKeyChecking=no '
63                     '-o ConnectTimeout=90 -o UserKnownHostsFile=/dev/null %s')
64    USERNAME_PROMPT = 'Username:'
65    HYRDA_RETRY_SLEEP_SECS = 10
66    HYDRA_MAX_CONNECT_RETRIES = 3
67    LOGOUT_CMD = 'logout'
68    CLI_CMD = 'CLI'
69    CLI_HELD = 'The administrator \[root\] has an active .* session.'
70    CLI_KILL_PREVIOUS = 'cancel'
71    CLI_PROMPT = 'cli>'
72    HYDRA_PROMPT = '#'
73    PORT_STATUS_CMD = 'portStatus'
74    QUIT_CMD = 'quit'
75    SESSION_KILL_CMD_FORMAT = 'administration sessions kill %s'
76    HYDRA_CONN_HELD_MSG_FORMAT = 'is being used'
77    CYCLE_SLEEP_TIME = 5
78
79    # Global Variables that will likely be changed by subclasses.
80    DEVICE_PROMPT = '$'
81    PASSWORD_PROMPT = 'Password:'
82    # The state change command can be any string format but must accept 2 vars:
83    # state followed by device/Plug name.
84    SET_STATE_CMD = '%s %s'
85    SUCCESS_MSG = None # Some RPM's may not return a success msg.
86
87    NEW_STATE_ON = 'ON'
88    NEW_STATE_OFF = 'OFF'
89    NEW_STATE_CYCLE = 'CYCLE'
90    TYPE = 'Should set TYPE in subclass.'
91
92
93    def __init__(self, rpm_hostname, hydra_hostname=None):
94        """
95        RPMController Constructor.
96        To be called by subclasses.
97
98        @param rpm_hostname: hostname of rpm device to be controlled.
99        """
100        self._dns_zone = rpm_config.get('CROS', 'dns_zone')
101        self.hostname = rpm_hostname
102        self.request_queue = Queue.Queue()
103        self._running = False
104        self.is_running_lock = threading.Lock()
105        # If a hydra name is provided by the subclass then we know we are
106        # talking to an rpm behind a hydra device.
107        self.hydra_hostname = hydra_hostname if hydra_hostname else None
108        self.behind_hydra = hydra_hostname is not None
109
110
111    def _start_processing_requests(self):
112        """
113        Check if there is a thread processing requests.
114        If not start one.
115        """
116        with self.is_running_lock:
117            if not self._running:
118                self._running = True
119                self._running_thread = threading.Thread(target=self._run)
120                self._running_thread.start()
121
122
123    def _stop_processing_requests(self):
124        """
125        Called if the request request_queue is empty.
126        Set running status to false.
127        """
128        with self.is_running_lock:
129            logging.debug('Request queue is empty. RPM Controller for %s'
130                          ' is terminating.', self.hostname)
131            self._running = False
132        if not self.request_queue.empty():
133            # This can occur if an item was pushed into the queue after we
134            # exited the while-check and before the _stop_processing_requests
135            # call was made. Therefore we need to start processing again.
136            self._start_processing_requests()
137
138
139    def _run(self):
140        """
141        Processes all queued up requests for this RPM Controller.
142        Callers should first request_queue up atleast one request and if this
143        RPM Controller is not running then call run.
144
145        Caller can either simply call run but then they will be blocked or
146        can instantiate a new thread to process all queued up requests.
147        For example:
148          threading.Thread(target=rpm_controller.run).start()
149
150        Requests are in the format of:
151          [powerunit_info, new_state, condition_var, result]
152        Run will set the result with the correct value.
153        """
154        while not self.request_queue.empty():
155            try:
156                result = multiprocessing.Value(ctypes.c_bool, False)
157                request = self.request_queue.get()
158                device_hostname = request['powerunit_info'].device_hostname
159                if (datetime.datetime.utcnow() > (request['start_time'] +
160                        datetime.timedelta(minutes=RPM_CALL_TIMEOUT_MINS))):
161                    logging.error('The request was waited for too long to be '
162                                  "processed. It is timed out and won't be "
163                                  'processed.')
164                    request['result_queue'].put(False)
165                    continue
166
167                is_timeout = multiprocessing.Value(ctypes.c_bool, False)
168                process = multiprocessing.Process(target=self._process_request,
169                                                  args=(request, result,
170                                                        is_timeout))
171                process.start()
172                process.join(SET_POWER_STATE_TIMEOUT_SECONDS +
173                             PROCESS_TIMEOUT_BUFFER)
174                if process.is_alive():
175                    logging.debug('%s: process (%s) still running, will be '
176                                  'terminated!', device_hostname, process.pid)
177                    process.terminate()
178                    is_timeout.value = True
179
180                if is_timeout.value:
181                    raise error.TimeoutException(
182                            'Attempt to set power state is timed out after %s '
183                            'seconds.' % SET_POWER_STATE_TIMEOUT_SECONDS)
184                if not result.value:
185                    logging.error('Request to change %s to state %s failed.',
186                                  device_hostname, request['new_state'])
187            except Exception as e:
188                logging.error('Request to change %s to state %s failed: '
189                              'Raised exception: %s', device_hostname,
190                              request['new_state'], e)
191                result.value = False
192
193            # Put result inside the result Queue to allow the caller to resume.
194            request['result_queue'].put(result.value)
195        self._stop_processing_requests()
196
197
198    def _process_request(self, request, result, is_timeout):
199        """Process the request to change a device's outlet state.
200
201        The call of set_power_state is made in a new running process. If it
202        takes longer than SET_POWER_STATE_TIMEOUT_SECONDS, the request will be
203        timed out.
204
205        @param request: A request to change a device's outlet state.
206        @param result: A Value object passed to the new process for the caller
207                       thread to retrieve the result.
208        @param is_timeout: A Value object passed to the new process for the
209                           caller thread to retrieve the information about if
210                           the set_power_state call timed out.
211        """
212        try:
213            logging.getLogger().handlers = []
214            is_timeout_value, result_value = retry.timeout(
215                     rpm_logging_config.set_up_logging_to_server,
216                     timeout_sec=10)
217            if is_timeout_value:
218                raise Exception('Setup local log server handler timed out.')
219        except Exception as e:
220            # Fail over to log to a new file.
221            LOG_FILENAME_FORMAT = rpm_config.get('GENERAL',
222                                                 'dispatcher_logname_format')
223            log_filename_format = LOG_FILENAME_FORMAT.replace(
224                    'dispatcher', 'controller_%d' % os.getpid())
225            logging.getLogger().handlers = []
226            rpm_logging_config.set_up_logging_to_file(
227                    log_dir='./logs',
228                    log_filename_format=log_filename_format,
229                    use_log_server=False)
230            logging.info('Failed to set up logging through log server: %s', e)
231        kwargs = {'powerunit_info':request['powerunit_info'],
232                  'new_state':request['new_state']}
233        try:
234            is_timeout_value, result_value = retry.timeout(
235                    self.set_power_state,
236                    args=(),
237                    kwargs=kwargs,
238                    timeout_sec=SET_POWER_STATE_TIMEOUT_SECONDS)
239            result.value = result_value
240            is_timeout.value = is_timeout_value
241        except Exception as e:
242            # This method runs in a subprocess. Must log the exception,
243            # otherwise exceptions raised in set_power_state just get lost.
244            # Need to convert e to a str type, because our logging server
245            # code doesn't handle the conversion very well.
246            logging.error('Request to change %s to state %s failed: '
247                          'Raised exception: %s',
248                          request['powerunit_info'].device_hostname,
249                          request['new_state'], str(e))
250            raise e
251
252
253    def queue_request(self, powerunit_info, new_state):
254        """
255        Queues up a requested state change for a device's outlet.
256
257        Requests are in the format of:
258          [powerunit_info, new_state, condition_var, result]
259        Run will set the result with the correct value.
260
261        @param powerunit_info: And PowerUnitInfo instance.
262        @param new_state: ON/OFF/CYCLE - state or action we want to perform on
263                          the outlet.
264        """
265        request = {}
266        request['powerunit_info'] = powerunit_info
267        request['new_state'] = new_state
268        request['start_time'] = datetime.datetime.utcnow()
269        # Reserve a spot for the result to be stored.
270        request['result_queue'] = Queue.Queue()
271        # Place in request_queue
272        self.request_queue.put(request)
273        self._start_processing_requests()
274        # Block until the request is processed.
275        result = request['result_queue'].get(block=True)
276        return result
277
278
279    def _kill_previous_connection(self):
280        """
281        In case the port to the RPM through the hydra serial concentrator is in
282        use, terminate the previous connection so we can log into the RPM.
283
284        It logs into the hydra serial concentrator over ssh, launches the CLI
285        command, gets the port number and then kills the current session.
286        """
287        ssh = self._authenticate_with_hydra(admin_override=True)
288        if not ssh:
289            return
290        ssh.expect(RPMController.PASSWORD_PROMPT, timeout=60)
291        ssh.sendline(rpm_config.get('HYDRA', 'admin_password'))
292        ssh.expect(RPMController.HYDRA_PROMPT)
293        ssh.sendline(RPMController.CLI_CMD)
294        cli_prompt_re = re.compile(RPMController.CLI_PROMPT)
295        cli_held_re = re.compile(RPMController.CLI_HELD)
296        response = ssh.expect_list([cli_prompt_re, cli_held_re], timeout=60)
297        if response == 1:
298            # Need to kill the previous adminstator's session.
299            logging.error("Need to disconnect previous administrator's CLI "
300                          "session to release the connection to RPM device %s.",
301                          self.hostname)
302            ssh.sendline(RPMController.CLI_KILL_PREVIOUS)
303            ssh.expect(RPMController.CLI_PROMPT)
304        ssh.sendline(RPMController.PORT_STATUS_CMD)
305        ssh.expect(': %s' % self.hostname)
306        ports_status = ssh.before
307        port_number = ports_status.split(' ')[-1]
308        ssh.expect(RPMController.CLI_PROMPT)
309        ssh.sendline(RPMController.SESSION_KILL_CMD_FORMAT % port_number)
310        ssh.expect(RPMController.CLI_PROMPT)
311        self._logout(ssh, admin_logout=True)
312
313
314    def _hydra_login(self, ssh):
315        """
316        Perform the extra steps required to log into a hydra serial
317        concentrator.
318
319        @param ssh: pexpect.spawn object used to communicate with the hydra
320                    serial concentrator.
321
322        @return: True if the login procedure is successful. False if an error
323                 occurred. The most common case would be if another user is
324                 logged into the device.
325        """
326        try:
327            response = ssh.expect_list(
328                    [re.compile(RPMController.PASSWORD_PROMPT),
329                     re.compile(RPMController.HYDRA_CONN_HELD_MSG_FORMAT)],
330                    timeout=15)
331        except pexpect.TIMEOUT:
332            # If there was a timeout, this ssh tunnel could be set up to
333            # not require the hydra password.
334            ssh.sendline('')
335            try:
336                ssh.expect(re.compile(RPMController.USERNAME_PROMPT))
337                logging.debug('Connected to rpm through hydra. Logging in.')
338                return True
339            except pexpect.ExceptionPexpect:
340                return False
341        if response == 0:
342            try:
343                ssh.sendline(rpm_config.get('HYDRA','password'))
344                ssh.sendline('')
345                response = ssh.expect_list(
346                        [re.compile(RPMController.USERNAME_PROMPT),
347                         re.compile(RPMController.HYDRA_CONN_HELD_MSG_FORMAT)],
348                        timeout=60)
349            except pexpect.EOF:
350                # Did not receive any of the expect responses, retry.
351                return False
352            except pexpect.TIMEOUT:
353                logging.debug('Timeout occurred logging in to hydra.')
354                return False
355        # Send the username that the subclass will have set in its
356        # construction.
357        if response == 1:
358            logging.debug('SSH Terminal most likely serving another'
359                          ' connection, retrying.')
360            # Kill the connection for the next connection attempt.
361            try:
362                self._kill_previous_connection()
363            except pexpect.ExceptionPexpect:
364                logging.error('Failed to disconnect previous connection, '
365                              'retrying.')
366                raise
367            return False
368        logging.debug('Connected to rpm through hydra. Logging in.')
369        return True
370
371
372    def _authenticate_with_hydra(self, admin_override=False):
373        """
374        Some RPM's are behind a hydra serial concentrator and require their ssh
375        connection to be tunneled through this device. This can fail if another
376        user is logged in; therefore this will retry multiple times.
377
378        This function also allows us to authenticate directly to the
379        administrator interface of the hydra device.
380
381        @param admin_override: Set to True if we are trying to access the
382                               administrator interface rather than tunnel
383                               through to the RPM.
384
385        @return: The connected pexpect.spawn instance if the login procedure is
386                 successful. None if an error occurred. The most common case
387                 would be if another user is logged into the device.
388        """
389        if admin_override:
390            username = rpm_config.get('HYDRA', 'admin_username')
391        else:
392            username = '%s:%s' % (rpm_config.get('HYDRA','username'),
393                                  self.hostname)
394        cmd = RPMController.SSH_LOGIN_CMD % (username, self.hydra_hostname)
395        num_attempts = 0
396        while num_attempts < RPMController.HYDRA_MAX_CONNECT_RETRIES:
397            try:
398                ssh = pexpect.spawn(cmd)
399            except pexpect.ExceptionPexpect:
400                return None
401            if admin_override:
402                return ssh
403            if self._hydra_login(ssh):
404                return ssh
405            # Authenticating with hydra failed. Sleep then retry.
406            time.sleep(RPMController.HYRDA_RETRY_SLEEP_SECS)
407            num_attempts += 1
408        logging.error('Failed to connect to the hydra serial concentrator after'
409                      ' %d attempts.', RPMController.HYDRA_MAX_CONNECT_RETRIES)
410        return None
411
412
413    def _login(self):
414        """
415        Log in into the RPM Device.
416
417        The login process should be able to connect to the device whether or not
418        it is behind a hydra serial concentrator.
419
420        @return: ssh - a pexpect.spawn instance if the connection was successful
421                 or None if it was not.
422        """
423        if self.behind_hydra:
424            # Tunnel the connection through the hydra.
425            ssh = self._authenticate_with_hydra()
426            if not ssh:
427                return None
428            ssh.sendline(self._username)
429        else:
430            # Connect directly to the RPM over SSH.
431            hostname = '%s.%s' % (self.hostname, self._dns_zone)
432            cmd = RPMController.SSH_LOGIN_CMD % (self._username, hostname)
433            try:
434                ssh = pexpect.spawn(cmd)
435            except pexpect.ExceptionPexpect:
436                return None
437        # Wait for the password prompt
438        try:
439            ssh.expect(self.PASSWORD_PROMPT, timeout=60)
440            ssh.sendline(self._password)
441            ssh.expect(self.DEVICE_PROMPT, timeout=60)
442        except pexpect.ExceptionPexpect:
443            return None
444        return ssh
445
446
447    def _logout(self, ssh, admin_logout=False):
448        """
449        Log out of the RPM device.
450
451        Send the device specific logout command and if the connection is through
452        a hydra serial concentrator, kill the ssh connection.
453
454        @param admin_logout: Set to True if we are trying to logout of the
455                             administrator interface of a hydra serial
456                             concentrator, rather than an RPM.
457        @param ssh: pexpect.spawn instance to use to send the logout command.
458        """
459        if admin_logout:
460            ssh.sendline(RPMController.QUIT_CMD)
461            ssh.expect(RPMController.HYDRA_PROMPT)
462        ssh.sendline(self.LOGOUT_CMD)
463        if self.behind_hydra and not admin_logout:
464            # Terminate the hydra session.
465            ssh.sendline('~.')
466            # Wait a bit so hydra disconnects completely. Launching another
467            # request immediately can cause a timeout.
468            time.sleep(5)
469
470
471    def set_power_state(self, powerunit_info, new_state):
472        """
473        Set the state of the dut's outlet on this RPM.
474
475        For ssh based devices, this will create the connection either directly
476        or through a hydra tunnel and call the underlying _change_state function
477        to be implemented by the subclass device.
478
479        For non-ssh based devices, this method should be overloaded with the
480        proper connection and state change code. And the subclass will handle
481        accessing the RPM devices.
482
483        @param powerunit_info: An instance of PowerUnitInfo.
484        @param new_state: ON/OFF/CYCLE - state or action we want to perform on
485                          the outlet.
486
487        @return: True if the attempt to change power state was successful,
488                 False otherwise.
489        """
490        ssh = self._login()
491        if not ssh:
492            return False
493        if new_state == self.NEW_STATE_CYCLE:
494            logging.debug('Beginning Power Cycle for device: %s',
495                          powerunit_info.device_hostname)
496            result = self._change_state(powerunit_info, self.NEW_STATE_OFF, ssh)
497            if not result:
498                return result
499            time.sleep(RPMController.CYCLE_SLEEP_TIME)
500            result = self._change_state(powerunit_info, self.NEW_STATE_ON, ssh)
501        else:
502            # Try to change the state of the device's power outlet.
503            result = self._change_state(powerunit_info, new_state, ssh)
504
505        # Terminate hydra connection if necessary.
506        self._logout(ssh)
507        ssh.close(force=True)
508        return result
509
510
511    def _change_state(self, powerunit_info, new_state, ssh):
512        """
513        Perform the actual state change operation.
514
515        Once we have established communication with the RPM this method is
516        responsible for changing the state of the RPM outlet.
517
518        @param powerunit_info: An instance of PowerUnitInfo.
519        @param new_state: ON/OFF - state or action we want to perform on
520                          the outlet.
521        @param ssh: The ssh connection used to execute the state change commands
522                    on the RPM device.
523
524        @return: True if the attempt to change power state was successful,
525                 False otherwise.
526        """
527        outlet = powerunit_info.outlet
528        device_hostname = powerunit_info.device_hostname
529        if not outlet:
530            logging.error('Request to change outlet for device: %s to new '
531                          'state %s failed: outlet is unknown, please '
532                          'make sure POWERUNIT_OUTLET exist in the host\'s '
533                          'attributes in afe.', device_hostname, new_state)
534        ssh.sendline(self.SET_STATE_CMD % (new_state, outlet))
535        if self.SUCCESS_MSG:
536            # If this RPM device returns a success message check for it before
537            # continuing.
538            try:
539                ssh.expect(self.SUCCESS_MSG, timeout=60)
540            except pexpect.ExceptionPexpect:
541                logging.error('Request to change outlet for device: %s to new '
542                              'state %s failed.', device_hostname, new_state)
543                return False
544        logging.debug('Outlet for device: %s set to %s', device_hostname,
545                      new_state)
546        return True
547
548
549    def type(self):
550        """
551        Get the type of RPM device we are interacting with.
552        Class attribute TYPE should be set by the subclasses.
553
554        @return: string representation of RPM device type.
555        """
556        return self.TYPE
557
558
559class SentryRPMController(RPMController):
560    """
561    This class implements power control for Sentry Switched CDU
562    http://www.servertech.com/products/switched-pdus/
563
564    Example usage:
565      rpm = SentrySwitchedCDU('chromeos-rack1-rpm1')
566      rpm.queue_request('chromeos-rack1-host1', 'ON')
567
568    @var _username: username used to access device.
569    @var _password: password used to access device.
570    """
571
572
573    DEVICE_PROMPT = 'Switched CDU:'
574    SET_STATE_CMD = '%s %s'
575    SUCCESS_MSG = 'Command successful'
576    NUM_OF_OUTLETS = 17
577    TYPE = 'Sentry'
578
579
580    def __init__(self, hostname, hydra_hostname=None):
581        super(SentryRPMController, self).__init__(hostname, hydra_hostname)
582        self._username = rpm_config.get('SENTRY', 'username')
583        self._password = rpm_config.get('SENTRY', 'password')
584
585
586    def _setup_test_user(self, ssh):
587        """Configure the test user for the RPM
588
589        @param ssh: Pexpect object to use to configure the RPM.
590        """
591        # Create and configure the testing user profile.
592        testing_user = rpm_config.get('SENTRY','testing_user')
593        testing_password = rpm_config.get('SENTRY','testing_password')
594        ssh.sendline('create user %s' % testing_user)
595        response = ssh.expect_list([re.compile('not unique'),
596                                    re.compile(self.PASSWORD_PROMPT)])
597        if not response:
598            return
599        # Testing user is not set up yet.
600        ssh.sendline(testing_password)
601        ssh.expect('Verify Password:')
602        ssh.sendline(testing_password)
603        ssh.expect(self.SUCCESS_MSG)
604        ssh.expect(self.DEVICE_PROMPT)
605        ssh.sendline('add outlettouser all %s' % testing_user)
606        ssh.expect(self.SUCCESS_MSG)
607        ssh.expect(self.DEVICE_PROMPT)
608
609
610    def _clear_outlet_names(self, ssh):
611        """
612        Before setting the outlet names, we need to clear out all the old
613        names so there are no conflicts. For example trying to assign outlet
614        2 a name already assigned to outlet 9.
615        """
616        for outlet in range(1, self.NUM_OF_OUTLETS):
617            outlet_name = 'Outlet_%d' % outlet
618            ssh.sendline(self.SET_OUTLET_NAME_CMD % (outlet, outlet_name))
619            ssh.expect(self.SUCCESS_MSG)
620            ssh.expect(self.DEVICE_PROMPT)
621
622
623    def setup(self, outlet_naming_map):
624        """
625        Configure the RPM by adding the test user and setting up the outlet
626        names.
627
628        Note the rpm infrastructure does not rely on the outlet name to map a
629        device to its outlet any more. We keep this method in case there is
630        a need to label outlets for other reasons. We may deprecate
631        this method if it has been proved the outlet names will not be used
632        in any scenario.
633
634        @param outlet_naming_map: Dictionary used to map the outlet numbers to
635                                  host names. Keys must be ints. And names are
636                                  in the format of 'hostX'.
637
638        @return: True if setup completed successfully, False otherwise.
639        """
640        ssh = self._login()
641        if not ssh:
642            logging.error('Could not connect to %s.', self.hostname)
643            return False
644        try:
645            self._setup_test_user(ssh)
646            # Set up the outlet names.
647            # Hosts have the same name format as the RPM hostname except they
648            # end in hostX instead of rpmX.
649            dut_name_format = re.sub('-rpm[0-9]*', '', self.hostname)
650            if self.behind_hydra:
651                # Remove "chromeosX" from DUTs behind the hydra due to a length
652                # constraint on the names we can store inside the RPM.
653                dut_name_format = re.sub('chromeos[0-9]*-', '', dut_name_format)
654            dut_name_format = dut_name_format + '-%s'
655            self._clear_outlet_names(ssh)
656            for outlet, name in outlet_naming_map.items():
657                dut_name = dut_name_format % name
658                ssh.sendline(self.SET_OUTLET_NAME_CMD % (outlet, dut_name))
659                ssh.expect(self.SUCCESS_MSG)
660                ssh.expect(self.DEVICE_PROMPT)
661        except pexpect.ExceptionPexpect as e:
662            logging.error('Setup failed. %s', e)
663            return False
664        finally:
665            self._logout(ssh)
666        return True
667
668
669class WebPoweredRPMController(RPMController):
670    """
671    This class implements RPMController for the Web Powered units
672    produced by Digital Loggers Inc.
673
674    @var _rpm: dli_urllib.Powerswitch instance used to interact with RPM.
675    """
676
677
678    TYPE = 'Webpowered'
679
680
681    def __init__(self, hostname, powerswitch=None):
682        username = rpm_config.get('WEBPOWERED', 'username')
683        password = rpm_config.get('WEBPOWERED', 'password')
684        # Call the constructor in RPMController. However since this is a web
685        # accessible device, there should not be a need to tunnel through a
686        # hydra serial concentrator.
687        super(WebPoweredRPMController, self).__init__(hostname)
688        self.hostname = '%s.%s' % (self.hostname, self._dns_zone)
689        if not powerswitch:
690            self._rpm = dli_urllib.Powerswitch(hostname=self.hostname,
691                                               userid=username,
692                                               password=password)
693        else:
694            # Should only be used in unit_testing
695            self._rpm = powerswitch
696
697
698    def _get_outlet_state(self, outlet):
699        """
700        Look up the state for a given outlet on the RPM.
701
702        @param outlet: the outlet to look up.
703
704        @return state: the outlet's current state.
705        """
706        status_list = self._rpm.statuslist()
707        for outlet_name, hostname, state in status_list:
708            if outlet_name == outlet:
709                return state
710        return None
711
712
713    def set_power_state(self, powerunit_info, new_state):
714        """
715        Since this does not utilize SSH in any manner, this will overload the
716        set_power_state in RPMController and completes all steps of changing
717        the device's outlet state.
718        """
719        device_hostname = powerunit_info.device_hostname
720        outlet = powerunit_info.outlet
721        if not outlet:
722            logging.error('Request to change outlet for device %s to '
723                          'new state %s failed: outlet is unknown. Make sure '
724                          'POWERUNIT_OUTLET exists in the host\'s '
725                          'attributes in afe' , device_hostname, new_state)
726            return False
727        state = self._get_outlet_state(outlet)
728        expected_state = new_state
729        if new_state == self.NEW_STATE_CYCLE:
730            logging.debug('Beginning Power Cycle for device: %s',
731                          device_hostname)
732            self._rpm.off(outlet)
733            logging.debug('Outlet for device: %s set to OFF', device_hostname)
734            # Pause for 5 seconds before restoring power.
735            time.sleep(RPMController.CYCLE_SLEEP_TIME)
736            self._rpm.on(outlet)
737            logging.debug('Outlet for device: %s set to ON', device_hostname)
738            expected_state = self.NEW_STATE_ON
739        if new_state == self.NEW_STATE_OFF:
740            self._rpm.off(outlet)
741            logging.debug('Outlet for device: %s set to OFF', device_hostname)
742        if new_state == self.NEW_STATE_ON:
743            self._rpm.on(outlet)
744            logging.debug('Outlet for device: %s set to ON', device_hostname)
745        # Lookup the final state of the outlet
746        return self._is_plug_state(powerunit_info, expected_state)
747
748
749    def _is_plug_state(self, powerunit_info, expected_state):
750        state = self._get_outlet_state(powerunit_info.outlet)
751        if expected_state not in state:
752            logging.error('Outlet for device: %s did not change to new state'
753                          ' %s', powerunit_info.device_hostname, expected_state)
754            return False
755        return True
756
757
758class CiscoPOEController(RPMController):
759    """
760    This class implements power control for Cisco POE switch.
761
762    Example usage:
763      poe = CiscoPOEController('chromeos1-poe-switch1')
764      poe.queue_request('chromeos1-rack5-host12-servo', 'ON')
765    """
766
767
768    SSH_LOGIN_CMD = ('ssh -o StrictHostKeyChecking=no '
769                     '-o UserKnownHostsFile=/dev/null %s')
770    POE_USERNAME_PROMPT = 'User Name:'
771    POE_PROMPT = '%s#'
772    EXIT_CMD = 'exit'
773    END_CMD = 'end'
774    CONFIG = 'configure terminal'
775    CONFIG_PROMPT = '%s\(config\)#'
776    CONFIG_IF = 'interface %s'
777    CONFIG_IF_PROMPT = '%s\(config-if\)#'
778    SET_STATE_ON = 'power inline auto'
779    SET_STATE_OFF = 'power inline never'
780    CHECK_INTERFACE_STATE = 'show interface status %s'
781    INTERFACE_STATE_MSG = 'Port\s+.*%s(\s+(\S+)){6,6}'
782    CHECK_STATE_TIMEOUT = 60
783    CMD_TIMEOUT = 30
784    LOGIN_TIMEOUT = 60
785    PORT_UP = 'Up'
786    PORT_DOWN = 'Down'
787    TYPE = 'CiscoPOE'
788
789
790    def __init__(self, hostname):
791        """
792        Initialize controller class for a Cisco POE switch.
793
794        @param hostname: the Cisco POE switch host name.
795        """
796        super(CiscoPOEController, self).__init__(hostname)
797        self._username = rpm_config.get('CiscoPOE', 'username')
798        self._password = rpm_config.get('CiscoPOE', 'password')
799        # For a switch, e.g. 'chromeos2-poe-switch8',
800        # the device prompt looks like 'chromeos2-poe-sw8#'.
801        short_hostname = self.hostname.replace('switch', 'sw')
802        self.poe_prompt = self.POE_PROMPT % short_hostname
803        self.config_prompt = self.CONFIG_PROMPT % short_hostname
804        self.config_if_prompt = self.CONFIG_IF_PROMPT % short_hostname
805
806
807    def _login(self):
808        """
809        Log in into the Cisco POE switch.
810
811        Overload _login in RPMController, as it always prompts for a user name.
812
813        @return: ssh - a pexpect.spawn instance if the connection was successful
814                 or None if it was not.
815        """
816        hostname = '%s.%s' % (self.hostname, self._dns_zone)
817        cmd = self.SSH_LOGIN_CMD % (hostname)
818        try:
819            ssh = pexpect.spawn(cmd)
820        except pexpect.ExceptionPexpect:
821            logging.error('Could not connect to switch %s', hostname)
822            return None
823        # Wait for the username and password prompt.
824        try:
825            ssh.expect(self.POE_USERNAME_PROMPT, timeout=self.LOGIN_TIMEOUT)
826            ssh.sendline(self._username)
827            ssh.expect(self.PASSWORD_PROMPT, timeout=self.LOGIN_TIMEOUT)
828            ssh.sendline(self._password)
829            ssh.expect(self.poe_prompt, timeout=self.LOGIN_TIMEOUT)
830        except pexpect.ExceptionPexpect:
831            logging.error('Could not log into switch %s', hostname)
832            return None
833        return ssh
834
835
836    def _enter_configuration_terminal(self, interface, ssh):
837        """
838        Enter configuration terminal of |interface|.
839
840        This function expects that we've already logged into the switch
841        and the ssh is prompting the switch name. The work flow is
842            chromeos1-poe-sw1#
843            chromeos1-poe-sw1#configure terminal
844            chromeos1-poe-sw1(config)#interface fa36
845            chromeos1-poe-sw1(config-if)#
846        On success, the function exits with 'config-if' prompt.
847        On failure, the function exits with device prompt,
848        e.g. 'chromeos1-poe-sw1#' in the above case.
849
850        @param interface: the name of the interface.
851        @param ssh: pexpect.spawn instance to use.
852
853        @return: True on success otherwise False.
854        """
855        try:
856            # Enter configure terminal.
857            ssh.sendline(self.CONFIG)
858            ssh.expect(self.config_prompt, timeout=self.CMD_TIMEOUT)
859            # Enter configure terminal of the interface.
860            ssh.sendline(self.CONFIG_IF % interface)
861            ssh.expect(self.config_if_prompt, timeout=self.CMD_TIMEOUT)
862            return True
863        except pexpect.ExceptionPexpect, e:
864            ssh.sendline(self.END_CMD)
865            logging.exception(e)
866        return False
867
868
869    def _exit_configuration_terminal(self, ssh):
870        """
871        Exit interface configuration terminal.
872
873        On success, the function exits with device prompt,
874        e.g. 'chromeos1-poe-sw1#' in the above case.
875        On failure, the function exists with 'config-if' prompt.
876
877        @param ssh: pexpect.spawn instance to use.
878
879        @return: True on success otherwise False.
880        """
881        try:
882            ssh.sendline(self.END_CMD)
883            ssh.expect(self.poe_prompt, timeout=self.CMD_TIMEOUT)
884            return True
885        except pexpect.ExceptionPexpect, e:
886            logging.exception(e)
887        return False
888
889
890    def _verify_state(self, interface, expected_state, ssh):
891        """
892        Check whehter the current state of |interface| matches expected state.
893
894        This function tries to check the state of |interface| multiple
895        times until its state matches the expected state or time is out.
896
897        After the command of changing state has been executed,
898        the state of an interface doesn't always change immediately to
899        the expected state but requires some time. As such, we need
900        a retry logic here.
901
902        @param interface: the name of the interface.
903        @param expect_state: the expected state, 'ON' or 'OFF'
904        @param ssh: pexpect.spawn instance to use.
905
906        @return: True if the state of |interface| swiches to |expected_state|,
907                 otherwise False.
908        """
909        expected_state = (self.PORT_UP if expected_state == self.NEW_STATE_ON
910                          else self.PORT_DOWN)
911        try:
912            start = time.time()
913            while((time.time() - start) < self.CHECK_STATE_TIMEOUT):
914                ssh.sendline(self.CHECK_INTERFACE_STATE % interface)
915                state_matcher = '.*'.join([self.INTERFACE_STATE_MSG % interface,
916                                           self.poe_prompt])
917                ssh.expect(state_matcher, timeout=self.CMD_TIMEOUT)
918                state = ssh.match.group(2)
919                if state == expected_state:
920                    return True
921        except pexpect.ExceptionPexpect, e:
922            logging.exception(e)
923        return False
924
925
926    def _logout(self, ssh, admin_logout=False):
927        """
928        Log out of the Cisco POE switch after changing state.
929
930        Overload _logout in RPMController.
931
932        @param admin_logout: ignored by this method.
933        @param ssh: pexpect.spawn instance to use to send the logout command.
934        """
935        ssh.sendline(self.EXIT_CMD)
936
937
938    def _change_state(self, powerunit_info, new_state, ssh):
939        """
940        Perform the actual state change operation.
941
942        Overload _change_state in RPMController.
943
944        @param powerunit_info: An PowerUnitInfo instance.
945        @param new_state: ON/OFF - state or action we want to perform on
946                          the outlet.
947        @param ssh: The ssh connection used to execute the state change commands
948                    on the POE switch.
949
950        @return: True if the attempt to change power state was successful,
951                 False otherwise.
952        """
953        interface = powerunit_info.outlet
954        device_hostname = powerunit_info.device_hostname
955        if not interface:
956            logging.error('Could not change state: the interface on %s for %s '
957                          'was not given.', self.hostname, device_hostname)
958            return False
959
960        # Enter configuration terminal.
961        if not self._enter_configuration_terminal(interface, ssh):
962            logging.error('Could not enter configuration terminal for %s',
963                          interface)
964            return False
965        # Change the state.
966        if new_state == self.NEW_STATE_ON:
967            ssh.sendline(self.SET_STATE_ON)
968        elif new_state == self.NEW_STATE_OFF:
969            ssh.sendline(self.SET_STATE_OFF)
970        else:
971            logging.error('Unknown state request: %s', new_state)
972            return False
973        # Exit configuraiton terminal.
974        if not self._exit_configuration_terminal(ssh):
975            logging.error('Skipping verifying outlet state for device: %s, '
976                          'because could not exit configuration terminal.',
977                          device_hostname)
978            return False
979        # Verify if the state has changed successfully.
980        if not self._verify_state(interface, new_state, ssh):
981            logging.error('Could not verify state on interface %s', interface)
982            return False
983
984        logging.debug('Outlet for device: %s set to %s',
985                      device_hostname, new_state)
986        return True
987