3999379869

Build Type

push

github

Committed by Jeremy Frasier

Commit Message

Bump version from 1.0.2 to 1.1.0

Run Details

16 of 113 branches covered (14.16%)

Branch coverage included in aggregate %.

1 of 1 new or added line in 1 file covered. (100.0%)

66 existing lines in 2 files now uncovered.

59 of 239 relevant lines covered (24.69%)

0.74 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

84.81

/src/awssh/awssh.py

"""awssh simplifies secure shell connections over AWS simple systems manager.

EXIT STATUS
    This utility exits with the same status as the underlying aws or ssh process:
    0   No error.
    >0  An error occurred.

Usage:
  awssh [options] <profile> <instance-id> [<command>...]
  awssh (-h | --help)

Options:
  -c --credentials=FILENAME     Shared credentials file name.
  -n --no-ssh                   Open an SSM shell session without using ssh.
  -r --region=REGION            AWS region name to use.
  -s --ssh-args=ARGUMENTS       Arguments to send to ssh.
  --log-level=LEVEL             If specified, then the log level will be set to
                                the specified value.  Valid values are "debug", "info",
                                "warning", "error", and "critical". [default: info]
  -h --help                     Show this message.
"""


# Standard Python Libraries
import logging
import os
from pathlib import Path
import signal
import subprocess  # nosec: B404 subprocess use is required for this tool
import sys
from typing import Any, Dict, Optional

# Third-Party Libraries
import docopt

# I cannot find a package that includes type stubs for schema, so I must add
# "type: ignore" to tell mypy to ignore this library
from schema import And, Schema, SchemaError, Use  # type: ignore

from . import CREDENTIAL_DIR
from ._version import __version__

# Options required for ssh to use ssm
# We will construct this for the user so they don't have to have it in their ssh config
# Note: They will probably still want to specify their own "User" option.
DEFAULT_SSH_OPTIONS = {
    "GSSAPIAuthentication": "yes",
    "GSSAPIDelegateCredentials": "yes",
    "ProxyCommand": """sh -c "aws ssm start-session --target %h --document-name AWS-StartSSHSession --parameters 'portNumber=%p'" """,
    "StrictHostKeyChecking": "no",
    "UserKnownHostsFile": "/dev/null",
}


def main() -> int:
    """Set up logging and prepare and SSM/ssh command."""
    args: Dict[str, str] = docopt.docopt(__doc__, version=__version__)

    # Validate and convert arguments as needed
    schema: Schema = Schema(
        {
            "--log-level": And(
                str,
                Use(str.lower),
                lambda n: n in ("debug", "info", "warning", "error", "critical"),
                error="Possible values for --log-level are "
                + "debug, info, warning, error, and critical.",
            ),
            str: object,  # Don't care about other keys, if any
        }
    )

    try:
        validated_args: Dict[str, Any] = schema.validate(args)
    except SchemaError as err:
        # Exit because one or more of the arguments were invalid
        print(err, file=sys.stderr)
        sys.exit(1)

    # Assign validated arguments to variables
    command: list[str] = validated_args["<command>"]
    credential_file: Optional[Path] = None
    if validated_args["--credentials"]:
        credential_file = CREDENTIAL_DIR / Path(validated_args["--credentials"])
    instance_id: str = validated_args["<instance-id>"]
    log_level: str = validated_args["--log-level"]
    no_ssh: bool = validated_args["--no-ssh"]
    profile: str = validated_args["<profile>"]
    region: str = validated_args["--region"]
    ssh_args: list[str]
    if validated_args["--ssh-args"]:
        ssh_args = validated_args["--ssh-args"].split()
    else:
        ssh_args = []

    if os.environ.get("AWSSH_USER"):
        ssh_args.append(f'-o User={os.environ["AWSSH_USER"]}')

    # Set up logging
    logging.basicConfig(
        format="%(asctime)-15s %(levelname)s %(message)s", level=log_level.upper()
    )

    returncode: int = _run_subprocess(
        credential_file, profile, region, instance_id, no_ssh, ssh_args, command
    )

    # Stop logging and clean up
    logging.shutdown()
    return returncode


