hooks: keypair: add some features and rewrite tests

stacker/hooks/keypair.py

@@ -1,22 +1,184 @@
 from __future__ import print_function
 from __future__ import division
 from __future__ import absolute_import
-from builtins import input
+
 import logging
 import os
+import sys
+
+from botocore.exceptions import ClientError
 
 from stacker.session_cache import get_session
+from stacker.hooks import utils
+from stacker.ui import get_raw_input
 
-from . import utils
 
 logger = logging.getLogger(__name__)
 
+KEYPAIR_LOG_MESSAGE = "keypair: %s (%s) %s"
+
+
+def get_existing_key_pair(ec2, keypair_name):
+    resp = ec2.describe_key_pairs()
+    keypair = next((kp for kp in resp["KeyPairs"]
+                    if kp["KeyName"] == keypair_name), None)
+
+    if keypair:
+        logger.info(KEYPAIR_LOG_MESSAGE,
+                    keypair["KeyName"],
+                    keypair["KeyFingerprint"],
+                    "exists")
+        return {
+            "status": "exists",
+            "key_name": keypair["KeyName"],
+            "fingerprint": keypair["KeyFingerprint"],
+        }
 
-def find(lst, key, value):
-    for i, dic in enumerate(lst):
-        if dic[key] == value:
-            return lst[i]
-    return False
+    logger.info("keypair: \"%s\" not found", keypair_name)
+    return None
+
+
+def import_key_pair(ec2, keypair_name, public_key_data):
+    keypair = ec2.import_key_pair(
+        KeyName=keypair_name,
+        PublicKeyMaterial=public_key_data.strip(),
+        DryRun=False)
+    logger.info(KEYPAIR_LOG_MESSAGE,
+                keypair["KeyName"],
+                keypair["KeyFingerprint"],
+                "imported")
+    return keypair
+
+
+def read_public_key_file(path):
+    try:
+        with open(utils.full_path(path), 'rb') as f:
+            data = f.read()
+
+        if not data.startswith(b"ssh-rsa"):
+            raise ValueError(
+                "Bad public key data, must be an RSA key in SSH authorized "
+                "keys format (beginning with `ssh-rsa`)")
+
+        return data.strip()
+    except (ValueError, IOError, OSError) as e:
+        logger.error("Failed to read public key file {}: {}".format(
+            path, e))
+        return None
+
+
+def create_key_pair_from_public_key_file(ec2, keypair_name, public_key_path):
+    public_key_data = read_public_key_file(public_key_path)
+    if not public_key_data:
+        return None
+
+    keypair = import_key_pair(ec2, keypair_name, public_key_data)
+    return {
+        "status": "imported",
+        "key_name": keypair["KeyName"],
+        "fingerprint": keypair["KeyFingerprint"],
+    }
+
+
+def create_key_pair_in_ssm(ec2, ssm, keypair_name, parameter_name,
+                           kms_key_id=None):
+    keypair = create_key_pair(ec2, keypair_name)
+    try:
+        kms_key_label = 'default'
+        kms_args = {}
+        if kms_key_id:
+            kms_key_label = kms_key_id
+            kms_args = {"KeyId": kms_key_id}
+
+        logger.info("Storing generated key in SSM parameter \"%s\" "
+                    "using KMS key \"%s\"", parameter_name, kms_key_label)
+
+        ssm.put_parameter(
+            Name=parameter_name,
+            Description="SSH private key for KeyPair \"{}\" "
+                        "(generated by Stacker)".format(keypair_name),
+            Value=keypair["KeyMaterial"],
+            Type="SecureString",
+            Overwrite=False,
+            **kms_args)
+    except ClientError:
+        # Erase the key pair if we failed to store it in SSM, since the
+        # private key will be lost anyway
+
+        logger.exception("Failed to store generated key in SSM, deleting "
+                         "created key pair as private key will be lost")
+        ec2.delete_key_pair(KeyName=keypair_name, DryRun=False)
+        return None
+
+    return {
+        "status": "created",
+        "key_name": keypair["KeyName"],
+        "fingerprint": keypair["KeyFingerprint"],
+    }
+
+
+def create_key_pair(ec2, keypair_name):
+    keypair = ec2.create_key_pair(KeyName=keypair_name, DryRun=False)
+    logger.info(KEYPAIR_LOG_MESSAGE,
+                keypair["KeyName"],
+                keypair["KeyFingerprint"],
+                "created")
+    return keypair
+
+
+def create_key_pair_local(ec2, keypair_name, dest_dir):
+    dest_dir = utils.full_path(dest_dir)
+    if not os.path.isdir(dest_dir):
+        logger.error("\"%s\" is not a valid directory", dest_dir)
+        return None
+
+    file_name = "{0}.pem".format(keypair_name)
+    key_path = os.path.join(dest_dir, file_name)
+    if os.path.isfile(key_path):
+        # This mimics the old boto2 keypair.save error
+        logger.error("\"%s\" already exists in \"%s\" directory",
+                     file_name, dest_dir)
+        return None
+
+    # Open the file before creating the key pair to catch errors early
+    with open(key_path, "wb") as f:
+        keypair = create_key_pair(ec2, keypair_name)
+        f.write(keypair["KeyMaterial"].encode("ascii"))
+
+    return {
+        "status": "created",
+        "key_name": keypair["KeyName"],
+        "fingerprint": keypair["KeyFingerprint"],
+        "file_path": key_path
+    }
+
+
+def interactive_prompt(keypair_name, ):
+    if not sys.stdin.isatty():
+        return None, None
+
+    try:
+        while True:
+            action = get_raw_input(
+                "import or create keypair \"%s\"? (import/create/cancel) " % (
+                    keypair_name,
+                )
+            )
+
+            if action.lower() == "cancel":
+                break
+
+            if action.lower() in ("i", "import"):
+                path = get_raw_input("path to keypair file: ")
+                return "import", path.strip()
+
+            if action.lower() == "create":
+                path = get_raw_input("directory to save keyfile: ")
+                return "create", path.strip()
+    except (EOFError, KeyboardInterrupt):
+        return None, None
+
+    return None, None
 
 
 def ensure_keypair_exists(provider, context, **kwargs):
