From 57ceb8c33b7268071d6a57b936e3aac8df1c43eb Mon Sep 17 00:00:00 2001 From: Aron Novak Date: Thu, 8 Dec 2022 10:35:31 +0100 Subject: [PATCH 1/4] support custom commands --- eb_ssm/eb_ssm.py | 52 ++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 42 insertions(+), 10 deletions(-) diff --git a/eb_ssm/eb_ssm.py b/eb_ssm/eb_ssm.py index a610c74..23a5449 100644 --- a/eb_ssm/eb_ssm.py +++ b/eb_ssm/eb_ssm.py @@ -16,7 +16,7 @@ class SSMWrapper: def __init__(self): args = self._parse_args() - + # environment_name may be None self.environment_name = args.environment_name or get_current_branch_environment() self.profile = self._raise_if_none( @@ -30,6 +30,11 @@ def __init__(self): "Please specify a specific region in the command or eb configuration.", ) + self.command = args.command or "bash -l" + + self.non_interactive = args.non_interactive == 'True' + + def _parse_args(self): parser = argparse.ArgumentParser(description="SSH onto an Elastic Beanstalk Server") parser.add_argument( @@ -48,8 +53,18 @@ def _parse_args(self): default=None, help="use a specific region", ) + parser.add_argument( + "-c", "--command", + default=None, + help="command to execute", + ) + parser.add_argument( + "-n", "--non-interactive", + default=False, + help="Do not ask any questions, minimize output", + ) return parser.parse_args() - + def _raise_if_none(self, value, default_value, error_message): """ Return value if it is not None. If value is None, return default_value if it is not None. @@ -62,11 +77,11 @@ def _raise_if_none(self, value, default_value, error_message): else: io.log_error(error_message) sys.exit() - + def ssh(self): aws.set_region(self.region) aws.set_profile(self.profile) - + if self.environment_name is None: environment_names = get_all_environment_names() if environment_names: @@ -76,25 +91,42 @@ def ssh(self): else: io.log_error("The current Elastic Beanstalk application has no environments") sys.exit() - + instances = get_instance_ids(self.environment_name) - if len(instances) == 1: + if len(instances) == 1 or self.non_interactive is True: instance = instances[0] else: io.echo() io.echo('Select an instance to ssh into') instance = utils.prompt_for_item_in_list(instances) - + params = [ "aws", "ssm", "start-session", "--document-name", "AWS-StartInteractiveCommand", - "--parameters", "command='bash -l'", + "--parameters", "command='" + self.command + "'", "--profile", self.profile, "--region", self.region, "--target", instance, ] - - os.system(" ".join(params)) + cmd = " ".join(params) + + # Filter the output of os.system(cmd to remove strange output from Session Manager in non_interactive mode. + # The output might begin with: + # + # Starting session with SessionId: [id] + # + # and end with: + # Exiting session with sessionId: [id]. + # @see https://github.com/aws/session-manager-plugin/pull/20 + # @see https://github.com/aws/amazon-ssm-agent/issues/358 + if self.non_interactive is True: + # A strange workaround until we have the PRs above. + output = os.popen(cmd).read() + output = output.splitlines() + output = output[2:-4] + print("\n".join(output)) + else: + os.system(cmd) def main(): From 97829e8f1cee960c378031cb4c95db9558f9ced7 Mon Sep 17 00:00:00 2001 From: Aron Novak Date: Fri, 21 Apr 2023 14:58:37 +0200 Subject: [PATCH 2/4] drop output filtering --- eb_ssm/eb_ssm.py | 18 +----------------- 1 file changed, 1 insertion(+), 17 deletions(-) diff --git a/eb_ssm/eb_ssm.py b/eb_ssm/eb_ssm.py index 23a5449..0d53ecc 100644 --- a/eb_ssm/eb_ssm.py +++ b/eb_ssm/eb_ssm.py @@ -110,23 +110,7 @@ def ssh(self): ] cmd = " ".join(params) - # Filter the output of os.system(cmd to remove strange output from Session Manager in non_interactive mode. - # The output might begin with: - # - # Starting session with SessionId: [id] - # - # and end with: - # Exiting session with sessionId: [id]. - # @see https://github.com/aws/session-manager-plugin/pull/20 - # @see https://github.com/aws/amazon-ssm-agent/issues/358 - if self.non_interactive is True: - # A strange workaround until we have the PRs above. - output = os.popen(cmd).read() - output = output.splitlines() - output = output[2:-4] - print("\n".join(output)) - else: - os.system(cmd) + os.system(cmd) def main(): From 03812d1e34b55d850df912c394b06ec5bf996462 Mon Sep 17 00:00:00 2001 From: Aron Novak Date: Fri, 21 Apr 2023 15:02:54 +0200 Subject: [PATCH 3/4] introduce constant for default command --- eb_ssm/eb_ssm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/eb_ssm/eb_ssm.py b/eb_ssm/eb_ssm.py index 0d53ecc..c98f0bf 100644 --- a/eb_ssm/eb_ssm.py +++ b/eb_ssm/eb_ssm.py @@ -11,6 +11,7 @@ get_default_region, get_instance_ids) LOG = minimal_logger(__name__) +DEFAULT_COMMAND="bash -l" class SSMWrapper: @@ -30,7 +31,7 @@ def __init__(self): "Please specify a specific region in the command or eb configuration.", ) - self.command = args.command or "bash -l" + self.command = args.command or DEFAULT_COMMAND self.non_interactive = args.non_interactive == 'True' From c9c7ee99161eef7c858bb9812212b587abc16513 Mon Sep 17 00:00:00 2001 From: Aron Novak Date: Fri, 21 Apr 2023 15:18:34 +0200 Subject: [PATCH 4/4] fixes --- eb_ssm/eb_ssm.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/eb_ssm/eb_ssm.py b/eb_ssm/eb_ssm.py index c98f0bf..6b24452 100644 --- a/eb_ssm/eb_ssm.py +++ b/eb_ssm/eb_ssm.py @@ -33,7 +33,7 @@ def __init__(self): self.command = args.command or DEFAULT_COMMAND - self.non_interactive = args.non_interactive == 'True' + self.instance_number = args.number def _parse_args(self): @@ -60,9 +60,9 @@ def _parse_args(self): help="command to execute", ) parser.add_argument( - "-n", "--non-interactive", - default=False, - help="Do not ask any questions, minimize output", + "-n", "--number", + default=None, + help="Specify the instance to connect to by number.", ) return parser.parse_args() @@ -94,8 +94,11 @@ def ssh(self): sys.exit() instances = get_instance_ids(self.environment_name) - if len(instances) == 1 or self.non_interactive is True: - instance = instances[0] + if len(instances) == 1: + self.instance_number = 0 + + if self.instance_number is not None: + instance = instances[int(self.instance_number)] else: io.echo() io.echo('Select an instance to ssh into')