def _run_subprocess(
    credential_file: Optional[Path],
    profile: str,
    region: Optional[str],
    instance_id: str,
    no_ssh: bool,
    ssh_args: list[str],
    remote_command: list[str],
) -> int:
    # Setup a modified environment for our awssh command
    awssh_env = os.environ.copy()
    if credential_file:
        awssh_env["AWS_SHARED_CREDENTIALS_FILE"] = str(credential_file)
    if region:
        awssh_env["AWS_DEFAULT_REGION"] = str(region)
    awssh_env["AWS_PROFILE"] = str(profile)
    if no_ssh:
        command = (
            ["aws", "ssm", "start-session"]
            + ["--target", instance_id]
            + ["--document-name", "SSM-SessionManagerRunShell"]
        )
    else:
        command = (
            ["ssh"]
            + [f"-o {key}={value}" for key, value in DEFAULT_SSH_OPTIONS.items()]
            + ssh_args
            + [instance_id]
            + remote_command
        )
    if logging.getLogger().isEnabledFor(logging.DEBUG):
        logging.debug("environment: ")
        for key in sorted(os.environ):
            logging.debug("%s=%s", key, os.environ[key])
        logging.debug("command: %s", " ".join(i for i in command))
    # Ignore interrupt signals and allow the child process to handle them
    interrupt_action = signal.signal(signal.SIGINT, signal.SIG_IGN)
    completed_process = (
        subprocess.run(  # nosec: B603 subprocess input is carefully validated
            args=command, env=awssh_env
        )
    )
    # Restore the interrupt signal handler to previous state
    signal.signal(signal.SIGINT, interrupt_action)
    return completed_process.returncode

