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)