class CreateSession(): def __init__(self, host, username, password, logger=None): self.cli = CLI() self.logger = logger self.mode = CommandMode(r'$') # for example r'%\s*$' enable_action_map = { "[Pp]assword for {}".format(username): lambda session, logger: session.send_line(password, logger) } self.elevated_mode = CommandMode(r'(?:(?!\)).)#\s*$', enter_command='sudo su', exit_command='exit', enter_action_map=enable_action_map) self.mode.add_child_node(self.elevated_mode) self.elevated_mode.add_parent_mode(self.mode) self.session_types = [ SSHSession(host=host, username=username, password=password) ] self.session = self.cli.get_session(command_mode=self.elevated_mode, new_sessions=self.session_types) def send_terminal_command(self, command, password=None): outp = [] out = None with self.session as my_session: if isinstance(command, list): for single_command in command: if password: single_command = '{command}'.format( command=single_command, password=password) # single_command = 'echo {password} | sudo -S sh -c "{command}"'.format(command=single_command, # password=password) self.logger.info( 'sending command {}'.format(single_command)) current_outp = my_session.send_command(single_command) outp.append(current_outp) self.logger.info('got output {}'.format(current_outp)) out = '\n'.join(outp) else: if password: command = '{command}'.format(command=command) out = my_session.send_command(command) return out
class TestCommandMode(TestCase): def setUp(self): self._prompt = Mock() self._enter_command = Mock() self._enter_action_map = Mock() self._enter_error_map = Mock() self._exit_command = Mock() self._exit_action_map = Mock() self._exit_error_map = Mock() self._enter_actions = Mock() self._session = Mock() self._command_mode = CommandMode(self._prompt, enter_command=self._enter_command, enter_action_map=self._enter_action_map, enter_error_map=self._enter_error_map, exit_command=self._exit_command, exit_action_map=self._exit_action_map, exit_error_map=self._exit_error_map, enter_actions=self._enter_actions) self._logger = Mock() def test_init(self): attributes = ['_enter_error_map', 'prompt', '_enter_command', '_exit_error_map', '_exit_command', 'child_nodes', '_exit_action_map', '_enter_actions', 'parent_node', '_enter_action_map'] self.assertTrue(len(set(attributes).difference(set(self._command_mode.__dict__.keys()))) == 0) def test_add_parent_mode(self): mode = Mock() self._command_mode.add_parent_mode(mode) mode.add_child_node.assert_called_once_with(self._command_mode) def test_step_up_send_command(self): cli_service = Mock() self._command_mode.step_up(cli_service) cli_service.send_command.assert_called_once_with(self._enter_command, expected_string=self._prompt, action_map=self._enter_action_map, error_map=self._enter_error_map) def test_step_up_set_command_mode(self): cli_service = Mock() self._command_mode.step_up(cli_service) self.assertTrue(cli_service.command_mode == self._command_mode) def test_step_up_call_enter_actions(self): cli_service = Mock() enter_actions = Mock() self._command_mode.enter_actions = enter_actions self._command_mode.step_up(cli_service) enter_actions.assert_called_once_with(cli_service) def test_step_down_sent_command(self): cli_service = Mock() parent_prompt = Mock() parent_node = Mock() parent_node.prompt = parent_prompt self._command_mode.parent_node = parent_node self._command_mode.step_down(cli_service) cli_service.send_command.assert_called_once_with(self._exit_command, expected_string=parent_prompt, action_map=self._exit_action_map, error_map=self._exit_error_map) def test_step_down_set_mode(self): cli_service = Mock() parent_prompt = Mock() parent_node = Mock() parent_node.prompt = parent_prompt self._command_mode.parent_node = parent_node self._command_mode.step_down(cli_service) self.assertTrue(cli_service.command_mode == parent_node)