1	"""awssh simplifies secure shell connections over AWS simple systems manager.
2
3	EXIT STATUS
4	This utility exits with the same status as the underlying aws or ssh process:
5	0 No error.
6	>0 An error occurred.
7
8	Usage:
9	awssh [options] <profile> <instance-id> [<command>...]
10	awssh (-h \| --help)
11
12	Options:
13	-c --credentials=FILENAME Shared credentials file name.
14	-n --no-ssh Open an SSM shell session without using ssh.
15	-r --region=REGION AWS region name to use.
16	-s --ssh-args=ARGUMENTS Arguments to send to ssh.
17	--log-level=LEVEL If specified, then the log level will be set to
18	the specified value. Valid values are "debug", "info",
19	"warning", "error", and "critical". [default: info]
20	-h --help Show this message.
21	"""
22
23
24	# Standard Python Libraries
25	import logging	3✔
26	import os	3✔
27	from pathlib import Path	3✔
28	import signal	3✔
29	import subprocess # nosec: B404 subprocess use is required for this tool	3✔
30	import sys	3✔
31	from typing import Any, Dict, Optional	3✔
32
33	# Third-Party Libraries
34	import docopt	3✔
35
36	# I cannot find a package that includes type stubs for schema, so I must add
37	# "type: ignore" to tell mypy to ignore this library
38	from schema import And, Schema, SchemaError, Use # type: ignore	3✔
39
40	from . import CREDENTIAL_DIR	3✔
41	from ._version import __version__	3✔
42
43	# Options required for ssh to use ssm
44	# We will construct this for the user so they don't have to have it in their ssh config
45	# Note: They will probably still want to specify their own "User" option.
46	DEFAULT_SSH_OPTIONS = {	3✔
47	"GSSAPIAuthentication": "yes",
48	"GSSAPIDelegateCredentials": "yes",
49	"ProxyCommand": """sh -c "aws ssm start-session --target %h --document-name AWS-StartSSHSession --parameters 'portNumber=%p'" """,
50	"StrictHostKeyChecking": "no",
51	"UserKnownHostsFile": "/dev/null",
52	}
53
54
55	def main() -> int:	3✔
56	"""Set up logging and prepare and SSM/ssh command."""
57	args: Dict[str, str] = docopt.docopt(__doc__, version=__version__)	3✔
58
59	# Validate and convert arguments as needed
60	schema: Schema = Schema(	3✔
61	{
62	"--log-level": And(
63	str,
64	Use(str.lower),
65	lambda n: n in ("debug", "info", "warning", "error", "critical"),
66	error="Possible values for --log-level are "
67	+ "debug, info, warning, error, and critical.",
68	),
69	str: object, # Don't care about other keys, if any
70	}
71	)
72
73	try:	3✔
74	validated_args: Dict[str, Any] = schema.validate(args)	3✔
75	except SchemaError as err:	3✔
76	# Exit because one or more of the arguments were invalid
77	print(err, file=sys.stderr)	3✔
78	sys.exit(1)	3✔
79
80	# Assign validated arguments to variables
81	command: list[str] = validated_args["<command>"]	3✔
82	credential_file: Optional[Path] = None	3✔
83	if validated_args["--credentials"]:	3!
UNCOV 84	credential_file = CREDENTIAL_DIR / Path(validated_args["--credentials"])	×
85	instance_id: str = validated_args["<instance-id>"]	3✔
86	log_level: str = validated_args["--log-level"]	3✔
87	no_ssh: bool = validated_args["--no-ssh"]	3✔
88	profile: str = validated_args["<profile>"]	3✔
89	region: str = validated_args["--region"]	3✔
90	ssh_args: list[str]
91	if validated_args["--ssh-args"]:	3!
UNCOV 92	ssh_args = validated_args["--ssh-args"].split()	×
93	else:
94	ssh_args = []	3✔
95
96	if os.environ.get("AWSSH_USER"):	3!
UNCOV 97	ssh_args.append(f'-o User={os.environ["AWSSH_USER"]}')	×
98
99	# Set up logging
100	logging.basicConfig(	3✔
101	format="%(asctime)-15s %(levelname)s %(message)s", level=log_level.upper()
102	)
103
104	returncode: int = _run_subprocess(	3✔
105	credential_file, profile, region, instance_id, no_ssh, ssh_args, command
106	)
107
108	# Stop logging and clean up
109	logging.shutdown()	3✔
110	return returncode	3✔
111
112
113	def _run_subprocess(	3✔
114	credential_file: Optional[Path],
115	profile: str,
116	region: Optional[str],
117	instance_id: str,
118	no_ssh: bool,
119	ssh_args: list[str],
120	remote_command: list[str],
121	) -> int:
122	# Setup a modified environment for our awssh command
123	awssh_env = os.environ.copy()	3✔
124	if credential_file:	3!
UNCOV 125	awssh_env["AWS_SHARED_CREDENTIALS_FILE"] = str(credential_file)	×
126	if region:	3!
UNCOV 127	awssh_env["AWS_DEFAULT_REGION"] = str(region)	×
128	awssh_env["AWS_PROFILE"] = str(profile)	3✔
129	if no_ssh:	3!
UNCOV 130	command = (	×
131	["aws", "ssm", "start-session"]
132	+ ["--target", instance_id]
133	+ ["--document-name", "SSM-SessionManagerRunShell"]
134	)
135	else:
136	command = (	3✔
137	["ssh"]
138	+ [f"-o {key}={value}" for key, value in DEFAULT_SSH_OPTIONS.items()]
139	+ ssh_args
140	+ [instance_id]
141	+ remote_command
142	)
143	if logging.getLogger().isEnabledFor(logging.DEBUG):	3✔
144	logging.debug("environment: ")	3✔
145	for key in sorted(os.environ):	3✔
146	logging.debug("%s=%s", key, os.environ[key])	3✔
147	logging.debug("command: %s", " ".join(i for i in command))	3✔
148	# Ignore interrupt signals and allow the child process to handle them
149	interrupt_action = signal.signal(signal.SIGINT, signal.SIG_IGN)	3✔
150	completed_process = (	3✔
151	subprocess.run( # nosec: B603 subprocess input is carefully validated
152	args=command, env=awssh_env
153	)
154	)
155	# Restore the interrupt signal handler to previous state
156	signal.signal(signal.SIGINT, interrupt_action)	3✔
157	return completed_process.returncode	3✔

cisagov / awssh / 3999379869

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous