Merge pull request #67 from StartTheShift/sg_vpc_support

Security Group VPC support
This commit is contained in:
Steve Pulec 2013-12-08 20:18:01 -08:00
commit 71d655eca1
5 changed files with 91 additions and 30 deletions

View File

@ -5,6 +5,7 @@ python:
env:
matrix:
#- BOTO_VERSION=2.13.3
- BOTO_VERSION=2.19.0
- BOTO_VERSION=2.12.0
- BOTO_VERSION=2.11.0
- BOTO_VERSION=2.10.0

View File

@ -12,3 +12,4 @@ Moto is written by Steve Pulec with contributions from:
* [Konstantinos Koukopoulos](https://github.com/kouk)
* [attili](https://github.com/attili)
* [JJ Zeng](https://github.com/jjofseattle)
* [Jon Haddad](https://github.com/rustyrazorblade)

View File

@ -1,4 +1,5 @@
import copy
import itertools
from collections import defaultdict
from boto.ec2.instance import Instance as BotoInstance, Reservation
@ -300,45 +301,50 @@ class SecurityRule(object):
class SecurityGroup(object):
def __init__(self, group_id, name, description):
def __init__(self, group_id, name, description, vpc_id=None):
self.id = group_id
self.name = name
self.description = description
self.ingress_rules = []
self.egress_rules = []
self.vpc_id = vpc_id
class SecurityGroupBackend(object):
def __init__(self):
self.groups = {}
# the key in the dict group is the vpc_id or None (non-vpc)
self.groups = defaultdict(dict)
super(SecurityGroupBackend, self).__init__()
def create_security_group(self, name, description, force=False):
def create_security_group(self, name, description, vpc_id=None, force=False):
group_id = random_security_group_id()
if not force:
existing_group = self.get_security_group_from_name(name)
existing_group = self.get_security_group_from_name(name, vpc_id)
if existing_group:
return None
group = SecurityGroup(group_id, name, description)
self.groups[group_id] = group
group = SecurityGroup(group_id, name, description, vpc_id=vpc_id)
self.groups[vpc_id][group_id] = group
return group
def describe_security_groups(self):
return self.groups.values()
return itertools.chain(*[x.values() for x in self.groups.values()])
def delete_security_group(self, name_or_group_id):
if name_or_group_id in self.groups:
# Group Id
return self.groups.pop(name_or_group_id)
else:
# Group Name
group = self.get_security_group_from_name(name_or_group_id)
def delete_security_group(self, name=None, group_id=None):
if group_id:
# loop over all the SGs, find the right one
for vpc in self.groups.values():
if group_id in vpc:
return vpc.pop(group_id)
elif name:
# Group Name. Has to be in standard EC2, VPC needs to be identified by group_id
group = self.get_security_group_from_name(name, None)
if group:
return self.groups.pop(group.id)
return self.groups[None].pop(group.id)
def get_security_group_from_name(self, name):
for group_id, group in self.groups.iteritems():
def get_security_group_from_name(self, name, vpc_id):
for group_id, group in self.groups[vpc_id].iteritems():
if group.name == name:
return group
@ -347,20 +353,20 @@ class SecurityGroupBackend(object):
default_group = ec2_backend.create_security_group("default", "The default security group", force=True)
return default_group
def authorize_security_group_ingress(self, group_name, ip_protocol, from_port, to_port, ip_ranges=None, source_group_names=None):
group = self.get_security_group_from_name(group_name)
def authorize_security_group_ingress(self, group_name, ip_protocol, from_port, to_port, ip_ranges=None, source_group_names=None, vpc_id=None):
group = self.get_security_group_from_name(group_name, vpc_id)
source_groups = []
for source_group_name in source_group_names:
source_groups.append(self.get_security_group_from_name(source_group_name))
source_groups.append(self.get_security_group_from_name(source_group_name, vpc_id))
security_rule = SecurityRule(ip_protocol, from_port, to_port, ip_ranges, source_groups)
group.ingress_rules.append(security_rule)
def revoke_security_group_ingress(self, group_name, ip_protocol, from_port, to_port, ip_ranges=None, source_group_names=None):
group = self.get_security_group_from_name(group_name)
def revoke_security_group_ingress(self, group_name, ip_protocol, from_port, to_port, ip_ranges=None, source_group_names=None, vpc_id=None):
group = self.get_security_group_from_name(group_name, vpc_id)
source_groups = []
for source_group_name in source_group_names:
source_groups.append(self.get_security_group_from_name(source_group_name))
source_groups.append(self.get_security_group_from_name(source_group_name, vpc_id))
security_rule = SecurityRule(ip_protocol, from_port, to_port, ip_ranges, source_groups)
if security_rule in group.ingress_rules:
@ -536,12 +542,12 @@ class SpotInstanceRequest(object):
self.security_groups = []
if security_groups:
for group_name in security_groups:
group = ec2_backend.get_security_group_from_name(group_name)
group = ec2_backend.get_security_group_from_name(group_name, None)
if group:
self.security_groups.append(group)
else:
# If not security groups, add the default
default_group = ec2_backend.get_security_group_from_name("default")
default_group = ec2_backend.get_security_group_from_name("default", None)
self.security_groups.append(default_group)

View File

@ -31,7 +31,8 @@ class SecurityGroups(object):
def create_security_group(self):
name = self.querystring.get('GroupName')[0]
description = self.querystring.get('GroupDescription')[0]
group = ec2_backend.create_security_group(name, description)
vpc_id = self.querystring.get("VpcId", [None])[0]
group = ec2_backend.create_security_group(name, description, vpc_id=vpc_id)
if not group:
# There was an exisitng group
return "There was an existing security group with name {0}".format(name), dict(status=409)
@ -40,9 +41,16 @@ class SecurityGroups(object):
def delete_security_group(self):
# TODO this should raise an error if there are instances in the group. See http://docs.aws.amazon.com/AWSEC2/latest/APIReference/ApiReference-query-DeleteSecurityGroup.html
name = self.querystring.get('GroupName')[0]
group = ec2_backend.delete_security_group(name)
name = self.querystring.get('GroupName')
sg_id = self.querystring.get('GroupId')
if name:
group = ec2_backend.delete_security_group(name[0])
elif sg_id:
group = ec2_backend.delete_security_group(group_id=sg_id[0])
# needs name or group now
if not group:
# There was no such group
return "There was no security group with name {0}".format(name), dict(status=404)
@ -83,7 +91,7 @@ DESCRIBE_SECURITY_GROUPS_RESPONSE = """<DescribeSecurityGroupsResponse xmlns="ht
<groupId>{{ group.id }}</groupId>
<groupName>{{ group.name }}</groupName>
<groupDescription>{{ group.description }}</groupDescription>
<vpcId/>
<vpcId>{{ group.vpc_id or ""}}</vpcId>
<ipPermissions>
{% for rule in group.ingress_rules %}
<item>

View File

@ -20,6 +20,42 @@ def test_create_and_describe_security_group():
all_groups.should.have.length_of(1)
all_groups[0].name.should.equal('test security group')
@mock_ec2
def test_create_and_describe_vpc_security_group():
conn = boto.connect_ec2('the_key', 'the_secret')
vpc_id = 'vpc-5300000c'
security_group = conn.create_security_group('test security group', 'this is a test security group', vpc_id=vpc_id)
security_group.vpc_id.should.equal(vpc_id)
security_group.name.should.equal('test security group')
security_group.description.should.equal('this is a test security group')
# Trying to create another group with the same name in the same VPC should throw an error
conn.create_security_group.when.called_with('test security group', 'this is a test security group', vpc_id).should.throw(EC2ResponseError)
all_groups = conn.get_all_security_groups()
all_groups[0].vpc_id.should.equal(vpc_id)
all_groups.should.have.length_of(1)
all_groups[0].name.should.equal('test security group')
@mock_ec2
def test_create_two_security_groups_with_same_name_in_different_vpc():
conn = boto.connect_ec2('the_key', 'the_secret')
vpc_id = 'vpc-5300000c'
vpc_id2 = 'vpc-5300000d'
sg1 = conn.create_security_group('test security group', 'this is a test security group', vpc_id)
sg2 = conn.create_security_group('test security group', 'this is a test security group', vpc_id2)
all_groups = conn.get_all_security_groups()
all_groups.should.have.length_of(2)
all_groups[0].name.should.equal('test security group')
all_groups[1].name.should.equal('test security group')
@mock_ec2
def test_deleting_security_groups():
@ -37,9 +73,18 @@ def test_deleting_security_groups():
conn.get_all_security_groups().should.have.length_of(1)
# Delete by group id
conn.delete_security_group(security_group1.id)
conn.delete_security_group(group_id=security_group1.id)
conn.get_all_security_groups().should.have.length_of(0)
@mock_ec2
def test_delete_security_group_in_vpc():
conn = boto.connect_ec2('the_key', 'the_secret')
vpc_id = "vpc-12345"
security_group1 = conn.create_security_group('test1', 'test1', vpc_id)
# this should not throw an exception
conn.delete_security_group(group_id=security_group1.id)
@mock_ec2
def test_authorize_ip_range_and_revoke():