Merging changes from internal development, mostly focuses around debug handling and session management

Author: LordAlfredo

ec2instanceconnectcli/EC2InstanceConnectCLI.py

@@ -1,3 +1,4 @@
+import logging
 import sys
 import time
 from subprocess import Popen
@@ -86,7 +87,11 @@ def invoke_command(self):
         """
         try:
             for bundle in self.instance_bundles:
-                bundle['session'] = self._get_botocore_session(profile_name=bundle['profile'], region=bundle['region'])
+                session = self._get_botocore_session(profile_name=bundle['profile'], region=bundle['region'])
+                # enable debug logging on botocore session if command line debug option is set
+                if self.logger.getEffectiveLevel() == logging.DEBUG:
+                    session.set_debug_logger()
+                bundle['session'] = session
 
             self.call_ec2()
             self.handle_keys()

ec2instanceconnectcli/EC2InstanceConnectLogger.py

@@ -5,14 +5,14 @@ class EC2InstanceConnectLogger(object):
     Common logger for all the EC2InstanceConnect components.
     """
 
-    def __init__(self, verbose=None):
+    def __init__(self, debug=False):
         """
-        :param verbose: Specifies if debug messages should be enabled
-        :type verbose: bool or None
+        :param debug: Specifies if debug messages should be enabled
+        :type debug: bool
         """
         self.logger = logging.getLogger('EC2InstanceConnect')
         log_level = logging.ERROR
-        if verbose:
+        if debug:
             log_level = logging.DEBUG
         self.logger.setLevel(log_level)
         ch = logging.StreamHandler()

ec2instanceconnectcli/input_parser.py

@@ -35,9 +35,9 @@ def parseargs(args, mode='ssh'):
     """
 
     if len(args) < 2:
-        raise AssertionError('No target was provided')
+        raise AssertionError('Missing target')
     if len(args[1]) < 1:
-        raise AssertionError('No target was provided')
+        raise AssertionError('Missing target')
 
     """
     Our flags.  As these are via argparse they're free.
@@ -61,8 +61,22 @@ def parseargs(args, mode='ssh'):
     # Process the instance & target data to give us an actual picture of what end hosts we're working with
     instance_bundles = _parse_instance_bundles(instance_bundles)
 
+    #validate instance_bundles
+    _validate_instance_bundles(instance_bundles, mode)
+
+    #validate instance_bundles
+    _validate_instance_bundles(instance_bundles, mode)
+
     return instance_bundles, flags, command
 
+def _validate_instance_bundles(instance_bundles, mode):
+    """
+    For supported modes ensure that instance_id exists.
+    """
+    for bundle in instance_bundles:
+        if mode in ['ssh', 'sftp']:
+            if not INSTANCE_ID_RE.match(bundle['instance_id']):
+                raise AssertionError('Missing instance_id')
 
 def _parse_command_flags(raw_command, instance_bundles, is_ssh=False):
     """
@@ -123,7 +137,7 @@ def _parse_command_flags(raw_command, instance_bundles, is_ssh=False):
 
     if used == len(raw_command) and len(raw_command) != 1:
         # EVERYTHING was a flag or flag value
-        raise AssertionError('No target was provided')
+        raise AssertionError('Missing target')
 
     # Target
     instance_bundles[0]['target'] = raw_command[command_index]
@@ -152,7 +166,7 @@ def _parse_instance_bundles(instance_bundles):
         if '@' in bundle['target']:
             if len(bundle['target'].split('@')) > 2:
                 # Host details includes an @.  Invalid.
-                raise AssertionError('Target DNS or IP address is invalid')
+                raise AssertionError('Invalid target')
             # A user was specified
             bundle['username'], bundle['target'] = bundle['target'].split('@')
         if ':' in bundle['target']:
@@ -185,13 +199,13 @@ def _parse_instance_bundles(instance_bundles):
 
         # Validate instance ID format
         if len(bundle['instance_id']) > 0 and not INSTANCE_ID_RE.match(bundle['instance_id'].lower()):
-            raise AssertionError('Instance ID {0} is not valid'.format(bundle['instance_id']))
+            raise AssertionError('Invalid instance_id')
 
         # Validate DNS/IP/hostname/etc
         if bundle.get('target', None):
             if not _is_valid_target(bundle.get('target', '')):
                 # It might be an IP
-                raise AssertionError('Target DNS or IP address is invalid')
+                raise AssertionError('Invalid target')
 
     return instance_bundles
 

ec2instanceconnectcli/mops.py

@@ -7,13 +7,9 @@
 from ec2instanceconnectcli.EC2InstanceConnectLogger import EC2InstanceConnectLogger
 from ec2instanceconnectcli import input_parser
 
-
-DEFAULT_USER = ''
 DEFAULT_INSTANCE = ''
-DEFAULT_PORT = 22
 DEFAULT_PROFILE = None
 
-
 def main(program, mode):
     """
     Parses system arguments and sets defaults
