Inspired by https://gist.github.com/robbyt/2493423 This StarCluster plugin grants all tcp, udp and icmp privileges for 10.0.0.0/8 between the current cluster's security group and the given security group, in both directions for cidr block 10.0.0.0/8. This would be particularly useful for using StarCluster within Amazon VPC.

add_to_security_group.py

"""
Based on https://gist.github.com/robbyt/2493423

This StarCluster plugin grants all tcp, udp and icmp privileges for 10.0.0.0/8
between the current cluster's security group and the given security group,
in both directions
"""
from starcluster.clustersetup import ClusterSetup
from starcluster.logger import log


class AddToSecGroup(ClusterSetup):
    def __init__(self, sec_group):
        self.sec_group = sec_group
        self.rules = [dict(ip_protocol='tcp', from_port=1, to_port=65535,
                           cidr_ip='10.0.0.0/8'),
                      dict(ip_protocol='udp', from_port=1, to_port=65535,
                           cidr_ip='10.0.0.0/8'),
                      dict(ip_protocol='icmp', from_port=-1, to_port=-1,
                           cidr_ip='10.0.0.0/8'),
                      ]

    def _modify_rule(self, to_group, from_group, method):
        """Call to_group.authorize(from_group, **rule) or
        to_group.revoke(from_group, **rule) for each rule in rules"""
        assert method in ['authorize', 'revoke'], \
            "Method must be either 'authorize' or 'revoke'"
        log.info("%s access from security group %s to %s" %
                (method, str(to_group), str(from_group)))
        for params in self.rules:
            getattr(to_group, method)(src_group=from_group, **params)

    def _get_security_group_instances(self, node):
        cluster_group = node.cluster_groups[0]
        sec_group = node.ec2.get_security_group(self.sec_group)
        return(cluster_group, sec_group)

    def run(self, nodes, master, user, user_shell, volumes):
        """Grant self.rules between (ie to AND from) this cluster's
        security group and the given security group, self.sec_group

        First attempt to revoke the rule in case it already exists"""
        cg, sg = self._get_security_group_instances(master)
        for sec_group1, sec_group2 in [(cg, sg), (sg, cg)]:
            try:
                self._modify_rule(sec_group1, sec_group2, 'revoke')
            except:
                pass  # remove the rule if it exists.
            self._modify_rule(sec_group1, sec_group2, 'authorize')

    def on_shutdown(self, nodes, master, user, user_shell, volumes):
        """Revoke self.rules from this cluster's security group
        and the given security group"""
        cg, sg = self._get_security_group_instances(master)

        for sec_group1, sec_group2 in [(cg, sg), (sg, cg)]:
            try:
                self._modify_rule(sec_group1, sec_group2, 'revoke')
            except Exception as e:
                log.warn('Failed to revoke permissions with error: %s' % e)