update_security_group.py

#!/usr/bin/env python

'''clear "ssh-temp" and create a new rule for all traffic from current IP'''

def get_my_public_ip():
    '''returns outbound (public) IP address in string format'''
    from httplib import HTTPConnection
    ipconn = HTTPConnection('icanhazip.com')
    ipconn.request('GET', '/')
    myip = ipconn.getresponse()
    myip = myip.read()
    ipconn.close()
    return myip.strip()

def get_security_group(sg_name):
    '''return a boto security group object from string name'''
    from boto.ec2 import connect_to_region
    conn = connect_to_region('us-west-2')
    try:
        return conn.get_all_security_groups(filters={'group-name': sg_name})[0]
    except IndexError:
        return None

def revoke_all_rules_from_group(sg_name):
    '''delete all rules from sg_name'''
    group = get_security_group(sg_name)
    if group != None:
        for rule in group.rules:
            for grant in rule.grants:
                group.revoke(
                        ip_protocol=rule.ip_protocol,
                        from_port=rule.from_port,
                        to_port=rule.to_port,
                        cidr_ip = grant
                        )
    else:
        raise KeyError('No group {} found.'.format(sg_name))

def grant_my_ip(sg_name):
    '''grant all traffic access to sg_name from current IP'''
    my_ip = get_my_public_ip()
    group = get_security_group(sg_name)
    group.authorize(
            ip_protocol='-1',
            from_port=None,
            to_port=None,
            cidr_ip=my_ip+'/32'
            )

def revoke_my_ip(sg_name):
    '''revoke access to sg_name from current IP'''
    my_ip = get_my_public_ip()
    group = get_security_group(sg_name)
    for rule in group.rules:
        for grant in rule.grants:
            if grant.split('/')[0] == my_ip:
                group.revoke(
                        ip_protocol=rule.ip_protocol,
                        from_port=rule.from_port,
                        to_port=rule.to_port,
                        cidr_ip = grant
                        )

def main():
    '''clear "ssh-temp" and create a new rule for all traffic from current IP'''
    import sys
    if len(sys.argv) != 3:
        sys.exit(1)
    action = sys.argv[1]
    group = sys.argv[2]
    if action == 'reset':
        revoke_all_rules_from_group(group)
        grant_my_ip(group)
    elif action == 'add':
        grant_my_ip(group)
    elif action == 'del':
        revoke_my_ip(group)

if __name__ == '__main__':
    main()