From 94f78a493f82b6a19bedec5d35da4263a5900d31 Mon Sep 17 00:00:00 2001 From: Jiri Otoupal Date: Wed, 8 Nov 2023 16:31:40 +0100 Subject: [PATCH] Parallel update --- abst/__version__.py | 4 +- abst/bastion_support/bastion_scheduler.py | 40 ++++++++++++----- abst/bastion_support/oci_bastion.py | 12 +++-- abst/cli_commands/config_cli/commands.py | 8 ++-- abst/cli_commands/parallel/commands.py | 54 ++++++++++++++++++++--- abst/config.py | 4 +- abst/main.py | 4 +- abst/tools.py | 28 ++++++++---- 8 files changed, 114 insertions(+), 40 deletions(-) diff --git a/abst/__version__.py b/abst/__version__.py index 25ea9a2..8b1e6c8 100644 --- a/abst/__version__.py +++ b/abst/__version__.py @@ -10,11 +10,11 @@ "CLI Command making OCI Bastion and kubernetes usage simple and fast" ) -__version__ = "2.2.7" +__version__ = "2.2.8" __author__ = "Jiri Otoupal" __author_email__ = "jiri-otoupal@ips-database.eu" __license__ = "MIT" __url__ = "https://github.com/jiri-otoupal/abst" __pypi_repo__ = "https://pypi.org/project/abst/" -__version_name__ = "Simplified Regional Parrot ?" +__version_name__ = "Octopus X0" diff --git a/abst/bastion_support/bastion_scheduler.py b/abst/bastion_support/bastion_scheduler.py index 250e971..799e3ac 100644 --- a/abst/bastion_support/bastion_scheduler.py +++ b/abst/bastion_support/bastion_scheduler.py @@ -1,6 +1,8 @@ import signal +from pathlib import Path from threading import Thread from time import sleep +from typing import Optional import rich from click import clear @@ -110,23 +112,37 @@ def _run_indefinitely(cls, func, force: bool = False): @classmethod @load_stack_decorator - def run(cls, force=False): + def run(cls, force=False, set_dir: Optional[Path] = None): signal.signal(signal.SIGINT, BastionScheduler.kill_all) signal.signal(signal.SIGTERM, BastionScheduler.kill_all) rich.print("Will run all Bastions in parallel") thread_list = [] - for context_name in cls.__dry_stack: - if cls.stopped: - return - bastion = Bastion(None if context_name == "default" else context_name, region= - Bastion.load_json(Bastion.get_creds_path_resolve(context_name)).get("region", None)) - cls.__live_stack.add(bastion) - t = Thread(name=context_name, target=cls._run_indefinitely, - args=[bastion.create_forward_loop, force], daemon=True) - thread_list.append(t) - t.start() - rich.print(f"Started {context_name}") + if not set_dir: + for context_name in cls.__dry_stack: + if cls.stopped: + return + bastion = Bastion(None if context_name == "default" else context_name, region= + Bastion.load_json(Bastion.get_creds_path_resolve(context_name)).get("region", None)) + cls.__live_stack.add(bastion) + t = Thread(name=context_name, target=cls._run_indefinitely, + args=[bastion.create_forward_loop, force], daemon=True) + thread_list.append(t) + t.start() + rich.print(f"Started {context_name}") + else: + for context_path in filter(lambda p: not str(p.name).startswith("."), set_dir.iterdir()): + if cls.stopped: + return + context_name = context_path.name[:-5] + bastion = Bastion(None if context_name == "default" else context_name, region= + Bastion.load_json(context_path).get("region", None), direct_json_path=context_path) + cls.__live_stack.add(bastion) + t = Thread(name=context_name, target=cls._run_indefinitely, + args=[bastion.create_forward_loop, force], daemon=True) + thread_list.append(t) + t.start() + rich.print(f"Started {context_name}") cls.__display_loop() diff --git a/abst/bastion_support/oci_bastion.py b/abst/bastion_support/oci_bastion.py index 33a953c..cf4e803 100644 --- a/abst/bastion_support/oci_bastion.py +++ b/abst/bastion_support/oci_bastion.py @@ -16,7 +16,7 @@ from abst.config import default_creds_path, \ default_contexts_location, default_conf_path, \ - default_conf_contents, get_public_key + default_conf_contents, get_public_key, default_parallel_sets_location from abst.wrappers import mark_on_exit @@ -27,7 +27,7 @@ class Bastion: custom_ssh_options: str = "-o ServerAliveInterval=20" force_ssh_options: str = "-o StrictHostKeyChecking=no -o ServerAliveInterval=20 -o UserKnownHostsFile=/dev/null" - def __init__(self, context_name=None, region=None): + def __init__(self, context_name=None, region=None, direct_json_path=None): self.context_name = context_name self.region = region self.shell: bool = False @@ -36,6 +36,7 @@ def __init__(self, context_name=None, region=None): self.active_tunnel: subprocess.Popen = Optional[None] self.response: Optional[dict] = None self._current_status = None + self.direct_json_path = direct_json_path @property def current_status(self): @@ -355,7 +356,9 @@ def get_creds_path_resolve(cls, context_name) -> Path: return default_contexts_location / (context_name + ".json") def get_creds_path(self) -> Path: - if self.context_name is None: + if self.direct_json_path: + return self.direct_json_path + elif self.context_name is None: return default_creds_path else: return default_contexts_location / (self.context_name + ".json") @@ -391,9 +394,10 @@ def load_json(cls, path=default_creds_path) -> dict: return creds @classmethod - def create_default_location(cls): + def create_default_locations(cls): Path(default_creds_path.parent).mkdir(exist_ok=True) Path(default_contexts_location).mkdir(exist_ok=True) + Path(default_parallel_sets_location).mkdir(exist_ok=True) if not default_conf_path.exists(): cls.write_creds_json(default_conf_contents, default_conf_path) diff --git a/abst/cli_commands/config_cli/commands.py b/abst/cli_commands/config_cli/commands.py index aa49914..434742b 100644 --- a/abst/cli_commands/config_cli/commands.py +++ b/abst/cli_commands/config_cli/commands.py @@ -5,8 +5,8 @@ from abst.bastion_support.oci_bastion import Bastion from abst.cfg_func import __upgrade from abst.config import default_contexts_location -from abst.utils.misc_funcs import setup_calls from abst.tools import get_context_path +from abst.utils.misc_funcs import setup_calls @click.group(help="Group of commands for operations with config") @@ -22,7 +22,7 @@ def generate(debug, context_name): path = get_context_path(context_name) - Bastion.create_default_location() + Bastion.create_default_locations() td = Bastion.generate_sample_dict() creds_path = Bastion.write_creds_json(td, path) print( @@ -44,13 +44,13 @@ def fill(debug, context_name): if not path.exists(): rich.print("Generating sample Creds file") - Bastion.create_default_location() + Bastion.create_default_locations() td = Bastion.generate_sample_dict() Bastion.write_creds_json(td, path) if not default_contexts_location.exists(): rich.print("Generating contexts location") - Bastion.create_default_location() + Bastion.create_default_locations() rich.print(f"[green]Filling {str(path)}") rich.print("Please fill field one by one as displayed") diff --git a/abst/cli_commands/parallel/commands.py b/abst/cli_commands/parallel/commands.py index ec60ab7..de7cc2b 100644 --- a/abst/cli_commands/parallel/commands.py +++ b/abst/cli_commands/parallel/commands.py @@ -3,6 +3,7 @@ from InquirerPy import inquirer from abst.bastion_support.bastion_scheduler import BastionScheduler +from abst.config import default_parallel_sets_location from abst.tools import display_scheduled from abst.utils.misc_funcs import setup_calls @@ -36,6 +37,29 @@ def add(debug, context_name): display_scheduled() +@parallel.command("create", help="Create folder for a parallel set") +@click.option("--debug", is_flag=True, default=False) +@click.argument("set-name", default="default") +def create(debug, set_name): + """ + Create folder for a parallel set + + :param debug: + :param set_name: + :return: + """ + setup_calls(debug) + try: + set_dir = default_parallel_sets_location / set_name + set_dir.mkdir() + rich.print(f"[green]Set {set_name} created![/green]") + rich.print( + f"Feel free to copy any context to {set_dir} for it to be run in [yellow]{set_name}[/yellow] parallel set") + rich.print(f"You can run this set by 'abst parallel run {set_name}'") + except OSError: + rich.print("[red]Set already exists[/red]") + + @parallel.command("remove", help="Remove Bastion from stack") @click.option("--debug", is_flag=True, default=False) @click.argument("context-name", default="default") @@ -49,12 +73,18 @@ def remove(debug, context_name): @click.option("--debug", is_flag=True, default=False) @click.option("-y", is_flag=True, default=False, help="Automatically confirm") @click.option("-f", "--force", is_flag=True, default=False, help="Will force connections ignoring security policies") -def run(debug, y, force): +@click.argument("set_name", default=None, required=False, type=str) +def run(debug, y, force, set_name=None): setup_calls(debug) if force: rich.print("[red]Running in force mode[/red][gray] this mode is less secure as it is ignoring key checking and" " security policies involving known_hosts[/gray]") - display_scheduled() + if set_name: + set_dir = get_set_dir(set_name) + else: + set_dir = None + + display_scheduled(set_dir) if not y: try: confirm = inquirer.confirm( @@ -66,12 +96,24 @@ def run(debug, y, force): if not confirm: rich.print("[green]Cancelling, nothing started[/green]") exit(0) - BastionScheduler.run(force) + BastionScheduler.run(force, set_dir) + + +def get_set_dir(set_name): + set_dir = default_parallel_sets_location / set_name + if not set_dir.exists(): + rich.print(f"Parallel set {set_name} did not found in {set_dir}") + exit(1) + elif len(list(set_dir.iterdir())) == 0: + rich.print(f"[red]No contexts found in {set_dir}[/red]") + exit(1) + return set_dir @parallel.command("display", help="Display current Bastions is stack") @click.option("--debug", is_flag=True, default=False) -def display(debug): +@click.argument("set_name", default=None, required=False, type=str) +def display(debug, set_name): setup_calls(debug) - - display_scheduled() + set_dir = get_set_dir(set_name) + display_scheduled(set_dir) diff --git a/abst/config.py b/abst/config.py index 597e1b3..775681b 100644 --- a/abst/config.py +++ b/abst/config.py @@ -4,10 +4,12 @@ default_stack_location: Path = (Path().home().resolve() / ".abst" / "stack.json") default_conf_path: Path = (Path().home().resolve() / ".abst" / "config.json") default_creds_path: Path = (Path().home().resolve() / ".abst" / "creds.json") +default_contexts_location: Path = (Path().home().resolve() / ".abst" / "contexts") +default_parallel_sets_location: Path = (Path().home().resolve() / ".abst" / "sets") + default_config_keys: tuple = ( "host", "bastion-id", "default-name", "ssh-pub-path", "private-key-path", "target-ip", "local-port", "target-port", "ttl", "resource-id", "resource-os-username", "region") -default_contexts_location: Path = (Path().home().resolve() / ".abst" / "contexts") default_stack_contents: dict = {"stack": []} default_conf_contents: dict = {"used_context": None, "private-key-path": "~/.ssh/id_rsa", "ssh-pub-path": "~/.ssh/id_rsa" diff --git a/abst/main.py b/abst/main.py index a4699e0..c164624 100644 --- a/abst/main.py +++ b/abst/main.py @@ -48,7 +48,7 @@ def use(debug, context_name): contexts.insert(0, "default") used_context = inquirer.select("Select context to use:", contexts).execute() - Bastion.create_default_location() + Bastion.create_default_locations() conf = Bastion.load_config() conf["used_context"] = None if used_context == "default" else used_context @@ -89,7 +89,7 @@ def main(): Notifier.notify() except (ConnectionError, ConnectTimeout): return False - Bastion.create_default_location() + Bastion.create_default_locations() cli() diff --git a/abst/tools.py b/abst/tools.py index f2dd1ff..4c5810e 100644 --- a/abst/tools.py +++ b/abst/tools.py @@ -1,6 +1,9 @@ +from pathlib import Path +from typing import Optional + from abst.bastion_support.bastion_scheduler import BastionScheduler -from abst.config import default_creds_path, default_contexts_location from abst.bastion_support.oci_bastion import Bastion +from abst.config import default_creds_path, default_contexts_location def get_context_path(context_name): @@ -10,20 +13,27 @@ def get_context_path(context_name): return default_contexts_location / (context_name + ".json") -def display_scheduled(): +def display_scheduled(set_dir: Optional[Path] = None): from rich.table import Table from rich.console import Console console = Console() - table = Table(title="Current Bastions in Stack", highlight=True) + table = Table(title=f"Bastions of {set_dir.name} set" if set_dir else "Bastions of default stack", highlight=True) table.add_column("Name", justify="left", style="cyan", no_wrap=True) table.add_column("Local Port", style="magenta", no_wrap=True) table.add_column("Active", justify="right", style="green", no_wrap=True) table.add_column("Status", justify="right", style="green", no_wrap=True) - - for context_name in BastionScheduler.get_bastions(): - conf = Bastion.load_json(Bastion.get_creds_path_resolve(context_name)) - table.add_row(context_name, conf.get('local-port', 'Not Specified'), - conf.get('target-ip', 'Not Specified'), - conf.get('target-port', 'Not Specified')) + if set_dir: + for context_path in filter(lambda p: not str(p.name).startswith("."), set_dir.iterdir()): + conf = Bastion.load_json(context_path) + context_name = context_path.name[:-5] + table.add_row(context_name, conf.get('local-port', 'Not Specified'), + conf.get('target-ip', 'Not Specified'), + conf.get('target-port', 'Not Specified')) + else: + for context_path in BastionScheduler.get_bastions(): + conf = Bastion.load_json(Bastion.get_creds_path_resolve(context_path)) + table.add_row(context_path, conf.get('local-port', 'Not Specified'), + conf.get('target-ip', 'Not Specified'), + conf.get('target-port', 'Not Specified')) console.print(table)