-
-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #20 from jiri-otoupal/ssh_processes
Change of ssh and ssh pod
- Loading branch information
Showing
16 changed files
with
302 additions
and
22 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Empty file.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,42 @@ | ||
| import click | ||
| import rich | ||
| from InquirerPy import inquirer | ||
|
|
||
| from abst.cli_commands.ssh_cli.utils import filter_keys_by_substring, filter_keys_by_port, do_ssh | ||
| from abst.config import broadcast_shm_name | ||
| from abst.sharing.local_broadcast import LocalBroadcast | ||
| from abst.utils.misc_funcs import setup_calls | ||
|
|
||
|
|
||
| @click.command("ssh", help="Will SSH into pod with containing string name") | ||
| @click.argument("port", required=False, default=None) | ||
| @click.option("-n", "--name", default=None, help="Name of context running, it can be just part of the name") | ||
| @click.option("--debug", is_flag=True, default=False) | ||
| def ssh_lin(port, name, debug): | ||
| setup_calls(debug) | ||
|
|
||
| lb = LocalBroadcast(broadcast_shm_name) | ||
| data = lb.retrieve_json() | ||
|
|
||
| if name and len(keys := filter_keys_by_substring(data, name)) == 1: | ||
| do_ssh(keys[0], data[keys[0]]["username"], data[keys[0]]["port"]) | ||
| return | ||
|
|
||
| if port is not None and len(keys := filter_keys_by_port(data, port)) == 1: | ||
| do_ssh(keys[0], data[keys[0]]["username"], data[keys[0]]["port"]) | ||
| return | ||
|
|
||
| longest_key = max(len(key) for key in data.keys()) | ||
| questions = [{"name": f"{key.ljust(longest_key)} |= status: {data['status']}", "value": (key, data)} for key, data | ||
| in data.items()] | ||
|
|
||
| context_name, context = inquirer.select("Select context to ssh to:", questions).execute() | ||
|
|
||
| if "username" not in context: | ||
| rich.print("[red]Please fill resource-os-username into config for this feature to work[/red]") | ||
| return | ||
| elif "port" not in context: | ||
| rich.print("[red]Please fill local-port into config for this feature to work[/red]") | ||
| return | ||
|
|
||
| do_ssh(name, context["username"], context["port"]) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,22 @@ | ||
| import os | ||
|
|
||
| import rich | ||
|
|
||
|
|
||
| def filter_keys_by_substring(dictionary: dict, substring: str) -> list: | ||
| # Convert substring to lower a case for case-insensitive comparison | ||
| substring_lower = substring.lower() | ||
| # Use a comprehension to filter keys | ||
| filtered_keys = [key for key in dictionary.keys() if substring_lower in key.lower()] | ||
| return filtered_keys | ||
|
|
||
|
|
||
| def filter_keys_by_port(data: dict, port: int): | ||
| """Filter and return keys from the servers dictionary where the port matches the given input.""" | ||
| return [key for key, data in data.items() if port in str(data.get('port', None))] | ||
|
|
||
|
|
||
| def do_ssh(context_name, username, port): | ||
| rich.print(f"[green]Running SSH to {context_name}[/green] " | ||
| f"[yellow]{username}@localhost:{port}[/yellow]") | ||
| os.system(f'ssh {username}@127.0.0.1 -p {port}') |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Empty file.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,102 @@ | ||
| import json | ||
| import struct | ||
| from json import JSONDecodeError | ||
| from multiprocessing import shared_memory | ||
|
|
||
| from deepmerge import always_merger | ||
|
|
||
| from abst.config import max_json_shared | ||
|
|
||
|
|
||
| class LocalBroadcast: | ||
| _instance = None | ||
|
|
||
| def __new__(cls, name: str, size: int = max_json_shared): | ||
| if cls._instance is None: | ||
| cls._instance = super(LocalBroadcast, cls).__new__(cls) | ||
| cls._instance.__init_shared_memory(name[:14], size) | ||
| return cls._instance | ||
|
|
||
| def __init_shared_memory(self, name: str, size: int): | ||
| self._data_name = name | ||
| self._len_name = f"{name}_len" | ||
| self._size = size | ||
|
|
||
| try: | ||
| # Attempt to create the main shared memory block | ||
| self._data_shm = shared_memory.SharedMemory(name=self._data_name, create=True, size=size) | ||
| self._data_is_owner = True | ||
| except FileExistsError: | ||
| self._data_shm = shared_memory.SharedMemory(name=self._data_name, create=False) | ||
| self._data_is_owner = False | ||
|
|
||
| try: | ||
| self._len_shm = shared_memory.SharedMemory(name=self._len_name, create=True, size=8) | ||
| self._len_shm.buf[:8] = struct.pack('Q', 0) | ||
| self._len_is_owner = True | ||
| except FileExistsError: | ||
| self._len_shm = shared_memory.SharedMemory(name=self._len_name, create=False) | ||
| self._len_is_owner = False | ||
|
|
||
| def store_json(self, data: dict) -> int: | ||
| """ | ||
| Serialize and store JSON data in shared memory. | ||
| @return: Size of the serialized data in bytes | ||
| """ | ||
|
|
||
| data_before = self.retrieve_json() | ||
| for key, value in data.items(): | ||
| for s_key in value.keys(): | ||
| if type(data_before[key][s_key]) == type(data[key][s_key]): | ||
| data_before[key].pop(s_key) | ||
|
|
||
| data_copy = always_merger.merge(data, data_before) | ||
|
|
||
| serialized_data = json.dumps(data_copy).encode('utf-8') | ||
| if len(serialized_data) > self._size: | ||
| raise ValueError("Data exceeds allocated shared memory size.") | ||
|
|
||
| # Write the data length to the length shared memory | ||
| self._len_shm.buf[:8] = struct.pack('Q', len(serialized_data)) | ||
|
|
||
| # Write data to the main shared memory | ||
| self._data_shm.buf[:len(serialized_data)] = serialized_data | ||
| return len(serialized_data) | ||
|
|
||
| def retrieve_json(self) -> dict: | ||
| """ | ||
| Retrieve and deserialize JSON data from shared memory. | ||
| """ | ||
| # Read the data length from the length shared memory | ||
| data_length = self.get_used_space() | ||
|
|
||
| if data_length == -1: | ||
| return {} | ||
|
|
||
| # Read data from the main shared memory | ||
| data = bytes(self._data_shm.buf[:data_length]).decode('utf-8') | ||
| try: | ||
| return json.loads(data) | ||
| except JSONDecodeError: | ||
| return {} | ||
|
|
||
| def get_used_space(self) -> int: | ||
| """ | ||
| Get the size of the shared memory | ||
| @return: Number of bytes used | ||
| """ | ||
| if self._len_shm.buf is None: | ||
| return -1 | ||
| return struct.unpack('Q', self._len_shm.buf[:8])[0] | ||
|
|
||
| def close(self): | ||
| """Close and unlink the shared memory blocks.""" | ||
| try: | ||
| self._data_shm.close() | ||
| self._len_shm.close() | ||
| if self._data_is_owner: | ||
| self._data_shm.unlink() | ||
| if self._len_is_owner: | ||
| self._len_shm.unlink() | ||
| except FileNotFoundError: | ||
| pass |
Empty file.
Oops, something went wrong.