@@ -25,29 +21,38 @@ def main(program, mode):
     :type mode: basestring
     """
 
-    parser = argparse.ArgumentParser(description="""Usage:
-    * mssh [user@]instance_id [-u profile] [-z availability_zone] [standard ssh flags] [command]
-    * mssh [user@]dns_name -t instance_id [-u profile] [-z availability_zone] [standard ssh flags] [command]
-    * msftp [user@]instance_id [-u aws_profile] [-z availability_zone] [standard sftp flags]
-    * msftp [user@]dns_name -t instance_id [-u aws_profile] [-z availability_zone] [standard sftp flags]
-    As standard for SFTP, the target may be followed by [:file ...] or [:dir[/]]
-    """)
-    parser.add_argument('-r', '--region', action='store', help='AWS region', type=str)
-    parser.add_argument('-z', '--zone', action='store', help='Availability zone', type=str)
-    parser.add_argument('-v', '--verbose', help='Enable logging', action="store_true")
+    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.add_argument('-u', '--profile', action='store', help='AWS Config Profile', type=str, default=DEFAULT_PROFILE)
-    parser.add_argument('-t', '--instance_id', action='store', help='EC2 Instance ID.  Required if target is DNS name or IP address.', type=str, default=DEFAULT_INSTANCE)
+    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:
-        args = parser.parse_known_args()
-    except:
+        instance_bundles, flags, program_command = input_parser.parseargs(args, mode)
+    except Exception as e:
+        print(str(e))
         parser.print_help()
         sys.exit(1)
 
-    logger = EC2InstanceConnectLogger(args[0].verbose)
-    instance_bundles, flags, program_command = input_parser.parseargs(args, mode)
-
     #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())

ec2instanceconnectcli/mputty.py

@@ -6,9 +6,7 @@
 from ec2instanceconnectcli.EC2InstanceConnectLogger import EC2InstanceConnectLogger
 from ec2instanceconnectcli import input_parser
 
-DEFAULT_USER = ''
 DEFAULT_INSTANCE = ''
-DEFAULT_PORT = 22
 DEFAULT_PROFILE = None
 
 def main(program, mode):
@@ -22,45 +20,61 @@ def main(program, mode):
     :type mode: basestring
     """
 
-    parser = argparse.ArgumentParser(description="""Usage:
-    * mssh-putty [user@]instance_id [-u profile] [-r region] [-z availability_zone] [-i identity_file_ppk] [standard ssh flags] [command]
-    * mssh-putty [user@]dns_name -t instance_id [-u profile] [-r region] [-z availability_zone] [-i identity_file_ppk] [standard ssh flags] [command]
-    * msftp-putty [user@]instance_id [-u profile] [-r region] [-z availability_zone] [-i identity_file_ppk] [standard sftp flags]
-    * msftp-putty [user@]dns_name -t instance_id [-u profile] [-r region] [-z availability_zone] [-i identity_file_ppk] [standard sftp flags]
-    """)
-    parser.add_argument('-r', '--region', action='store', help='AWS region', type=str)
-    parser.add_argument('-z', '--zone', action='store', help='Availability zone', type=str)
-    parser.add_argument('-v', '--verbose', help='Enable logging', action="store_true")
-    parser.add_argument('-i', '--identity', action='store', help="Required. Identity file in ppk format", type=str)
-    parser.add_argument('-u', '--profile', action='store', help='AWS Config Profile', type=str, default=DEFAULT_PROFILE)
-    parser.add_argument('-t', '--instance_id', action='store', help='EC2 Instance ID.  Required if target is DNS name or IP address.', type=str, default=DEFAULT_INSTANCE)
+    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
 
-    try:
-        args = parser.parse_known_args()
-    except:
-        parser.print_help()
-        sys.exit(1)
+            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 = []
-    with open(args[0].identity, 'r') as f:
-        pub_key_lines = f.readlines()
+    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)
+        #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)]
-    pub_key = "ssh-rsa "
-    for pub_key_line in pub_key_lines:
-        pub_key += pub_key_line[:-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].verbose)
-    instance_bundles, flags, program_command = input_parser.parseargs(args, mode)
+    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()