@@ -28,84 +190,63 @@ def ensure_keypair_exists(provider, context, **kwargs):
         provider (:class:`stacker.providers.base.BaseProvider`): provider
             instance
         context (:class:`stacker.context.Context`): context instance
-
-    Returns: boolean for whether or not the hook succeeded.
+        keypair (str): name of the key pair to create
+        ssm_parameter_name (str, optional): path to an SSM store parameter to
+            receive the generated private key, instead of importing it or
+            storing it locally.
+        ssm_key_id (str, optional): ID of a KMS key to encrypt the SSM
+            parameter with. If omitted, the default key will be used.
+        public_key_path (str, optional): path to a public key file to be
+            imported instead of generating a new key. Incompatible with the SSM
+            options, as the private key will not be available for storing.
+
+    Returns:
+        In case of failure ``False``, otherwise a dict containing:
+            status (str): one of "exists", "imported" or "created"
+            key_name (str): name of the key pair
+            fingerprint (str): fingerprint of the key pair
+            file_path (str, optional): if a new key was created, the path to
+                the file where the private key was stored
 
     """
-    session = get_session(provider.region)
-    client = session.client("ec2")
-    keypair_name = kwargs.get("keypair")
-    resp = client.describe_key_pairs()
-    keypair = find(resp["KeyPairs"], "KeyName", keypair_name)
-    message = "keypair: %s (%s) %s"
+
+    keypair_name = kwargs["keypair"]
+    ssm_parameter_name = kwargs.get("ssm_parameter_name")
+    ssm_key_id = kwargs.get("ssm_key_id")
+    public_key_path = kwargs.get("public_key_path")
+
+    if public_key_path and ssm_parameter_name:
+        logger.error("public_key_path and ssm_parameter_name cannot be "
+                     "specified at the same time")
+        return False
+
+    session = get_session(region=provider.region,
+                          profile=kwargs.get("profile"))
+    ec2 = session.client("ec2")
+
+    keypair = get_existing_key_pair(ec2, keypair_name)
     if keypair:
-        logger.info(message,
-                    keypair["KeyName"],
-                    keypair["KeyFingerprint"],
-                    "exists")
-        return {
-            "status": "exists",
-            "key_name": keypair["KeyName"],
-            "fingerprint": keypair["KeyFingerprint"],
-        }
+        return keypair
 
-    logger.info("keypair: \"%s\" not found", keypair_name)
-    create_or_upload = input(
-        "import or create keypair \"%s\"? (import/create/Cancel) " % (
-            keypair_name,
-        ),
-    )
-    if create_or_upload == "import":
-        path = input("path to keypair file: ")
-        full_path = utils.full_path(path)
-        if not os.path.exists(full_path):
-            logger.error("Failed to find keypair at path: %s", full_path)
-            return False
-
-        with open(full_path) as read_file:
-            contents = read_file.read()
-
-        keypair = client.import_key_pair(KeyName=keypair_name,
-                                         PublicKeyMaterial=contents)
-        logger.info(message,
-                    keypair["KeyName"],
-                    keypair["KeyFingerprint"],
-                    "imported")
-        return {
-            "status": "imported",
-            "key_name": keypair["KeyName"],
-            "fingerprint": keypair["KeyFingerprint"],
-            "file_path": full_path,
-        }
-    elif create_or_upload == "create":
-        path = input("directory to save keyfile: ")
-        full_path = utils.full_path(path)
-        if not os.path.exists(full_path) and not os.path.isdir(full_path):
-            logger.error("\"%s\" is not a valid directory", full_path)
-            return False
-
-        file_name = "{0}.pem".format(keypair_name)
-        if os.path.isfile(os.path.join(full_path, file_name)):
-            # This mimics the old boto2 keypair.save error
-            logger.error("\"%s\" already exists in \"%s\" directory",
-                         file_name,
-                         full_path)
-            return False
-
-        keypair = client.create_key_pair(KeyName=keypair_name)
-        logger.info(message,
-                    keypair["KeyName"],
-                    keypair["KeyFingerprint"],
-                    "created")
-        with open(os.path.join(full_path, file_name), "w") as f:
-            f.write(keypair["KeyMaterial"])
+    if public_key_path:
+        keypair = create_key_pair_from_public_key_file(
+            ec2, keypair_name, public_key_path)
 
-        return {
-            "status": "created",
-            "key_name": keypair["KeyName"],
-            "fingerprint": keypair["KeyFingerprint"],
-            "file_path": os.path.join(full_path, file_name)
-        }
+    elif ssm_parameter_name:
+        ssm = session.client('ssm')
+        keypair = create_key_pair_in_ssm(
+            ec2, ssm, keypair_name, ssm_parameter_name, ssm_key_id)
     else:
-        logger.warning("no action to find keypair, failing")
+        action, path = interactive_prompt(keypair_name)
+        if action == "import":
+            keypair = create_key_pair_from_public_key_file(
+                ec2, keypair_name, path)
+        elif action == "create":
+            keypair = create_key_pair_local(ec2, keypair_name, path)
+        else:
+            logger.warning("no action to find keypair, failing")
+
+    if not keypair:
         return False
+
+    return keypair

stacker/tests/conftest.py

@@ -4,7 +4,7 @@
 import os
 
 import pytest
-
+import py.path
 
 logger = logging.getLogger(__name__)
 
@@ -35,3 +35,10 @@ def aws_credentials():
             os.environ[key] = value
 
     saved_env.clear()
+
+
+@pytest.fixture(scope="package")
+def stacker_fixture_dir():
+    path = os.path.join(os.path.dirname(os.path.realpath(__file__)),
+                        'fixtures')
+    return py.path.local(path)

stacker/tests/fixtures/keypair/fingerprint

@@ -0,0 +1 @@
+d7:50:1f:78:55:5f:22:c1:f6:88:c6:5d:82:4f:94:4f

stacker/tests/fixtures/keypair/id_rsa

@@ -0,0 +1,27 @@
+-----BEGIN OPENSSH PRIVATE KEY-----
+b3BlbnNzaC1rZXktdjEAAAAABG5vbmUAAAAEbm9uZQAAAAAAAAABAAABFwAAAAdzc2gtcn
+NhAAAAAwEAAQAAAQEA7rF34ExOHgT+dDYJUswkhBpyC+vnK+ptx+nGQDTkPj9aP1uAXbXA
+C97KK+Ihou0jniYKPJMHsjEK4a7eh2ihoK6JkYs9+y0MeGCAHAYuGXdNt5jv1e0XNgoYdf
+JloC0pgOp4Po9+4qeuOds8bb9IxwM/aSaJWygaSc22ZTzeOWQk5PXJNH0lR0ZelUUkj0HK
+aouuV6UX/t+czTghgnNZgDjk5sOfUNmugN7fJi+6/dWjOaukDkJttfZXLRTPDux0SZw4Jo
+RqZ40cBNS8ipLVk24BWeEjVlNl6rrFDtO4yrkscz7plwXlPiRLcdCdbamcCZaRrdkftKje
+5ypz5dvocQAAA9DJ0TBmydEwZgAAAAdzc2gtcnNhAAABAQDusXfgTE4eBP50NglSzCSEGn
+IL6+cr6m3H6cZANOQ+P1o/W4BdtcAL3sor4iGi7SOeJgo8kweyMQrhrt6HaKGgromRiz37
+LQx4YIAcBi4Zd023mO/V7Rc2Chh18mWgLSmA6ng+j37ip6452zxtv0jHAz9pJolbKBpJzb
+ZlPN45ZCTk9ck0fSVHRl6VRSSPQcpqi65XpRf+35zNOCGCc1mAOOTmw59Q2a6A3t8mL7r9
+1aM5q6QOQm219lctFM8O7HRJnDgmhGpnjRwE1LyKktWTbgFZ4SNWU2XqusUO07jKuSxzPu
+mXBeU+JEtx0J1tqZwJlpGt2R+0qN7nKnPl2+hxAAAAAwEAAQAAAQAwMUSy1LUw+nElpYNc
+ZDs7MNu17HtQMpTXuCt+6y7qIoBmKmNQiFGuE91d3tpLuvVmCOgoMsdrAtvflR741/dKKf
+M8n5B0FjReWZ2ECvtjyOK4HvjNiIEXOBKYPcim/ndSwARnHTHRMWnL5KfewLBA/jbfVBiH
+fyFPpWkeJ5v2mg3EDCkTCj7mBZwXYkX8uZ1IN6CZJ9kWNaPO3kloTlamgs6pd/5+OmMGWc
+/vhfJQppaJjW58y7D7zCpncHg3Yf0HZsgWRTGJO93TxuyzDlAXITVGwqcz7InTVQZS1XTx
+3FNmIpb0lDtVrKGxwvR/7gP6DpxMlKkzoCg3j1o8tHvBAAAAgQDuZCVAAqQFrY4ZH2TluP
+SFulXuTiT4mgQivAwI6ysMxjpX1IGBTgDvHXJ0xyW4LN7pCvg8hRAhsPlaNBX24nNfOGmn
+QMYp/qAZG5JP2vEJmDUKmEJ77Twwmk+k0zXfyZyfo7rgpF4c5W2EFnV7xiMtBTKbAj4HMn
+qGPYDPGpySTwAAAIEA+w72mMctM2yd9Sxyg5b7ZlhuNyKW1oHcEvLoEpTtru0f8gh7C3HT
+C0SiuTOth2xoHUWnbo4Yv5FV3gSoQ/rd1sWbkpEZMwbaPGsTA8bkCn2eItsjfrQx+6oY1U
+HgZDrkjbByB3KQiq+VioKsrUmgfT/UgBq2tSnHqcYB56Eqj0sAAACBAPNkMvCstNJGS4FN
+nSCGXghoYqKHivZN/IjWP33t/cr72lGp1yCY5S6FCn+JdNrojKYk2VXOSF5xc3fZllbr7W
+hmhXRr/csQkymXMDkJHnsdhpMeoEZm7wBjUx+hE1+QbNF63kZMe9sjm5y/YRu7W7H6ngme
+kb5FW97sspLYX8WzAAAAF2RhbmllbGt6YUBkYW5pZWwtcGMubGFuAQID
+-----END OPENSSH PRIVATE KEY-----

stacker/tests/fixtures/keypair/id_rsa.pub

@@ -0,0 +1 @@
+ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAQEAklOUpkDHrfHY17SbrmTIpNLTGK9Tjom/BWDSUGPl+nafzlHDTYW7hdI4yZ5ew18JH4JW9jbhUFrviQzM7xlELEVf4h9lFX5QVkbPppSwg0cda3Pbv7kOdJ/MTyBlWXFCR+HAo3FXRitBqxiX1nKhXpHAZsMciLq8V6RjsNAQwdsdMFvSlVK/7XAt3FaoJoAsncM1Q9x5+3V0Ww68/eIFmb1zuUFljQJKprrX88XypNDvjYNby6vw/Pb0rwert/EnmZ+AW4OZPnTPI89ZPmVMLuayrD2cE86Z/il8b+gw3r3+1nKatmIkjn2so1d01QraTlMqVSsbxNrRFi9wrf+M7Q==