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