bergundy · May 10, 2018 · May 11, 2018 · May 11, 2018 · Oct 9, 2018 · Oct 9, 2018
Showing with 456 additions and 118 deletions.

+43 −0 .gitignore

+2 −0 Dockerfile

+13 −0 Jenkinsfile

+26 −0 Makefile

+13 −3 README.md

+1 −0 VERSION

+209 −91 ec2grep/__init__.py

+123 −0 poetry.lock

+26 −0 pyproject.toml

+0 −24 setup.py
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,43 @@
+.idea
+*.iml
+out
+outfds
+**/.bundle/
+*COMMIT_MSG
+
+*.py[cod]
+
+# C extensions
+*.so
+*.iml
+
+# Packages
+*.egg
+*.egg-info
+dist
+build
+eggs
+parts
+bin
+var
+sdist
+develop-eggs
+.installed.cfg
+lib
+lib64
+
+# Installer logs
+pip-log.txt
+
+# Unit test / coverage reports
+.coverage
+.tox
+nosetests.xml
+
+# Translations
+*.mo
+
+# Mr Developer
+.mr.developer.cfg
+.project
+.pydevproject
diff --git a/Dockerfile b/Dockerfile
@@ -0,0 +1,2 @@
+ARG PYTHON_VERSION=3.8
+FROM 174522763890.dkr.ecr.us-east-1.amazonaws.com/ubuntu-python-pyproject:${PYTHON_VERSION}
diff --git a/Jenkinsfile b/Jenkinsfile
@@ -0,0 +1,13 @@
+// vim: filetype=groovy
+node('general') {
+    def base = load('/var/jenkins_home/workspace/Infra/build-scripts/build/Jenkinsfile')
+    base.execute([
+        customStages: base.get_pypi_stages(
+            /* is_poetry */true,
+            /* python_version */'3.8',
+            /* should_bump_version */true,
+            /* additional_package_roots */null,
+            /* is_wheel */true,
+        ),
+    ])
+}
diff --git a/Makefile b/Makefile
@@ -0,0 +1,26 @@
+ECR_URL ?= 174522763890.dkr.ecr.us-east-1.amazonaws.com
+
+.PHONY: build
+build: 
+	@echo "======= Building Python ========"
+	docker build --pull .
+
+ci-test:
+	echo "No tests"
+
+dist:
+	echo "Actual dist is done in the post-dist stage"
+
+ci-lint:
+	@echo "======= Linting ========"
+	docker run \
+		-v $(PWD)/:/app/ \
+		-v $(PWD)/test-reports:/test-reports/ \
+		${ECR_URL}/python-black
+
+black:
+	docker run \
+		-v $(PWD)/:/app/ \
+		--entrypoint black \
+		${ECR_URL}/python-black \
+		/app/
diff --git a/README.md b/README.md
@@ -1,26 +1,30 @@
 ## ec2grep - EC2 cli tool
 
 ### Install
-Preferally for the default system interpreter instead of a virtualenv.
+
+Make sure you have `pipx` available in your system (`brew install pipx` if not).
 
 ```bash
-pip install git+https://github.com/bergundy/ec2grep
+pipx install ec2grep --index-url https://artifactory.frdstr.com/artifactory/api/pypi/pypi/simple
 ```
 
 ### Usage
 
 ##### ls
 Basic usage, find by name tag, external / internal IP, DNS
+
 ```bash
 ec2 ls my-hostname
 ```
 
-Custom formatter (name, ip, extended)
+Using a custom formatter (name, ip, extended)
+
 ```bash
 ec2 ls --format=name my-hostname
 ```
 
 ##### ssh
+
 Open an SSH session
 ```bash
 ec2 ssh my-hostname
@@ -31,7 +35,13 @@ With arguments
 ec2 ssh my-hostname -- w
 ```
 
+Parallel command
+```bash
+ec2 ssh --parallel -F ~/.ssh/gozer my-hostname -- "<command>"
+```
+
 ##### custom region
+
 ```bash
 ec2 --region us-west-2 ls my-hostname
 ```
diff --git a/VERSION b/VERSION
@@ -0,0 +1 @@
+0.1
diff --git a/ec2grep/__init__.py b/ec2grep/__init__.py
@@ -1,126 +1,244 @@
-#!/usr/bin/env python
-from concurrent.futures import wait, ThreadPoolExecutor as Executor
-from functools import partial
-from operator import itemgetter
+#! /usr/bin/env python
 
+import concurrent.futures
+import functools
+import itertools
 import os
+import pprint
+import subprocess
 import sys
-import click
+import threading
+
 import boto3
