initial checkin

[sfa.git] / sfa / openstack / security_group.py

diff --git a/sfa/openstack/security_group.py b/sfa/openstack/security_group.py
new file mode 100644 (file)
index 0000000..f723312
--- /dev/null
+++ b/sfa/openstack/security_group.py
@@ -0,0 +1,107 @@
+from sfa.util.sfalogging import logger
+
+class SecurityGroup:
+
+    def __init__(self, driver):
+        self.driver = driver
+
+        
+    def create_security_group(self, name):
+        conn = self.driver.euca_shell.get_euca_connection()
+        try:
+            conn.create_security_group(name=group_name)
+        except Exception, ex:
+            logger.log_exc("Failed to add security group")
+
+    def delete_security_group(self, name):
+        conn = self.driver.euca_shell.get_euca_connection()
+        try:
+            conn.delete_security_group(name=group_name)
+        except Exception, ex:
+            logger.log_exc("Failed to delete security group")
+
+
+    def _validate_port_range(self, port_range):
+        from_port = to_port = None
+        if isinstance(port_range, str):
+            ports = port_range.split('-')
+            if len(ports) > 1:
+                from_port = int(ports[0])
+                to_port = int(ports[1])
+            else:
+                from_port = to_port = int(ports[0])
+        else:
+            from_port = to_port = None
+        return (from_port, to_port)
+
+    def _validate_icmp_type_code(self, icmp_type_code):
+        from_port = to_port = None
+        if isinstance(icmp_type_code, str):
+            code_parts = icmp_type_code.split(':')
+            if len(code_parts) > 1:
+                try:
+                    from_port = int(code_parts[0])
+                    to_port = int(code_parts[1])
+                except ValueError:
+                    logger.error('port must be an integer.')
+        return (from_port, to_port)
+
+
+    def add_rule_to_group(self, group_name=None, protocol='tcp', cidr_ip='0.0.0.0/0',
+                          port_range=None, icmp_type_code=None,
+                          source_group_name=None, source_group_owner_id=None):
+
+
+        from_port, to_port = self._validate_port_range(port_range)
+        icmp_type = self._validate_icmp_type_code(icmp_type_code)
+        if icmp_type:
+            from_port, to_port = icmp_type[0], icmp_type[1]
+
+        if group_name:
+            if cidr_ip:
+                euca.validate_address(cidr_ip)
+            if protocol:
+               euca.validate_protocol(protocol)
+            conn = self.driver.euca_shell.get_euca_connection()
+            try:
+                conn.authorize_security_group(
+                    group_name=group_name,
+                    src_security_group_name=source_group_name,
+                    src_security_group_owner_id=source_group_owner_id,
+                    ip_protocol=protocol,
+                    from_port=from_port,
+                    to_port=to_port,
+                    cidr_ip=ip,
+                    )
+            except Exception, ex:
+                logger.log_exc("Failed to add rule to group %s" % group_name)
+
+
+    def remove_rule_from_group(self, group_name=None, protocol='tcp', cidr_ip='0.0.0.0/0',
+                          port_range=None, icmp_type_code=None,
+                          source_group_name=None, source_group_owner_id=None):
+
+        from_port, to_port = self._validate_port_range(port_range)
+        icmp_type = self._validate_icmp_type_code(icmp_type_code)
+        if icmp_type:
+            from_port, to_port = icmp_type[0], icmp_type[1]
+
+        if group_name:
+            if cidr_ip:
+                euca.validate_address(cidr_ip)
+            if protocol:
+               euca.validate_protocol(protocol)
+            conn = self.driver.euca_shell.get_euca_connection()
+            try:
+                conn.revoke_security_group(
+                    group_name=group_name,
+                    src_security_group_name=source_group_name,
+                    src_security_group_owner_id=source_group_owner_id,
+                    ip_protocol=protocol,
+                    from_port=from_port,
+                    to_port=to_port,
+                    cidr_ip=ip,
+                    )
+            except Exception, ex:
+                logger.log_exc("Failed to remove rule from group %s" % group_name) 
+