from starcluster.clustersetup import ClusterSetup
from starcluster.logger import log
from starcluster.config import get_easy_ec2
from starcluster import awsutils

class AddToSecGroup(ClusterSetup):
    def __init__(self, sec_group):
        self.sec_group = sec_group

    def _grant(self, to_group, from_group):
        log.info("Granting access from security group %s to %s" % 
                  (str(to_group), str(from_group)))
        try:
            to_group.authorize('tcp', 1, 65535, src_group=from_group)
            to_group.authorize('udp', 1, 65535, src_group=from_group)
        except:
            log.info("Error Adding security group")

    def _revoke(self, to_group, from_group):
        log.info("Revoking access from security group %s to %s" % 
                  (str(to_group), str(from_group)))
        try:
            to_group.revoke('tcp', 1, 65535, src_group=from_group)
            to_group.revoke('udp', 1, 65535, src_group=from_group)
        except:
            log.info("Error Revoking security group")

    def _get_group_objects(self, node):
        cluster_group = node.cluster_groups[0]
        sec_group = node.ec2.get_or_create_group(self.sec_group, "StarCluster")
        return(cluster_group, sec_group)

    def run(self, nodes, master, user, user_shell, volumes):
        cg, sg = self._get_group_objects(master)
        self._grant(sg, cg)
        self._grant(cg, sg)

    def on_shutdown(self, nodes, master, user, user_shell, volumes):
        cg, sg = self._get_group_objects(master)
        self._revoke(sg, cg)
        self._revoke(cg, sg)