-import itertools
+import click
 import six
 
-executor = Executor(4)
-name = (lambda i: {tag['Key']: tag['Value'] for tag in i.get('Tags', [])}.get('Name', ''))
-public_ip = itemgetter('PublicIpAddress')
-private_ip = itemgetter('PrivateIpAddress')
-extended_public = (lambda i: "{} ({})".format(name(i), public_ip(i)))
-extended_private = (lambda i: "{} ({})".format(name(i), private_ip(i)))
-DEFAULT_ATTRIBUTES = [
-    'tag:Name',
-    'network-interface.addresses.association.public-ip',
-    'network-interface.addresses.private-ip-address',
-    'network-interface.private-dns-name'
-]
-
-
-def _get_instances(ec2, filter_):
-    response = ec2.describe_instances(Filters=[filter_])
-    reservations = response['Reservations']
-    return list(itertools.chain.from_iterable(r['Instances'] for r in reservations))
-
-
-def match_instances(region_name, query, attributes=DEFAULT_ATTRIBUTES):
-    ec2 = boto3.client('ec2', region_name=region_name)
-    get_instances = partial(_get_instances, ec2)
-    instance_lists = executor.map(get_instances, [
-        {'Name': attr, 'Values': ['*{}*'.format(query)]} for attr in attributes
-    ])
-    chained = (i for i in itertools.chain.from_iterable(instance_lists) if 'PublicIpAddress' in i)
-    return sorted(chained, key=name)
-
-
-def die(*args):
-    click.echo(*args, err=True)
-    sys.exit(1)
-
-
-def read_number(min_value, max_value):
-    while True:
-        try:
-            choice = six.moves.input()
-            choice = int(choice)
-            if not (min_value <= choice <= max_value):
-                raise ValueError("Invalid input")
-            return choice
-        except ValueError as e:
-            click.echo("{}".format(e), err=True)
-            continue
-
-
+DEFAULT_ATTRIBUTES = (
+    "tag:Name",
+    "network-interface.addresses.association.public-ip",
+    "network-interface.addresses.private-ip-address",
+    "network-interface.private-dns-name",
+)
+
+
+name = lambda i: {tag["Key"]: tag["Value"] for tag in i.get("Tags", [])}.get("Name", "")
+initiator = lambda i: {tag["Key"]: tag["Value"] for tag in i.get("Tags", [])}.get("Initiator", "")
+private_ip = lambda i: i.get("PrivateIpAddress", None)
+public_ip = lambda i: i.get("PublicIpAddress", None)
+extended_public = lambda i: "{} ({})".format(name(i), public_ip(i))
+extended_private = lambda i: "{} ({})".format(name(i), private_ip(i))
+extended = lambda i: "{} (public: {}, private: {})".format(name(i), public_ip(i), private_ip(i))
+extended_initiator = lambda i: "{} (name: {})".format(initiator(i), name(i))
 formatters = {
-    'extended_public': extended_public,
-    'extended_private': extended_private,
-    'public_ip': public_ip,
-    'private_ip': private_ip,
-    'name': name
+    "extended": extended,
+    "extended_public": extended_public,
+    "extended_private": extended_private,
+    "public_ip": public_ip,
+    "private_ip": private_ip,
+    "name": name,
+    "initiator": initiator,
+    "extended_initiator": extended_initiator,
 }
 
 
 @click.group()
-@click.option('--region', '-r', default='us-east-1')
+@click.option("--region", "-r", default="us-east-1")
 @click.pass_context
 def cli(ctx, region):
-    ctx.obj = {'region': region}
+    ctx.obj = {"region": region}
+
+
+@cli.command()
+@click.option("--user", "-u")
+@click.option("--key", "-i", help="SSH key")
+@click.option("--ssh-config", "-F", help="SSH config file")
+@click.option("--prefer-public-ip", "-p", is_flag=True, default=False)
+@click.option("--parallel", is_flag=True, default=False)
+@click.argument("query")
+@click.argument("ssh_args", nargs=-1, type=click.UNPROCESSED)
+@click.pass_context
+def ssh(ctx, user, key, ssh_config, prefer_public_ip, parallel, query, ssh_args):
+    validate_identity_parameters(user, key, ssh_config)
+
+    def op(host, host_count):
+        command = ["ssh", "-oStrictHostKeyChecking=no", "-t"]
+        command.extend(["-i", key]) if key else command.extend(["-F", ssh_config])
+        if user:
+            command.extend(["-l", user])
+        command.extend([host] + list(ssh_args))
+        return command
+
+    run_op_for_hosts(ctx, prefer_public_ip, parallel, query, "ssh", op)
 
 
 @cli.command()
