-
Notifications
You must be signed in to change notification settings - Fork 6
/
awser.py
executable file
·144 lines (112 loc) · 3.62 KB
/
awser.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
#!/usr/bin/env python3
"""Quickly SSH to EC2 instances by name
"""
import sys
import argparse
import readline # used by input() for more intuitive editing of input
from subprocess import call
import boto3
session = boto3.session.Session()
AWS_DEFAULT_REGION = session.region_name or 'us-west-2'
SSH_CONNECT_TIMEOUT = 10
def parse_args():
"""
Parse arguments passed to the command
"""
parser = argparse.ArgumentParser(
description="Quickly SSH to EC2 instances by name")
parser.add_argument(
'keywords',
nargs="+",
help="Keyword(s) to filter list of servers, i.e. 'core 21a'")
parser.add_argument(
'-r',
'--region',
default=AWS_DEFAULT_REGION,
help="Overrides the default AWS region.")
parser.add_argument(
'-p', '--profile', default=None, help="Specifies which profile to use.")
parser.add_argument(
'-u', '--user', default=None, help="Specifies a user for SSH.")
parser.add_argument(
'-i',
'--identity',
default=None,
help="Selects a file from which the identity (private key) is read.")
parser.add_argument(
'-d',
'--public-dns',
action='store_true',
help="Use public DNS name instead of IP address to connect.")
return parser.parse_args()
def get_hosts(name_filter, region):
"""Retrieve list of hosts from AWS and parse to a reasonable data structure
"""
hosts = []
client = boto3.client('ec2', region_name=region)
response = client.describe_instances(
Filters=[{
'Name': 'tag:Name',
'Values': [name_filter]
}, {
'Name': 'instance-state-name',
'Values': ['running']
}])
for r in response['Reservations']:
for i in r['Instances']:
hosts.append([
i,
next(x['Value'] for x in i['Tags'] if x['Key'] == 'Name')
])
return hosts
def ssh(ip, user, identity):
"""Helper method to SSH to a host
"""
ssh_args = [
"-o", "StrictHostKeyChecking=no", "-o",
"ConnectTimeout=%s" % SSH_CONNECT_TIMEOUT
]
if identity:
ssh_args += ["-i", identity]
if user:
ssh_args += [user + "@" + ip]
else:
ssh_args += [ip]
return call(["ssh"] + ssh_args)
def main():
"""Main function
"""
args = parse_args()
if args.profile:
boto3.setup_default_session(profile_name=args.profile)
name_filter = "*"
for a in args.keywords:
name_filter += a + "*"
hosts = get_hosts(name_filter, args.region)
if not hosts:
sys.exit("No hosts found matching those keywords.")
host_name_key = 'PrivateIpAddress'
if args.public_dns:
host_name_key = 'PublicDnsName'
if len(hosts) == 1:
print("Logging in to %s..." % hosts[0][1])
ssh(hosts[0][0][host_name_key], args.user, args.identity)
else:
choice = None
for i, host in enumerate(hosts, 1):
print("{0}) {1} - {2}".format(i, host[1], host[0][host_name_key]))
print("Enter) exit")
try:
choice = int(input("Choose wisely: "))
except (SyntaxError, NameError, ValueError):
sys.exit("You chose... poorly")
if choice == 0 or choice > len(hosts):
sys.exit("You chose... poorly")
else:
print("Logging in to %s..." % hosts[choice - 1][1])
ssh(hosts[choice - 1][0][host_name_key], args.user, args.identity)
if __name__ == "__main__":
try:
main()
except KeyboardInterrupt:
sys.exit("\nReceived SIGINT, exiting...")