diff --git a/virtualbox/pybox/exceptions/__init__.py b/virtualbox/pybox/exceptions/__init__.py new file mode 100644 index 0000000..6c31278 --- /dev/null +++ b/virtualbox/pybox/exceptions/__init__.py @@ -0,0 +1 @@ +from .ssh_exception import InvalidSSHConnection \ No newline at end of file diff --git a/virtualbox/pybox/exceptions/ssh_exception.py b/virtualbox/pybox/exceptions/ssh_exception.py new file mode 100644 index 0000000..29278a1 --- /dev/null +++ b/virtualbox/pybox/exceptions/ssh_exception.py @@ -0,0 +1,2 @@ +class InvalidSSHConnection(Exception): + pass \ No newline at end of file diff --git a/virtualbox/pybox/install_vbox.py b/virtualbox/pybox/install_vbox.py index c99746a..46b7f41 100755 --- a/virtualbox/pybox/install_vbox.py +++ b/virtualbox/pybox/install_vbox.py @@ -34,11 +34,14 @@ from consts.networking import NICs, OAM, MGMT, Serial from consts.timeout import HostTimeout from consts import env +from exceptions import InvalidSSHConnection + from Parser import handle_args # Global vars V_BOX_OPTIONS = None +SSH_CONNECTIONS = {} # Network OAM_CONFIG = [getattr(OAM, attr) for attr in dir(OAM) if not attr.startswith('__')] @@ -972,6 +975,21 @@ def _connect_to_ssh(node='floating'): LOG.error("#### Failed SSH connection\nError: %s", repr(exc)) raise + SSH_CONNECTIONS[node] = ssh + return ssh + +def ssh_handler(node='floating'): + """ + Handles the SSH connection. Tries to retrieve a already existing connection. + If it doesn't exist, or isn't active, creates a new one. + """ + try: + ssh = SSH_CONNECTIONS[node] + ssh_transport = ssh.get_transport() + if ssh_transport is None or not ssh_transport.is_active(): + raise InvalidSSHConnection() + except (InvalidSSHConnection, KeyError): + ssh = _connect_to_ssh(node) return ssh @@ -1000,19 +1018,21 @@ def connect_to_ssh(node='floating'): Returns: return code of decorated function """ - - try: - ssh = _connect_to_ssh(node) - kwargs['ssh_client'] = ssh - return func(*args, **kwargs) - finally: - if ssh: - ssh.close() + ssh = ssh_handler(node) + kwargs['ssh_client'] = ssh + return func(*args, **kwargs) return connect_to_ssh_wrapper return connect_to_ssh_decorator +def close_ssh_connections(ssh_clients: dict): + """ + Closes the connection with all current created SSH Clients. + """ + for _, client in ssh_clients.items(): + if client is not None: + client.close() def stage_test_success(): """Prints a log message indicating the execution of a test stage.""" @@ -1790,7 +1810,7 @@ def run_custom_script(script, timeout, console, mode): LOG.info(" console mode: %s", console) LOG.info(" user mode: %s", mode) if console == 'ssh': - ssh_client = _connect_to_ssh() + ssh_client = ssh_handler() # pylint: disable=W0703, C0103 _, __, return_code = run_ssh_cmd(ssh_client, f"./{script}", timeout=timeout, mode=mode) if return_code != 0: @@ -2491,6 +2511,7 @@ def signal_handler(): """ print('You pressed Ctrl+C!') + close_ssh_connections(SSH_CONNECTIONS) kpi.print_kpi_metrics() sys.exit(1) @@ -2522,99 +2543,102 @@ def log_heading_msg(msg, pattern='#', panel_size=20): # pylint: disable=invalid-name if __name__ == "__main__": - kpi.init_kpi_metrics() - signal.signal(signal.SIGINT, signal_handler) + try: + kpi.init_kpi_metrics() + signal.signal(signal.SIGINT, signal_handler) - load_config() + load_config() - if V_BOX_OPTIONS.list_stages: - print(f"Defined setups: {list(STAGES_CHAINS.keys())}") - if V_BOX_OPTIONS.setup_type and V_BOX_OPTIONS.setup_type in AVAILABLE_CHAINS: - AVAILABLE_CHAINS = [V_BOX_OPTIONS.setup_type] - for stg_chain in AVAILABLE_CHAINS: - stg_no = 1 - print(f"Stages for setup on: {stg_chain}") - for stage in STAGES_CHAINS[stg_chain]: - print(wrap_stage_help(stage, STAGE_CALLBACKS[stage][HELP], stg_no)) - stg_no += 1 - print("Available stages that can be used for --custom-stages:") - for stage in AVAILABLE_STAGES: - print(wrap_stage_help(stage, STAGE_CALLBACKS[stage][HELP])) - sys.exit(0) + if V_BOX_OPTIONS.list_stages: + print(f"Defined setups: {list(STAGES_CHAINS.keys())}") + if V_BOX_OPTIONS.setup_type and V_BOX_OPTIONS.setup_type in AVAILABLE_CHAINS: + AVAILABLE_CHAINS = [V_BOX_OPTIONS.setup_type] + for stg_chain in AVAILABLE_CHAINS: + stg_no = 1 + print(f"Stages for setup on: {stg_chain}") + for stage in STAGES_CHAINS[stg_chain]: + print(wrap_stage_help(stage, STAGE_CALLBACKS[stage][HELP], stg_no)) + stg_no += 1 + print("Available stages that can be used for --custom-stages:") + for stage in AVAILABLE_STAGES: + print(wrap_stage_help(stage, STAGE_CALLBACKS[stage][HELP])) + sys.exit(0) - init_logging(V_BOX_OPTIONS.labname, V_BOX_OPTIONS.logpath) - LOG.info("Logging to directory: %s", (get_log_dir() + "/")) + init_logging(V_BOX_OPTIONS.labname, V_BOX_OPTIONS.logpath) + LOG.info("Logging to directory: %s", (get_log_dir() + "/")) - LOG.info("Install manages: %s controllers, %s workers, %s storages.", - V_BOX_OPTIONS.controllers, V_BOX_OPTIONS.workers, V_BOX_OPTIONS.storages) + LOG.info("Install manages: %s controllers, %s workers, %s storages.", + V_BOX_OPTIONS.controllers, V_BOX_OPTIONS.workers, V_BOX_OPTIONS.storages) - # Setup stages to run based on config - install_stages = [] - if V_BOX_OPTIONS.custom_stages: - # Custom stages - install_stages = V_BOX_OPTIONS.custom_stages.split(',') - for stage in install_stages: - invalid_stages = [] - if stage not in AVAILABLE_STAGES: - invalid_stages.append(stage) - if invalid_stages: - LOG.warning("Following custom stages are not supported: %s.\n" \ - "Choose from: %s", invalid_stages, AVAILABLE_STAGES) - sys.exit(1) - else: - # List all stages between 'from-stage' to 'to-stage' - stages = STAGES_CHAINS[V_BOX_OPTIONS.setup_type] - from_stg_index = 0 - to_stg_index = None - if V_BOX_OPTIONS.from_stage: - if V_BOX_OPTIONS.from_stage == 'start': - from_stg_index = 0 - else: - from_stg_index = stages.index(V_BOX_OPTIONS.from_stage) - if V_BOX_OPTIONS.to_stage: - if V_BOX_OPTIONS.from_stage == 'end': - to_stg_index = -1 - else: - to_stg_index = stages.index(V_BOX_OPTIONS.to_stage) + 1 - if to_stg_index is not None: - install_stages = stages[from_stg_index:to_stg_index] + # Setup stages to run based on config + install_stages = [] + if V_BOX_OPTIONS.custom_stages: + # Custom stages + install_stages = V_BOX_OPTIONS.custom_stages.split(',') + for stage in install_stages: + invalid_stages = [] + if stage not in AVAILABLE_STAGES: + invalid_stages.append(stage) + if invalid_stages: + LOG.warning("Following custom stages are not supported: %s.\n" \ + "Choose from: %s", invalid_stages, AVAILABLE_STAGES) + sys.exit(1) else: - install_stages = stages[from_stg_index:] - LOG.info("Executing %s stage(s): %s.", len(install_stages), install_stages) + # List all stages between 'from-stage' to 'to-stage' + stages = STAGES_CHAINS[V_BOX_OPTIONS.setup_type] + from_stg_index = 0 + to_stg_index = None + if V_BOX_OPTIONS.from_stage: + if V_BOX_OPTIONS.from_stage == 'start': + from_stg_index = 0 + else: + from_stg_index = stages.index(V_BOX_OPTIONS.from_stage) + if V_BOX_OPTIONS.to_stage: + if V_BOX_OPTIONS.from_stage == 'end': + to_stg_index = -1 + else: + to_stg_index = stages.index(V_BOX_OPTIONS.to_stage) + 1 + if to_stg_index is not None: + install_stages = stages[from_stg_index:to_stg_index] + else: + install_stages = stages[from_stg_index:] + LOG.info("Executing %s stage(s): %s.", len(install_stages), install_stages) - validate(V_BOX_OPTIONS, install_stages) + validate(V_BOX_OPTIONS, install_stages) - stg_no = 0 - prev_stage = None - for stage in install_stages: - stg_no += 1 - stg_start_time = time.time() - try: - stg_msg = f"({stg_no}/{len(install_stages)}) Entering stage {stage}" - log_heading_msg(stg_msg) - STAGE_CALLBACKS[stage][CALLBACK]() + stg_no = 0 + prev_stage = None + for stage in install_stages: + stg_no += 1 + stg_start_time = time.time() + try: + stg_msg = f"({stg_no}/{len(install_stages)}) Entering stage {stage}" + log_heading_msg(stg_msg) + STAGE_CALLBACKS[stage][CALLBACK]() - # Take snapshot if configured - if V_BOX_OPTIONS.snapshot: - vboxmanage.take_snapshot( - V_BOX_OPTIONS.labname, - f"snapshot-AFTER-{stage}") + # Take snapshot if configured + if V_BOX_OPTIONS.snapshot: + vboxmanage.take_snapshot( + V_BOX_OPTIONS.labname, + f"snapshot-AFTER-{stage}") - # Compute KPIs - stg_duration = time.time() - stg_start_time - kpi.set_kpi_metric(stage, stg_duration) - kpi.print_kpi(stage) - kpi.print_kpi('total') - except Exception as stg_exc: - stg_duration = time.time() - stg_start_time - kpi.set_kpi_metric(stage, stg_duration) - LOG.error("INSTALL FAILED, ABORTING!") - kpi.print_kpi_metrics() - LOG.info("Exception details: %s", repr(stg_exc)) - raise - # Stage completed - prev_stage = stage + # Compute KPIs + stg_duration = time.time() - stg_start_time + kpi.set_kpi_metric(stage, stg_duration) + kpi.print_kpi(stage) + kpi.print_kpi('total') + except Exception as stg_exc: + stg_duration = time.time() - stg_start_time + kpi.set_kpi_metric(stage, stg_duration) + LOG.error("INSTALL FAILED, ABORTING!") + kpi.print_kpi_metrics() + LOG.info("Exception details: %s", repr(stg_exc)) + raise + # Stage completed + prev_stage = stage - LOG.info("INSTALL SUCCEEDED!") - kpi.print_kpi_metrics() + LOG.info("INSTALL SUCCEEDED!") + kpi.print_kpi_metrics() + finally: + close_ssh_connections(SSH_CONNECTIONS)