-@click.option('--key', '-i')
-@click.option('--prefer-public-ip', '-p', is_flag=True, default=False)
-@click.argument('query')
-@click.argument('ssh_args', nargs=-1, type=click.UNPROCESSED)
+@click.option("--user", "-u")
+@click.option("--key", "-i", help="SSH key")
+@click.option("--ssh-config", "-F", help="SSH config file")
+@click.option("--prefer-public-ip", "-p", is_flag=True, default=False)
+@click.option("--parallel", is_flag=True, default=False)
+@click.option("--download", "-d", is_flag=True, default=False)
+@click.argument("query")
+@click.argument("source", type=str, required=True)
+@click.argument("target", type=str, required=False)
+@click.argument("scp_args", nargs=-1, type=click.UNPROCESSED)
 @click.pass_context
-def ssh(ctx, key, prefer_public_ip, query, ssh_args):
+def scp(ctx, user, key, ssh_config, prefer_public_ip, parallel, query, download, source, target, scp_args):
+    validate_identity_parameters(user, key, ssh_config)
+
+    if download and not os.path.isdir(target):
+        die("Download target must be an existing directory")
+
+    def format_scp_remote_host_path(user, host, path):
+        return "{}@{}:{}".format(user, host, path) if user else "{}:{}".format(host, path)
+
+    def op(host, host_count):
+        command = ["scp", "-oStrictHostKeyChecking=no"]
+        command.extend(["-i", key]) if key else command.extend(["-F", ssh_config])
+        command.extend(list(scp_args))
+        if download:
+            if host_count > 1:
+                # when downloading from multiple hosts, use a specific directory per each host
+                host_specific_dir_name = host.replace(".", "-")
+                full_target = os.path.join(target, host_specific_dir_name)
+                os.mkdir(full_target)
+            else:
+                full_target = target
+            command.extend([format_scp_remote_host_path(user, host, source), full_target])
+        else:
+            command.extend([source, format_scp_remote_host_path(user, host, target)])
+        return command
+
+    run_op_for_hosts(ctx, prefer_public_ip, parallel, query, "scp", op)
+
+
+def run_op_for_hosts(ctx, prefer_public_ip, parallel, query, op_name, op):
     get_ip = public_ip if prefer_public_ip else private_ip
     fmt_match = extended_public if prefer_public_ip else extended_private
-    extra_args = []
 
