def test_mssh_no_target_no_public_ip(self, mock_instance_data,
mock_push_key, mock_run):
mock_file = "identity"
flag = '-f flag'
command = 'command arg'
logger = EC2InstanceConnectLogger()
instance_bundles = [{
'username': self.default_user,
'instance_id': self.instance_id,
'target': None,
'zone': self.availability_zone,
'region': self.region,
'profile': self.profile
}]
mock_instance_data.return_value = self.private_instance_info
mock_push_key.return_value = None
cli_command = EC2InstanceConnectCommand("ssh", instance_bundles,
mock_file, flag, command,
logger.get_logger())
cli = EC2InstanceConnectCLI(instance_bundles, "", cli_command,
logger.get_logger())
cli.invoke_command()
expected_command = 'ssh -o "IdentitiesOnly=yes" -i {0} {1} {2}@{3} {4}'.format(
mock_file, flag, self.default_user, self.private_ip, command)
# Check that we successfully get to the run
self.assertTrue(mock_instance_data.called)
self.assertTrue(mock_push_key.called)
mock_run.assert_called_with(expected_command)
def test_mssh_with_target(self, mock_instance_data, mock_push_key,
mock_run):
mock_file = 'identity'
flag = '-f flag'
command = 'command arg'
host = '0.0.0.0'
logger = EC2InstanceConnectLogger()
instance_bundles = [{
'username': self.default_user,
'instance_id': self.instance_id,
'target': host,
'zone': self.availability_zone,
'region': self.region,
'profile': self.profile
}]
mock_instance_data.return_value = self.instance_info
mock_push_key.return_value = None
cli_command = EC2InstanceConnectCommand("ssh", instance_bundles,
mock_file, flag, command,
logger.get_logger())
cli = EC2InstanceConnectCLI(instance_bundles, "", cli_command,
logger.get_logger())
cli.invoke_command()
expected_command = 'ssh -o "IdentitiesOnly=yes" -i {0} {1} {2}@{3} {4}'.format(
mock_file, flag, self.default_user, host, command)
# Check that we successfully get to the run
# Since both target and availability_zone are provided, mock_instance_data should not be called
self.assertFalse(mock_instance_data.called)
self.assertTrue(mock_push_key.called)
mock_run.assert_called_with(expected_command)
def test_mscp(self,
mock_instance_data,
mock_push_key,
mock_run):
mock_file = 'identity'
flag = '-f flag'
command = 'file2 file3'
logger = EC2InstanceConnectLogger()
instance_bundles = [{'username': self.default_user, 'instance_id': self.instance_id,
'target': None, 'zone': self.availability_zone, 'region': self.region,
'profile': self.profile, 'file': 'file1'},
{'username': self.default_user, 'instance_id': self.instance_id,
'target': None, 'zone': self.availability_zone, 'region': self.region,
'profile': self.profile, 'file': 'file4'}]
mock_instance_data.return_value = self.instance_info
mock_push_key.return_value = None
expected_command = "scp -i {0} {1} {2}@{3}:{4} {5} {6}@{7}:{8}".format(mock_file, flag, self.default_user,
self.public_ip, 'file1', command,
self.default_user,
self.public_ip, 'file4')
cli_command = EC2InstanceConnectCommand("scp", instance_bundles, mock_file, flag, command, logger.get_logger())
cli = EC2InstanceConnectCLI(instance_bundles, "", cli_command, logger.get_logger())
cli.invoke_command()
# Check that we successfully get to the run
self.assertTrue(mock_instance_data.called)
self.assertTrue(mock_push_key.called)
mock_run.assert_called_with(expected_command)
def main(program, mode):
"""
Parses system arguments and sets defaults
Calls `ssh` or `sftp` to SSH into the Instance or transfer files.
:param program: Client program to be used for SSH/SFTP operations.
:type program: basestring
:param mode: Identifies either SSH/SFTP operation.
:type mode: basestring
"""
usage = ""
if mode == "ssh":
usage="""
mssh [-t instance_id] [-u profile] [-z availability_zone] [-r region] [supported ssh flags] target [command]
target => [user@]instance_id | [user@]hostname
[supported ssh flags] => [-l login_name] [-p port]
"""
elif mode == "sftp":
usage="""
msftp [-u aws_profile] [-z availability_zone] [supported sftp flags] target
target => [user@]instance_id[:file ...][:dir[/]] | [user@]hostname[:file ...][:dir[/]]
[supported sftp flags] => [-P port] [-b batchfile]
"""
parser = argparse.ArgumentParser(usage=usage)
parser.add_argument('-r', '--region', action='store', help='AWS region', type=str, metavar='')
parser.add_argument('-z', '--zone', action='store', help='Availability zone', type=str, metavar='')
parser.add_argument('-u', '--profile', action='store', help='AWS Config Profile', type=str, default=DEFAULT_PROFILE, metavar='')
parser.add_argument('-t', '--instance_id', action='store', help='EC2 Instance ID. Required if target is hostname', type=str, default=DEFAULT_INSTANCE, metavar='')
parser.add_argument('-d', '--debug', action="store_true", help='Turn on debug logging')
args = parser.parse_known_args()
logger = EC2InstanceConnectLogger(args[0].debug)
try:
instance_bundles, flags, program_command = input_parser.parseargs(args, mode)
except Exception as e:
print(str(e))
parser.print_help()
sys.exit(1)
#Generate temp key
cli_key = EC2InstanceConnectKey(logger.get_logger())
cli_command = EC2InstanceConnectCommand(program, instance_bundles, cli_key.get_priv_key_file(), flags, program_command, logger.get_logger())
try:
# TODO: Handling for if the '-i' flag is passed
cli = EC2InstanceConnectCLI(instance_bundles, cli_key.get_pub_key(), cli_command, logger.get_logger())
cli.invoke_command()
except Exception as e:
print('Failed with:\n' + str(e))
sys.exit(1)
def main(program, mode):
"""
Parses system arguments sets defaults
Calls `putty or psftp` to SSH or do file operations using EC2InstanceConnect.
:param program: Client program to be used for SSH/SFTP operations.
:type program: basestring
:param mode: Identifies either SSH/SFTP operation.
:type mode: basestring
"""
usage = ""
if mode == "ssh":
usage = """
mssh-putty [-t instance_id] [-u profile] [-r region] [-z availability_zone] [-i identity_file_ppk] [supported putty flags] target
target => [user@]instance_id | [user@]dns_name
[supported putty flags] => [-l login_name] [ -P port] [-m remote_command_file]
"""
elif mode == "sftp":
usage = """
msftp-putty [-t instance_id] [-u profile] [-r region] [-z availability_zone] [-i identity_file_ppk] [supported psftp flags] target
target => [user@]instance_id | [user@]dns_name
[supported psftp flags] => [-l login_name] [ -P port] [-bc] [-b batchfile]
"""
parser = argparse.ArgumentParser(usage=usage)
parser.add_argument('-r',
'--region',
action='store',
help='AWS region',
type=str,
metavar='')
parser.add_argument('-z',
'--zone',
action='store',
help='Availability zone',
type=str,
metavar='')
parser.add_argument('-i',
'--identity',
action='store',
help="Required. Identity file in ppk format",
type=str,
required=True,
metavar='')
parser.add_argument('-u',
'--profile',
action='store',
help='AWS Config Profile',
type=str,
default=DEFAULT_PROFILE,
metavar='')
parser.add_argument(
'-t',
'--instance_id',
action='store',
help='EC2 Instance ID. Required if target is hostname',
type=str,
default=DEFAULT_INSTANCE,
metavar='')
parser.add_argument('-d',
'--debug',
action="store_true",
help='Turn on debug logging')
args = parser.parse_known_args()
#Read public key from ppk file.
#Public key is unencrypted and in rsa format.
pub_key_lines = []
pub_key = "ssh-rsa "
try:
with open(args[0].identity, 'r') as f:
pub_key_lines = f.readlines()
#Validate that the identity file format is ppk.
if pub_key_lines[0].find("PuTTY-User-Key-File-") == -1:
print("Not a valid Putty key.")
sys.exit(1)
#public key starts from 4th line in ppk file.
line_len = int(pub_key_lines[3].split(':')[1].strip())
pub_key_lines = pub_key_lines[4:(4 + line_len)]
for pub_key_line in pub_key_lines:
pub_key += pub_key_line[:-1]
except Exception as e:
print(str(e))
sys.exit(1)
logger = EC2InstanceConnectLogger(args[0].debug)
try:
instance_bundles, flags, program_command = input_parser.parseargs(
args, mode)
except Exception as e:
print(str(e))
parser.print_help()
sys.exit(1)
cli_command = EC2InstanceConnectCommand(program, instance_bundles,
args[0].identity,
flags, program_command,
logger.get_logger())
cli_command.get_command()
try:
mssh = EC2InstanceConnectCLI(instance_bundles, pub_key, cli_command,
logger.get_logger())
mssh.invoke_command()
except Exception as e:
print("Failed with:\n" + str(e))
sys.exit(1)