implemented SecurityGroup entity for AWS EC2

Author: Matous Mojzis

wrapanapi/entities/__init__.py

@@ -10,9 +10,11 @@
 from .server import Server, ServerState
 from .network import Network, NetworkMixin
 from .volume import Volume, VolumeMixin
+from .security_group import SecurityGroup, SecurityGroupMixin
 
 __all__ = [
     'Template', 'TemplateMixin', 'Vm', 'VmState', 'VmMixin', 'Instance',
     'PhysicalContainer', 'Server', 'ServerState', 'Stack', 'StackMixin',
-    'Network', 'NetworkMixin', 'Volume', 'VolumeMixin'
+    'Network', 'NetworkMixin', 'Volume', 'VolumeMixin', 'SecurityGroup',
+    'SecurityGroupMixin'
 ]

wrapanapi/entities/security_group.py

@@ -0,0 +1,83 @@
+"""
+wrapanapi.entities.security_group
+
+SecurityGroups
+"""
+
+from abc import ABCMeta, abstractmethod
+
+from wrapanapi.entities.base import Entity, EntityMixin
+from wrapanapi.exceptions import MultipleItemsError, NotFoundError
+
+
+class SecurityGroup(Entity, metaclass=ABCMeta):
+    """
+    Defines methods/properties pertaining to security groups
+    """
+    @abstractmethod
+    def get_details(self):
+        """
+        Return a dict with detailed info about this object
+
+        There's no specific prescription for how this dict should be formatted--
+        it will vary based on entity and provider type. It is recommended
+        that the values contain simple python data types instead of
+        complex classes so the data can be parsed easily.
+
+        Returns: dict
+        """
+
+
+class SecurityGroupMixin(EntityMixin, metaclass=ABCMeta):
+    """
+    Defines methods for systems that support security groups
+    """
+    @abstractmethod
+    def create_sec_group(self, **kwargs):
+        """
+        Creates security group
+
+        Returns: wrapanapi.entities.SecurityGroup for newly created SecurityGroup
+        """
+
+    @abstractmethod
+    def list_sec_groups(self, **kwargs):
+        """
+        Return a list of SecurityGroup entities.
+
+        Returns: list of SecurityGroup objects
+        """
+
+    @abstractmethod
+    def find_sec_groups(self, name, **kwargs):
+        """
+        Find a security group based on 'name' or other kwargs
+
+        Returns an empty list if no matches found
+
+        Returns: implementation of wrapanapi.network.SecurityGroup
+        """
+
+    @abstractmethod
+    def get_sec_group(self, name, **kwargs):
+        """
+        Get a security group based on name or other kwargs
+
+        Returns: SecurityGroup object
+        Raises:
+            MultipleItemsError if multiple matches found
+            NotFoundError if unable to find security group
+        """
+
+    def does_sec_group_exist(self, name):
+        """
+        Checks if a security group with 'name' exists on the system
+
+        If multiple security groups with the same name exists, this still returns 'True'
+        """
+        try:
+            return bool(self.get_sec_group(name))
+        except MultipleItemsError:
+            return True
+        except NotFoundError:
+            return False

wrapanapi/systems/ec2.py

@@ -12,8 +12,9 @@
     client as boto3client
 )
 
-from wrapanapi.entities import (Instance, Network, NetworkMixin, Stack, StackMixin,
-                                Template, TemplateMixin, VmMixin, VmState, Volume)
+from wrapanapi.entities import (Instance, Network, NetworkMixin, SecurityGroup, SecurityGroupMixin,
+                                Stack, StackMixin, Template, TemplateMixin, VmMixin, VmState,
+                                Volume)
 from wrapanapi.exceptions import (ActionTimedOutError, MultipleItemsError, NotFoundError)
 from wrapanapi.systems.base import System
 
@@ -449,7 +450,98 @@ def cleanup(self):
         return self.delete()
 
 