-    matches = match_instances(ctx.obj['region'], query)
+    choices = get_hosts_choice(ctx, fmt_match, query)
+
+    if parallel and len(choices) > 1:
+        message = "running {}:\n{}".format(op_name, pprint.pformat([fmt_match(c) for c in choices]))
+        click.echo(message)
+        click.echo()
+        processes = []
+        for choice in choices:
+            ip = get_ip(choice)
+            p = subprocess.Popen(op(ip, len(choices)), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            stdout_consumer = OutputConsumer(input=p.stdout, output=sys.stdout, ip=ip)
+            stderr_consumer = OutputConsumer(input=p.stderr, output=sys.stderr, ip=ip)
+            processes.append({"out": stdout_consumer, "err": stderr_consumer, "p": p})
+        for p in processes:
+            p["p"].wait()
+            p["out"].join()
+            p["err"].join()
+    else:
+        for choice in choices:
+            message = "running {} {}".format(op_name, fmt_match(choice))
+            click.echo()
+            click.echo(message)
+            click.echo(u"\u2500" * len(message))
+            subprocess.call(op(get_ip(choice), len(choices)))
+
+
+def get_hosts_choice(ctx, fmt_match, query):
+    matches = match_instances(ctx.obj["region"], query)
     if not matches:
         die("No matches found")
-
     if len(matches) > 1:
+        click.echo("[0] All")
         for i, inst in enumerate(matches):
-            click.echo("[{}] {}".format(i+1, fmt_match(inst)))
-        click.echo("pick an option [1-{}] ".format(len(matches)), nl=False)
-        index = read_number(1, len(matches)) - 1
-        choice = matches[index]
-        click.echo("")
+            click.echo("[{}] {}".format(i + 1, fmt_match(inst)))
+        click.echo("select servers [0-{}] (comma separated) ".format(len(matches)), nl=False)
+        indices = read_numbers(0, len(matches))
+        if len(indices) == 1 and indices[0] == 0:
+            choices = matches
+        else:
+            choices = [matches[index - 1] for index in indices]
+        click.echo()
     else:
-        choice = matches[0]
-
-    click.echo("sshing {}".format(fmt_match(choice)))
-
-    if key:
-        extra_args.extend(['-i', key])
-    os.execvp('ssh', ['ssh', '-oStrictHostKeyChecking=no'] + extra_args + [get_ip(choice)] + list(ssh_args))
+        choices = [matches[0]]
+    return choices
 
 
 @cli.command()
-@click.argument('query')
-@click.option('--delim', '-d', default='\n')
-@click.option('--formatter', '-f', default='extended_private')
-@click.option('--custom-format', '-c', is_flag=True, default=False)
+@click.option("--delim", "-d", default="\n")
+@click.option("--formatter", "-f", default="extended_private")
+@click.option("--custom-format", "-c", is_flag=True, default=False)
+@click.argument("query")
 @click.pass_context
-def ls(ctx, query, formatter, delim, custom_format):
-    matches = match_instances(ctx.obj['region'], query)
-    if not custom_format:
-        formatter = formatter.join('{}')
+def ls(ctx, formatter, delim, custom_format, query):
+    matches = match_instances(ctx.obj["region"], query)
     if not matches:
         die("No matches found")
+    if not custom_format:
+        formatter = formatter.join("{}")
+    click.echo(delim.join(formatter.format(**{k: f(m) for k, f in formatters.items()}) for m in matches))
+
+
+def _get_instances(ec2, filter_):
+    response = ec2.describe_instances(Filters=[filter_])
+    reservations = response["Reservations"]
+    return list(itertools.chain.from_iterable(r["Instances"] for r in reservations))
+
+
+def match_instances(region_name, query, attributes=DEFAULT_ATTRIBUTES):
+    ec2 = boto3.client("ec2", region_name=region_name)
+    get_instances = functools.partial(_get_instances, ec2)
+    with concurrent.futures.ThreadPoolExecutor(len(attributes)) as executor:
+        instance_lists = executor.map(
+            get_instances, [{"Name": attr, "Values": ["*{}*".format(query)]} for attr in attributes]
+        )
+    chained = (
+        i for i in itertools.chain.from_iterable(instance_lists) if "PublicIpAddress" in i or "PrivateIpAddress" in i
+    )
+    return sorted(chained, key=name)
+
+
+def die(*args):
+    click.echo(*args, err=True)
+    sys.exit(1)
+
+
+def read_numbers(min_value, max_value):
+    while True:
+        try:
+            choices = six.moves.input()
+            choices = [int(c.strip()) for c in choices.split(",")]
+            if not (all(min_value <= c <= max_value for c in choices)):
+                raise ValueError("Invalid input")
+            if len(choices) != 1 and 0 in choices:
+                raise ValueError("Invalid input")
+            return choices
+        except ValueError as e:
+            click.echo(str(e), err=True)
+            continue
+
+
+def validate_identity_parameters(user, key, ssh_config):
+    if not key and not ssh_config:
+        die("Must supply either SSH key or SSH config file")
+    if ssh_config and key:
+        die("The ssh-config option is mutually exclusive with the key option")
+
+
+class OutputConsumer(object):
+    def __init__(self, input, output, ip):
+        self.ip = ip
+        self.input = input
+        self.output = output
+        self.consumer = threading.Thread(target=self.consume_output)
+        self.consumer.daemon = True
+        self.consumer.start()
+
+    def consume_output(self):
+        for line in iter(self.input.readline, b""):
+            self.output.write("[{:<15}] {}".format(self.ip, line))
 
-    click.echo(delim.join(formatter.format(**{k: f(m) for k, f in formatters.iteritems()}) for m in matches))
+    def join(self):
+        self.consumer.join()
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     cli()
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -0,0 +1,26 @@
+[tool.poetry]
+name = "ec2grep"
+version = "0.1.0"
+description = "EC2 cli tool"
+authors = ["Roey Berman <roey.berman@gmail.com>"]
+keywords = ["ec2", "cli", "aws", "ssh"]
+
+packages = [
+    { include = "ec2grep" }
+]
+
+[tool.poetry.scripts]
+ec2 = "ec2grep:cli"
+
+[tool.poetry.dependencies]
+python = "^3.8"
+
+boto3 = "^1.15.18"
+click = "^7.1.2"
+
+[tool.black]
+line-length = 120
+
+[build-system]
+requires = ["poetry>=1.1.3"]
+build-backend = "poetry.masonry.api"
diff --git a/setup.py b/setup.py
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		ARG PYTHON_VERSION=3.8
		FROM 174522763890.dkr.ecr.us-east-1.amazonaws.com/ubuntu-python-pyproject:${PYTHON_VERSION}