diff --git a/.coveragerc b/.coveragerc index d315b87..601b59a 100644 --- a/.coveragerc +++ b/.coveragerc @@ -2,7 +2,7 @@ # https://coverage.readthedocs.io/en/latest/config.html [run] -source = src/example +source = src/awssh omit = branch = true diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 31d1120..af12b96 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -23,12 +23,12 @@ jobs: - id: setup-python uses: actions/setup-python@v2 with: - python-version: 3.9 + python-version: "3.10" # We need the Go version and Go cache location for the actions/cache step, # so the Go installation must happen before that. - uses: actions/setup-go@v2 with: - go-version: '1.16' + go-version: "1.16" - name: Store installed Go version id: go-version run: | @@ -112,9 +112,6 @@ jobs: fail-fast: false matrix: python-version: - - "3.6" - - "3.7" - - "3.8" - "3.9" - "3.10" steps: @@ -165,7 +162,7 @@ jobs: - id: setup-python uses: actions/setup-python@v2 with: - python-version: 3.9 + python-version: "3.10" - uses: actions/cache@v2 env: BASE_CACHE_KEY: "${{ github.job }}-${{ runner.os }}-\ @@ -199,9 +196,6 @@ jobs: fail-fast: false matrix: python-version: - - "3.6" - - "3.7" - - "3.8" - "3.9" - "3.10" steps: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index bf28c97..721cda6 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -109,7 +109,13 @@ repos: hooks: - id: mypy additional_dependencies: + - boto3-stubs + - pytest-mypy + - types-docopt + - types-requests - types-setuptools + args: + - --scripts-are-modules - repo: https://github.com/asottile/pyupgrade rev: v2.31.0 hooks: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 7d761ad..cea26ce 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -15,7 +15,7 @@ all of which should be in this repository. If you want to report a bug or request a new feature, the most direct method is to [create an -issue](https://github.com/cisagov/ssm-ssh/issues) in +issue](https://github.com/cisagov/awssh/issues) in this repository. We recommend that you first search through existing issues (both open and closed) to check if your particular issue has already been reported. If it has then you might want to add a comment @@ -25,7 +25,7 @@ one. ## Pull requests ## If you choose to [submit a pull -request](https://github.com/cisagov/ssm-ssh/pulls), +request](https://github.com/cisagov/awssh/pulls), you will notice that our continuous integration (CI) system runs a fairly extensive set of linters, syntax checkers, system, and unit tests. Your pull request may fail these checks, and that's OK. If you want @@ -135,9 +135,9 @@ can create and configure the Python virtual environment with these commands: ```console -cd ssm-ssh -pyenv virtualenv ssm-ssh -pyenv local ssm-ssh +cd awssh +pyenv virtualenv awssh +pyenv local awssh pip install --requirement requirements-dev.txt ``` diff --git a/README.md b/README.md index 9a36e08..adeeb74 100644 --- a/README.md +++ b/README.md @@ -1,25 +1,75 @@ -# ssm-ssh # - -[![GitHub Build Status](https://github.com/cisagov/ssm-ssh/workflows/build/badge.svg)](https://github.com/cisagov/ssm-ssh/actions) -[![Coverage Status](https://coveralls.io/repos/github/cisagov/ssm-ssh/badge.svg?branch=develop)](https://coveralls.io/github/cisagov/ssm-ssh?branch=develop) -[![Total alerts](https://img.shields.io/lgtm/alerts/g/cisagov/ssm-ssh.svg?logo=lgtm&logoWidth=18)](https://lgtm.com/projects/g/cisagov/ssm-ssh/alerts/) -[![Language grade: Python](https://img.shields.io/lgtm/grade/python/g/cisagov/ssm-ssh.svg?logo=lgtm&logoWidth=18)](https://lgtm.com/projects/g/cisagov/ssm-ssh/context:python) -[![Known Vulnerabilities](https://snyk.io/test/github/cisagov/ssm-ssh/develop/badge.svg)](https://snyk.io/test/github/cisagov/ssm-ssh) - -This is a generic skeleton project that can be used to quickly get a -new [cisagov](https://github.com/cisagov) Python library GitHub -project started. This skeleton project contains [licensing -information](LICENSE), as well as -[pre-commit hooks](https://pre-commit.com) and -[GitHub Actions](https://github.com/features/actions) configurations -appropriate for a Python library project. - -## New Repositories from a Skeleton ## - -Please see our [Project Setup guide](https://github.com/cisagov/development-guide/tree/develop/project_setup) -for step-by-step instructions on how to start a new repository from -a skeleton. This will save you time and effort when configuring a -new repository! +# awssh ☁️🔒🐚 # + +[![GitHub Build Status](https://github.com/cisagov/awssh/workflows/build/badge.svg)](https://github.com/cisagov/awssh/actions) +[![Coverage Status](https://coveralls.io/repos/github/cisagov/awssh/badge.svg?branch=develop)](https://coveralls.io/github/cisagov/awssh?branch=develop) +[![Total alerts](https://img.shields.io/lgtm/alerts/g/cisagov/awssh.svg?logo=lgtm&logoWidth=18)](https://lgtm.com/projects/g/cisagov/awssh/alerts/) +[![Language grade: Python](https://img.shields.io/lgtm/grade/python/g/cisagov/awssh.svg?logo=lgtm&logoWidth=18)](https://lgtm.com/projects/g/cisagov/awssh/context:python) +[![Known Vulnerabilities](https://snyk.io/test/github/cisagov/awssh/develop/badge.svg)](https://snyk.io/test/github/cisagov/awssh) + +This project provides a tool that simplifies secure shell connections over +[AWS Systems Manager](https://docs.aws.amazon.com/systems-manager/latest/userguide/what-is-systems-manager.html) +(formerly known as SSM). + +## Pre-requisites ## + +- The [AWS CLI](https://aws.amazon.com/cli/) installed on your system. +- A valid AWS profile that has permissions to start/stop SSM sessions. +- A [`bash`](https://www.gnu.org/software/bash/) shell. + +## Usage ## + +### Setup ### + +1. Install the `awssh` command line utility. One easy way to do this is + to run the `setup-env` script in the main directory. +1. Define environment variables: + - `AWSSH_PROFILE_FILTER`: A string that will match one or more profiles + in your AWS configuration file that have permission to start/stop SSM + sessions. + - `AWSSH_USER`: The username to use for ssh connections over SSM. + + ```bash + export AWSSH_PROFILE_FILTER="startstopssmsession" + export AWSSH_USER="lemmy.kilmister" + ``` + +1. Source the [`awssh-completion.bash`](tools/awssh-completion.bash) file in + your `bash` environment: + + ```bash + source tools/awssh-completion.bash + ``` + + If you skip this step, you won't get to enjoy any of that sweet, sweet + tab completion that will make life a lot easier for you. Don't say we + didn't warn you. + +### Start a SSM shell session without ssh ### + +```console +awssh --no-ssh my-aws-startstopssmsession-profile i-01234567890abcdef +``` + +### Start a SSM shell session with ssh ### + +```console +awssh my-aws-startstopssmsession-profile i-01234567890abcdef +``` + +Tab completion can be used to autocomplete the following items as you type +your `awssh` command: + +- Shared credentials file (following `-c`, `--credentials=FILENAME`), by + showing matching files in the `.aws` directory in your home directory + (e.g. `~/.aws/`) +- AWS region (`-r`, `--region`) +- AWS profile (``), provided your chosen (or default) credentials + file contains at least one profile that matches the string specified by the + `AWSSH_PROFILE_FILTER` environment variable +- AWS instance you want to open a session to (``); note that + if your instance is tagged with a name, you can start typing that name and + when you tab complete, the name will be transformed into the instance ID + (assuming you have typed enough of the name to identify a unique instance). ## Contributing ## diff --git a/bump_version.sh b/bump_version.sh index e1324b8..3dd1e89 100755 --- a/bump_version.sh +++ b/bump_version.sh @@ -6,7 +6,7 @@ set -o nounset set -o errexit set -o pipefail -VERSION_FILE=src/example/_version.py +VERSION_FILE=src/awssh/_version.py HELP_INFORMATION="bump_version.sh (show|major|minor|patch|prerelease|build|finalize)" diff --git a/setup.py b/setup.py index 17f8760..b0f1c5e 100644 --- a/setup.py +++ b/setup.py @@ -1,5 +1,5 @@ """ -This is the setup module for the example project. +This is the setup module for the awssh project. Based on: @@ -42,10 +42,10 @@ def get_version(version_file): setup( - name="example", + name="awssh", # Versions should comply with PEP440 - version=get_version("src/example/_version.py"), - description="Example Python library", + version=get_version("src/awssh/_version.py"), + description="A tool that simplifies secure shell connections over AWS Systems Manager", long_description=readme(), long_description_content_type="text/markdown", # Landing page for CISA's cybersecurity mission @@ -53,8 +53,8 @@ def get_version(version_file): # Additional URLs for this project per # https://packaging.python.org/guides/distributing-packages-using-setuptools/#project-urls project_urls={ - "Source": "https://github.com/cisagov/ssm-ssh", - "Tracker": "https://github.com/cisagov/ssm-ssh/issues", + "Source": "https://github.com/cisagov/awssh", + "Tracker": "https://github.com/cisagov/awssh/issues", }, # Author details author="Cybersecurity and Infrastructure Security Agency", @@ -74,23 +74,25 @@ def get_version(version_file): # Specify the Python versions you support here. In particular, ensure # that you indicate whether you support Python 2, Python 3 or both. "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.6", - "Programming Language :: Python :: 3.7", - "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", ], - python_requires=">=3.6", + python_requires=">=3.9", # What does your project relate to? - keywords="skeleton", + keywords="aws ssm ssh", packages=find_packages(where="src"), package_dir={"": "src"}, - package_data={"example": ["data/*.txt"]}, py_modules=[splitext(basename(path))[0] for path in glob("src/*.py")], include_package_data=True, - install_requires=["docopt", "schema", "setuptools >= 24.2.0"], + install_requires=[ + "boto3", + "docopt", + "schema", + "setuptools >= 24.2.0", + ], extras_require={ "test": [ + "boto3-stubs", "coverage", # coveralls 1.11.0 added a service number for calls from # GitHub Actions. This caused a regression which resulted in a 422 @@ -102,8 +104,16 @@ def get_version(version_file): "pre-commit", "pytest-cov", "pytest", + "types-docopt", + "types-requests", + "types-setuptools", + ] + }, + # Conveniently allows one to run the CLI tool as `awssh` + entry_points={ + "console_scripts": [ + "awssh = awssh.awssh:main", + "_awssh-completer = awssh.autocompleter:main", ] }, - # Conveniently allows one to run the CLI tool as `example` - entry_points={"console_scripts": ["example = example.example:main"]}, ) diff --git a/src/example/__init__.py b/src/awssh/__init__.py similarity index 57% rename from src/example/__init__.py rename to src/awssh/__init__.py index 98b5e04..d732b96 100644 --- a/src/example/__init__.py +++ b/src/awssh/__init__.py @@ -1,9 +1,14 @@ -"""The example library.""" +"""The awssh library.""" +# Standard Python Libraries +from pathlib import Path + # We disable a Flake8 check for "Module imported but unused (F401)" here because # although this import is not directly used, it populates the value # package_name.__version__, which is used to get version information about this # Python package. from ._version import __version__ # noqa: F401 -from .example import example_div -__all__ = ["example_div"] +CREDENTIAL_DIR = Path("~/.aws").expanduser() +DEFAULT_CREDENTIAL_FILE = Path(CREDENTIAL_DIR) / Path("credentials") + +__all__ = ["CREDENTIAL_DIR"] diff --git a/src/example/__main__.py b/src/awssh/__main__.py similarity index 73% rename from src/example/__main__.py rename to src/awssh/__main__.py index 11a3238..91c65b9 100644 --- a/src/example/__main__.py +++ b/src/awssh/__main__.py @@ -1,5 +1,5 @@ """Code to run if this package is used as a Python module.""" -from .example import main +from .awssh import main main() diff --git a/src/example/_version.py b/src/awssh/_version.py similarity index 70% rename from src/example/_version.py rename to src/awssh/_version.py index 5eb9b0e..de155d7 100644 --- a/src/example/_version.py +++ b/src/awssh/_version.py @@ -1,2 +1,2 @@ """This file defines the version of this module.""" -__version__ = "0.1.0" +__version__ = "1.0.0" diff --git a/src/awssh/autocompleter.py b/src/awssh/autocompleter.py new file mode 100644 index 0000000..309eb98 --- /dev/null +++ b/src/awssh/autocompleter.py @@ -0,0 +1,349 @@ +"""AWS SSH autocompleter script.""" + +# Standard Python Libraries +import configparser +from dataclasses import dataclass +import os +from pathlib import Path +import re +import shlex +from typing import Optional, TextIO + +# Third-Party Libraries +import boto3 + +from . import CREDENTIAL_DIR, DEFAULT_CREDENTIAL_FILE + +CREDENTIALS_OPTIONS = {"--credentials", "-c"} +HELP_OPTIONS = {"--help", "-h"} +LOG_LEVEL_OPTIONS = {"--log-level"} +LOG_LEVELS = {"debug", "info", "warning", "error", "critical"} +REGION_OPTIONS = {"--region", "-r"} +SSH_ARGS_OPTIONS = {"--ssh-args", "-s"} + +if os.environ.get("LC_CTYPE", "") == "UTF-8": + os.environ["LC_CTYPE"] = "en_US.UTF-8" + +# Debugging the completer can be touchy since standard out goes to bash. +# Set the environment below to a filename to enable logging. +LOG_FILE: Optional[TextIO] = None +if filename := os.environ.get("BASH_COMP_DEBUG_FILE"): + LOG_FILE = open(filename, "a", encoding="utf-8") + + +def log(message: str) -> None: + """Log a message to the log file if it is open. + + Args: + message (str): The message to log. + """ + if LOG_FILE: + print(message, file=LOG_FILE) + + +def get_cred_files() -> set[str]: + """Return a set of all credential files in the credential directory. + + Returns: + set[str]: A set of all credential files in the credential directory. + """ + result: set[str] = set() + for i in CREDENTIAL_DIR.iterdir(): + if i.is_file(): + result.add(i.name) + return result + + +def get_regions() -> set[str]: + """Return a set of all available AWS regions. + + Returns: + set[str]: A set of all available AWS regions. + """ + session = boto3.session.Session(region_name="us-east-1") + return set(session.get_available_regions("ec2")) + + +def get_profiles(cred_filename: Path, profile_filter: Optional[str] = None) -> set[str]: + """Return a set of profiles within a credential file that match the filter. + + Args: + cred_filename (Path): The path to the credential file. + profile_filter (Optional[str], optional): Regular expression to match profile names. Defaults to None. + + Returns: + set[str]: A set of profiles within a credential file that match the filter. + """ + if profile_filter: + profile_filter_re: re.Pattern[str] = re.compile(profile_filter) + config = configparser.ConfigParser() + config.read(cred_filename) + result: set[str] = set() + for section in config.sections(): + if not profile_filter_re or profile_filter_re.search(section): + result.add(section) + return result + + +def get_instances( + credential_file: Path, profile: str, region: str +) -> set[tuple[str, str]]: + """Return a set of instances and names that match the profile and region. + + Args: + credential_file (Path): The path to the credential file. + profile (str): The profile to use. + region (str): The AWS region to use. + + Returns: + set[tuple[str, str]]: A set of instances and names that match the profile and region. + """ + # Create session + # boto3 doesn't have a programatic way to set the credential file. yuck. + os.environ["AWS_SHARED_CREDENTIALS_FILE"] = str(credential_file) + session = boto3.session.Session(region_name=region, profile_name=profile) + ec2 = session.resource("ec2") + instances = {} + # Loop through instances that are running + for i in ec2.instances.filter( + Filters=[{"Name": "instance-state-name", "Values": ["running"]}] + ): + # convert tags into a proper dictionary + if i.tags is None: + # if an instance has no tags use the id as the Name + log(f"warning: instance does not have tags: {i}") + i.tag_dict = {"Name": i.id} + else: + i.tag_dict = {tag["Key"]: tag["Value"] for tag in i.tags} + instances[i.tag_dict["Name"]] = i + + result: set[tuple[str, str]] = set() + for name, instance in instances.items(): + result.add((instance.id, name)) + return result + + +def build_option_candidates(word_set: set[str]) -> set[str]: + """Return a set of command line option candidates for the given word set. + + Args: + word_set (set[str]): The set of words to build option candidates for. + + Returns: + set[str]: A set of command line option candidates for the given word set. + """ + candidates: set[str] = set() + if HELP_OPTIONS & word_set: + return candidates + if len(word_set) <= 1: + candidates |= HELP_OPTIONS + if not CREDENTIALS_OPTIONS & word_set: + candidates |= CREDENTIALS_OPTIONS + if not LOG_LEVEL_OPTIONS & word_set: + candidates |= LOG_LEVEL_OPTIONS + if not REGION_OPTIONS & word_set: + candidates |= REGION_OPTIONS + if not SSH_ARGS_OPTIONS & word_set: + candidates |= SSH_ARGS_OPTIONS + return candidates + + +@dataclass +class ParsedState: + """A class to hold the state of the command line parser.""" + + aws_region: Optional[str] = os.environ.get("AWS_REGION") + cred_file: Path = Path( + os.environ.get("AWS_SHARED_CREDENTIALS_FILE", DEFAULT_CREDENTIAL_FILE) + ) + instance: Optional[str] = None + profile: Optional[str] = None + ssh_args: Optional[str] = None + ssh_command: Optional[list[str]] = None + + +def parse_command_line(words: list[str]) -> ParsedState: + """Parse the command line and return the parsed state. + + Args: + words (list[str]): The command line words to parse. + + Returns: + ParsedState: The parsed state. + """ + state: ParsedState = ParsedState() + pos_args: list[str] = [] + i = iter(words[:-1]) + try: + while True: + word = next(i) + if word in CREDENTIALS_OPTIONS: + state.cred_file = Path(CREDENTIAL_DIR) / Path(next(i)) + elif word in REGION_OPTIONS: + state.aws_region = next(i) + elif word in LOG_LEVEL_OPTIONS: + next(i) # don't care + elif word in SSH_ARGS_OPTIONS: + state.ssh_args = next(i) + elif not word.startswith("-"): + pos_args.append(word) + except StopIteration: + pass + if pos_args: + state.profile = pos_args.pop(0) + if pos_args: + state.instance = pos_args.pop(0) + if pos_args: + state.ssh_command = pos_args + return state + + +def process_command_line( + command_line: str, command_index: int +) -> tuple[list[str], str, str]: + """Process the command line and return the command line words, the command, and the command index. + + Args: + command_line (str): The command line to process. + command_index (int): The index of the current command in the command line. + + Returns: + tuple[list[str], str, str]: The command line words, the command, and the command index. + """ + # Chop the command line at the current cursor location + command_line = command_line[:command_index] + log(f"chopped command line is: {command_line}") + words = shlex.split(command_line) + del words[0] # delete the command itself + if command_line[-1].isspace(): + # add a new empty word since we have a trailing space + words.append("") + cur_word = words[-1] if len(words) > 0 else "" + prev_word = words[-2] if len(words) > 1 else "" + log(f"words is {words}") + log(f"cur_word is {cur_word}, prev_word is {prev_word}") + return words, cur_word, prev_word + + +def autocomplete(command_line: str, command_index: int) -> tuple[set[str], str]: + """Autocomplete the command line and return the autocomplete set and the current word. + + Args: + command_line (str): The command line to autocomplete. + command_index (int): The index of the current command in the command line. + + Returns: + tuple(set[str], str): The autocomplete set and the current word. + """ + log("-" * 40) + log(f"command_line is {command_line}, command_index is {command_index}") + words: list[str] + cur_word: str + prev_word: str + words, cur_word, prev_word = process_command_line(command_line, command_index) + word_set: set[str] = set(words) + + # Calculate the state from the current commandline + state = parse_command_line(words) + log(f"state is {state.__dict__}") + + # Build a list of candidate completions + candidates: set[str] = set() + + # If we are working on an option, then suggest parameters for that option. + if prev_word in CREDENTIALS_OPTIONS: + candidates |= get_cred_files() + elif prev_word in LOG_LEVEL_OPTIONS: + candidates |= LOG_LEVELS + elif prev_word in REGION_OPTIONS: + candidates |= get_regions() + else: + # Not working on an option parameter + if not state.profile: + # If we don't have the positional profile argument, suggest options + candidates |= build_option_candidates(word_set) + + # If we have enough optional information start suggesting positional parameters + if state.cred_file.is_file() and state.aws_region: + profile_filter = os.environ.get("AWSSH_PROFILE_FILTER") + profiles = get_profiles(state.cred_file, profile_filter) + if state.profile not in profiles: + candidates |= profiles + + # If we have enough information suggest instance ids and names + if not state.instance and state.profile in profiles: + instances: set[tuple[str, str]] = get_instances( + state.cred_file, state.profile, state.aws_region + ) + # Build sets of id-name pairs, and names alone + instance_id_and_name: set[str] = {f"{i[0]} {i[1]}" for i in instances} + instance_names: set[str] = {i[1] for i in instances} + + # Include the names alone. They will be substituted with an + # instance id when only one candidate is left. + candidates |= instance_names + + # Filter instances now so we can strip off the name if there is only one + filtered_id_and_name = { + i for i in instance_id_and_name if cur_word in i + } + if len(filtered_id_and_name) == 1: + # If there is only one matching instance, remove the Name suffix + candidates = {filtered_id_and_name.pop().split()[0]} + else: + candidates |= instance_id_and_name + + return candidates, cur_word + + +def print_completions(comps: set[str], cur: str) -> None: + """Output the completions to stdout. + + Args: + comps (set[str]): The completions to output. + cur (str): The current word. + """ + log(f"completions: {sorted(comps)}") + contains: list[str] = [] + starts: list[str] = [] + for i in sorted(comps): + if i.startswith(cur): + starts.append(i) + if cur in i or cur == "": + contains.append(i) + log(f"starts with: {starts}") + log(f"contains: {contains}") + # Prefer completions in this order: + # 1. completions that start with the current word + # 2. completions that contain the current word + # 3. other candidates that matched some other way (e.g., by instance name) + for i in starts or contains or comps: + print(i) + + +def main() -> None: + """Execute main entry point for the awssh autocomplete script. + + Returns: + None + """ + # bash exports COMP_LINE and COMP_POINT, tcsh COMMAND_LINE only + command_line: str = ( + os.environ.get("COMP_LINE") or os.environ.get("COMMAND_LINE") or "" + ) + command_index: int = int(os.environ.get("COMP_POINT") or len(command_line)) + + try: + candidates: set[str] + cur_word: str + candidates, cur_word = autocomplete(command_line, command_index) + + print_completions(candidates, cur_word) + except KeyboardInterrupt: + # If the user hits Ctrl+C, we don't want to print + # a traceback to the user. + pass + + # Close log file if it exists + if LOG_FILE: + LOG_FILE.close() diff --git a/src/awssh/awssh.py b/src/awssh/awssh.py new file mode 100644 index 0000000..6f16846 --- /dev/null +++ b/src/awssh/awssh.py @@ -0,0 +1,152 @@ +"""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] [...] + 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 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[""] + credential_file: Optional[Path] = None + if validated_args["--credentials"]: + credential_file = CREDENTIAL_DIR / Path(validated_args["--credentials"]) + instance_id: str = validated_args[""] + log_level: str = validated_args["--log-level"] + no_ssh: bool = validated_args["--no-ssh"] + profile: str = validated_args[""] + 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)) + completed_process = ( + subprocess.run( # nosec: B603 subprocess input is carefully validated + args=command, env=awssh_env + ) + ) + return completed_process.returncode diff --git a/src/example/data/secret.txt b/src/example/data/secret.txt deleted file mode 100644 index c40a49b..0000000 --- a/src/example/data/secret.txt +++ /dev/null @@ -1 +0,0 @@ -Three may keep a secret, if two of them are dead. diff --git a/src/example/example.py b/src/example/example.py deleted file mode 100644 index d3eda19..0000000 --- a/src/example/example.py +++ /dev/null @@ -1,103 +0,0 @@ -"""example is an example Python library and tool. - -Divide one integer by another and log the result. Also log some information -from an environment variable and a package resource. - -EXIT STATUS - This utility exits with one of the following values: - 0 Calculation completed successfully. - >0 An error occurred. - -Usage: - example [--log-level=LEVEL] - example (-h | --help) - -Options: - -h --help Show this message. - --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] -""" - -# Standard Python Libraries -import logging -import os -import sys -from typing import Any, Dict - -# Third-Party Libraries -import docopt -import pkg_resources -from schema import And, Schema, SchemaError, Use - -from ._version import __version__ - -DEFAULT_ECHO_MESSAGE: str = "Hello World from the example default!" - - -def example_div(dividend: int, divisor: int) -> float: - """Print some logging messages.""" - logging.debug("This is a debug message") - logging.info("This is an info message") - logging.warning("This is a warning message") - logging.error("This is an error message") - logging.critical("This is a critical message") - return dividend / divisor - - -def main() -> None: - """Set up logging and call the example function.""" - 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.", - ), - "": Use(int, error=" must be an integer."), - "": And( - Use(int), - lambda n: n != 0, - error=" must be an integer that is not 0.", - ), - 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 - dividend: int = validated_args[""] - divisor: int = validated_args[""] - log_level: str = validated_args["--log-level"] - - # Set up logging - logging.basicConfig( - format="%(asctime)-15s %(levelname)s %(message)s", level=log_level.upper() - ) - - logging.info("%d / %d == %f", dividend, divisor, example_div(dividend, divisor)) - - # Access some data from an environment variable - message: str = os.getenv("ECHO_MESSAGE", DEFAULT_ECHO_MESSAGE) - logging.info('ECHO_MESSAGE="%s"', message) - - # Access some data from our package data (see the setup.py) - secret_message: str = ( - pkg_resources.resource_string("example", "data/secret.txt") - .decode("utf-8") - .strip() - ) - logging.info('Secret="%s"', secret_message) - - # Stop logging and clean up - logging.shutdown() diff --git a/tests/test_example.py b/tests/test_awssh.py similarity index 67% rename from tests/test_example.py rename to tests/test_awssh.py index f8dea67..7cab23c 100644 --- a/tests/test_example.py +++ b/tests/test_awssh.py @@ -1,5 +1,5 @@ #!/usr/bin/env pytest -vs -"""Tests for example.""" +"""Tests for awssh.""" # Standard Python Libraries import logging @@ -11,14 +11,7 @@ import pytest # cisagov Libraries -import example - -div_params = [ - (1, 1, 1), - (2, 2, 1), - (0, 1, 0), - (8, 2, 4), -] +from awssh import awssh log_levels = ( "debug", @@ -30,14 +23,14 @@ # define sources of version strings RELEASE_TAG = os.getenv("RELEASE_TAG") -PROJECT_VERSION = example.__version__ +PROJECT_VERSION = awssh.__version__ def test_stdout_version(capsys): """Verify that version string sent to stdout agrees with the module version.""" with pytest.raises(SystemExit): with patch.object(sys, "argv", ["bogus", "--version"]): - example.example.main() + awssh.main() captured = capsys.readouterr() assert ( captured.out == f"{PROJECT_VERSION}\n" @@ -54,7 +47,7 @@ def test_running_as_module(capsys): # package and running it, so there is nothing to use from this # import. As a result, we can safely ignore this warning. # cisagov Libraries - import example.__main__ # noqa: F401 + import awssh.__main__ # noqa: F401 captured = capsys.readouterr() assert ( captured.out == f"{PROJECT_VERSION}\n" @@ -81,7 +74,7 @@ def test_log_levels(level): ), "root logger should not have handlers yet" return_code = None try: - example.example.main() + awssh.main() except SystemExit as sys_exit: return_code = sys_exit.code assert return_code is None, "main() should return success" @@ -99,46 +92,7 @@ def test_bad_log_level(): with patch.object(sys, "argv", ["bogus", "--log-level=emergency", "1", "1"]): return_code = None try: - example.example.main() - except SystemExit as sys_exit: - return_code = sys_exit.code - assert return_code == 1, "main() should exit with error" - - -@pytest.mark.parametrize("dividend, divisor, quotient", div_params) -def test_division(dividend, divisor, quotient): - """Verify division results.""" - result = example.example_div(dividend, divisor) - assert result == quotient, "result should equal quotient" - - -@pytest.mark.slow -def test_slow_division(): - """Example of using a custom marker. - - This test will only be run if --runslow is passed to pytest. - Look in conftest.py to see how this is implemented. - """ - # Standard Python Libraries - import time - - result = example.example_div(256, 16) - time.sleep(4) - assert result == 16, "result should equal be 16" - - -def test_zero_division(): - """Verify that division by zero throws the correct exception.""" - with pytest.raises(ZeroDivisionError): - example.example_div(1, 0) - - -def test_zero_divisor_argument(): - """Verify that a divisor of zero is handled as expected.""" - with patch.object(sys, "argv", ["bogus", "1", "0"]): - return_code = None - try: - example.example.main() + awssh.main() except SystemExit as sys_exit: return_code = sys_exit.code assert return_code == 1, "main() should exit with error" diff --git a/tools/awssh-completion.bash b/tools/awssh-completion.bash new file mode 100644 index 0000000..18d4871 --- /dev/null +++ b/tools/awssh-completion.bash @@ -0,0 +1,3 @@ +#!/usr/bin/env bash + +complete -C '_awssh-completer' awssh