From d0c610c5ac4a822a8bd6dce3654ceceebc556a09 Mon Sep 17 00:00:00 2001 From: Nuwan Goonasekera Date: Thu, 14 Sep 2017 14:59:13 +0530 Subject: [PATCH] Added keypair filtering --- moto/ec2/models.py | 46 +++++++++++++++++++++----------- moto/ec2/responses/key_pairs.py | 23 ++++++---------- tests/test_ec2/test_key_pairs.py | 19 +++++++++++++ 3 files changed, 58 insertions(+), 30 deletions(-) diff --git a/moto/ec2/models.py b/moto/ec2/models.py index fca84f0a5..0df3797d8 100755 --- a/moto/ec2/models.py +++ b/moto/ec2/models.py @@ -795,16 +795,29 @@ class InstanceBackend(object): return reservations +class KeyPair(object): + def __init__(self, name, fingerprint, material): + self.name = name + self.fingerprint = fingerprint + self.material = material + + def get_filter_value(self, filter_name): + if filter_name == 'key-name': + return self.name + elif filter_name == 'fingerprint': + return self.fingerprint + + class KeyPairBackend(object): def __init__(self): - self.keypairs = defaultdict(dict) + self.keypairs = {} super(KeyPairBackend, self).__init__() def create_key_pair(self, name): if name in self.keypairs: raise InvalidKeyPairDuplicateError(name) - self.keypairs[name] = keypair = random_key_pair() - keypair['name'] = name + keypair = KeyPair(name, **random_key_pair()) + self.keypairs[name] = keypair return keypair def delete_key_pair(self, name): @@ -812,24 +825,27 @@ class KeyPairBackend(object): self.keypairs.pop(name) return True - def describe_key_pairs(self, filter_names=None): + def describe_key_pairs(self, key_names=None, filters=None): results = [] - for name, keypair in self.keypairs.items(): - if not filter_names or name in filter_names: - keypair['name'] = name - results.append(keypair) + if key_names: + results = [keypair for keypair in self.keypairs.values() + if keypair.name in key_names] + if len(key_names) > len(results): + unknown_keys = set(key_names) - set(results) + raise InvalidKeyPairNameError(unknown_keys) + else: + results = self.keypairs.values() - # TODO: Trim error message down to specific invalid name. - if filter_names and len(filter_names) > len(results): - raise InvalidKeyPairNameError(filter_names) - - return results + if filters: + return generic_filter(filters, results) + else: + return results def import_key_pair(self, key_name, public_key_material): if key_name in self.keypairs: raise InvalidKeyPairDuplicateError(key_name) - self.keypairs[key_name] = keypair = random_key_pair() - keypair['name'] = key_name + keypair = KeyPair(key_name, **random_key_pair()) + self.keypairs[key_name] = keypair return keypair diff --git a/moto/ec2/responses/key_pairs.py b/moto/ec2/responses/key_pairs.py index 936df2cd3..59268556e 100644 --- a/moto/ec2/responses/key_pairs.py +++ b/moto/ec2/responses/key_pairs.py @@ -11,7 +11,7 @@ class KeyPairs(BaseResponse): if self.is_not_dryrun('CreateKeyPair'): keypair = self.ec2_backend.create_key_pair(name) template = self.response_template(CREATE_KEY_PAIR_RESPONSE) - return template.render(**keypair) + return template.render(keypair=keypair) def delete_key_pair(self): name = self.querystring.get('KeyName')[0] @@ -23,11 +23,7 @@ class KeyPairs(BaseResponse): def describe_key_pairs(self): names = keypair_names_from_querystring(self.querystring) filters = filters_from_querystring(self.querystring) - if len(filters) > 0: - raise NotImplementedError( - 'Using filters in KeyPairs.describe_key_pairs is not yet implemented') - - keypairs = self.ec2_backend.describe_key_pairs(names) + keypairs = self.ec2_backend.describe_key_pairs(names, filters) template = self.response_template(DESCRIBE_KEY_PAIRS_RESPONSE) return template.render(keypairs=keypairs) @@ -37,7 +33,7 @@ class KeyPairs(BaseResponse): if self.is_not_dryrun('ImportKeyPair'): keypair = self.ec2_backend.import_key_pair(name, material) template = self.response_template(IMPORT_KEYPAIR_RESPONSE) - return template.render(**keypair) + return template.render(keypair=keypair) DESCRIBE_KEY_PAIRS_RESPONSE = """ @@ -54,12 +50,9 @@ DESCRIBE_KEY_PAIRS_RESPONSE = """ - {{ name }} - - {{ fingerprint }} - - {{ material }} - + {{ keypair.name }} + {{ keypair.fingerprint }} + {{ keypair.material }} """ @@ -71,6 +64,6 @@ DELETE_KEY_PAIR_RESPONSE = """ 471f9fdd-8fe2-4a84-86b0-bd3d3e350979 - {{ name }} - {{ fingerprint }} + {{ keypair.name }} + {{ keypair.fingerprint }} """ diff --git a/tests/test_ec2/test_key_pairs.py b/tests/test_ec2/test_key_pairs.py index ec979a871..0a7fb9f76 100644 --- a/tests/test_ec2/test_key_pairs.py +++ b/tests/test_ec2/test_key_pairs.py @@ -130,3 +130,22 @@ def test_key_pairs_import_exist(): cm.exception.code.should.equal('InvalidKeyPair.Duplicate') cm.exception.status.should.equal(400) cm.exception.request_id.should_not.be.none + + +@mock_ec2_deprecated +def test_key_pair_filters(): + conn = boto.connect_ec2('the_key', 'the_secret') + + _ = conn.create_key_pair('kpfltr1') + kp2 = conn.create_key_pair('kpfltr2') + kp3 = conn.create_key_pair('kpfltr3') + + kp_by_name = conn.get_all_key_pairs( + filters={'key-name': 'kpfltr2'}) + set([kp.name for kp in kp_by_name] + ).should.equal(set([kp2.name])) + + kp_by_name = conn.get_all_key_pairs( + filters={'fingerprint': kp3.fingerprint}) + set([kp.name for kp in kp_by_name] + ).should.equal(set([kp3.name]))