-class EC2System(System, VmMixin, TemplateMixin, StackMixin, NetworkMixin):
+class SecurityGroup(_TagMixin, _SharedMethodsMixin, SecurityGroup):
+    def __init__(self, system, raw=None, **kwargs):
+        """
+        Constructor for an SecurityGroup tied to a specific system.
+
+        Args:
+            system: an EC2System object
+            raw: the boto.ec2.volume.Volume object if already obtained, or None
+            uuid: unique ID of the volume
+        """
+        self._uuid = raw.id if raw else kwargs.get('uuid')
+        if not self._uuid:
+            raise ValueError("missing required kwarg: 'uuid'")
+
+        super(SecurityGroup, self).__init__(system, raw, **kwargs)
+
+        self._api = self.system.ec2_connection
+
+    @property
+    def name(self):
+        tag_value = self.get_tag_value('Name')
+        return tag_value if tag_value else self.raw.group_name
+
+    def set_rule(self, permission_list, traffic_type='inbound'):
+        """
+
+        Args:
+            permission_list= [{
+                'FromPort': from_port(-1 - all ports),
+                'ToPort': to_port(-1 - all ports),
+                'IpProtocol': 'tcp/udp/icmp/icmpv6/-1(all protocols)',
+                'IpRanges': [{'CidrIp': '0.0.0.0/32'}]
+                }]
+            traffic_type: inbound/outbound
+
+        Returns:
+
+        """
+        try:
+            if traffic_type == 'inbound':
+                self.raw.authorize_ingress(IpPermissions=permission_list)
+            else:
+                self.raw.authorize_egress(IpPermissions=permission_list)
+            self.refresh()
+            return True
+        except Exception:
+            return False
+
+    def unset_rule(self, permission_list, traffic_type='inbound'):
+        """
+
+        Args:
+            permission_list: [
+                FromPort: from_port(-1 - all ports),
+                ToPort: to_port(-1 - all ports),
+                IpProtocol: 'tcp/udp/icmp/icmpv6/-1(all protocols)',
+                IpRanges: [CidrIp: '0.0.0.0/32']
+                ]
+            type: inbound/outbound
+
+        Returns:
+
+        """
+        try:
+            if traffic_type == 'inbound':
+                self.raw.revoke_ingress(IpPermissions=permission_list)
+            else:
+                self.raw.revoke_egress(IpPermissions=permission_list)
+            self.refresh()
+            return True
+        except Exception:
+            return False
+
+    def delete(self):
+        """
+        Delete SecurityGroup
+        """
+        self.logger.info("Deleting SecurityGroup '%s', id: '%s'", self.name, self.uuid)
+        try:
+            self.raw.delete()
+            return True
+        except ActionTimedOutError:
+            return False
+
+    def cleanup(self):
+        """
+        Cleanup SecurityGroup
+        """
+        return self.delete()
+
+
+class EC2System(System, VmMixin, TemplateMixin, StackMixin, NetworkMixin, SecurityGroupMixin):
     """EC2 Management System, powered by boto
 
     Wraps the EC2 API
@@ -1512,3 +1604,44 @@ def get_snapshot_id_if_import_completed(self, task_id):
             return result.get("SnapshotId")
         else:
             return False
+
+    def create_sec_group(self, group_name, desc="SG created for automated test", vpc_id=None):
+        sg_kwargs = {'Description': desc, 'GroupName': group_name}
+        if vpc_id:
+            sg_kwargs['VpcId'] = vpc_id
+        result = self.ec2_connection.create_security_group(**sg_kwargs)
+        return SecurityGroup(system=self, uuid=result['GroupId'],
+                             raw=self.ec2_resource.SecurityGroup(result['GroupId']))
+
+    def get_sec_group(self, name=None, id=None):
+        return self._get_resource(SecurityGroup, self.find_sec_groups, name=name, id=id)
+
+    def list_sec_groups(self):
+        """
+        Returns a list of SecurityGroup objects
+        """
+        sec_group_list = [
+            EBSVolume(system=self, uuid=sec_group['GroupId'], raw=self.ec2_resource.SecurityGroup(
+                sec_group['GroupId']))
+            for sec_group in self.ec2_connection.describe_security_groups().get('SecurityGroups')
+        ]
+        return sec_group_list
+
+    def find_sec_groups(self, name=None, id=None):
+        """
+        Return list of all security groups with given name or id
+        Args:
+            name: name to search
+            id: id to search
+        Returns:
+            List of SecurityFGroup objects
+        """
+        if not name and not id or name and id:
+            raise ValueError("Either name or id must be set and not both!")
+        if id:
+            sec_groups = self.ec2_connection.describe_security_groups(GroupIds=[id])
+        else:
+            sec_groups = self.ec2_connection.describe_security_groups(GroupNames=[name])
+        return [
+            SecurityGroup(system=self, raw=self.ec2_resource.SecurityGroup(sec_group['GroupId']))
+            for sec_group in sec_groups.get('SecurityGroups')]