Changed to a structure with one security group for all rules with source security group

This involves generating all pairs of addresses from source and target groups. Hopefully this will not give too many addresses

iptables_configurator.py

@@ -34,12 +34,6 @@ def etcd_get_dirs(etcd_client, prefix):
     return result
 
 
-def escape_group_name(name):
-    name = name.replace(" ", "_")
-    name = name.replace("\t", "_")
-    return name
-
-
 class Generator(object):
     def __init__(self, cert, key, cacert, logfile=None):
         self.etcd_client = etcd.Etcd(ssl_key=key, ssl_cert=cert, verify=cacert)
@@ -49,25 +43,17 @@ class Generator(object):
             logging.getLogger("").addHandler(handler)
         self.range = [150, 200]
         self.prefix = "158.38.213."
-        self.groups = []
-
-    def short_group_id(self, group_id):
-        if group_id not in self.groups:
-            self.groups.append(group_id)
-        return self.groups.index(group_id)
 
     def output(self, line):
         self.output_file.write(line)
         self.output_file.write("\n")
 
     def create_ipset(self, name, set_type):
-        name = escape_group_name(name)
         for family in ("inet", "inet6"):
             self.output("ipset create {}_{} {} family {}".format(name, family,
                                                                  set_type, family))
 
-    def add_ipset_member(self, name, member, protocol="tcp", port=None, net=None):
-        name = escape_group_name(name)
+    def add_ipset_member(self, name, member, protocol="tcp", port=None, net=None, source=None):
         suffix = ""
         if port is not None:
             suffix += ",{}:{}".format(protocol, port)
@@ -75,19 +61,28 @@ class Generator(object):
             suffix += ",{}".format(net)
         member = member.lower()
         if member in self.addresses_v4 and (net is None or "." in net):
-            self.output("ipset add {}_inet {}{}".format(name, self.addresses_v4[member],
-                                                        suffix))
+            if source is not None:
+                if source in self.addresses_v4:
+                    suffix4 = "{},{}".format(suffix, self.addresses_v4[source])
+                    self.output("ipset add {}_inet {}{}".format(name, self.addresses_v4[member],
+                                                                suffix4))
+            else:
+                self.output("ipset add {}_inet {}{}".format(name, self.addresses_v4[member],
+                                                            suffix))
+
         if member in self.addresses_v6 and (net is None or ":" in net):
-            self.output("ipset add {}_inet6 {}{}".format(name, self.addresses_v6[member],
-                                                         suffix))
+            if source is not None:
+                if source in self.addresses_v6:
+                    suffix6 = "{},{}".format(suffix, self.addresses_v6[source])
+                    self.output("ipset add {}_inet6 {}{}".format(name, self.addresses_v6[member],
+                                                                 suffix6))
+            else:
+                self.output("ipset add {}_inet6 {}{}".format(name, self.addresses_v6[member],
+                                                             suffix))
 
     def process_security_group(self, group_id, name):
         rules = security_groups.get_group_rules(self.etcd_client, group_id)
         _, members = security_groups.get_group_members(self.etcd_client, group_id)
-        source_group = "source_{}".format(self.short_group_id(group_id))
-        self.create_ipset(source_group, "hash:ip")
-        for member in members:
-            self.add_ipset_member(source_group, member)
         for rule in rules:
             for member in members:
                 if rule["source_type"] == "any":
@@ -99,12 +94,9 @@ class Generator(object):
                                           rule["source_cidr"])
                 elif rule["source_type"] == "security_group":
                     source_group = rule["source_security_group"]
-                    set_name = "rules_from_group_{}".format(self.short_group_id(source_group))
-                    if source_group not in self.source_sets:
-                        self.create_ipset(set_name, "hash:ip,port")
-                        self.source_sets[source_group] = set_name
-                    self.add_ipset_member(set_name, member, rule["protocol"],
-                                          rule["destination_port"])
+                    for source in security_groups.get_group_members(self.etcd_client, source_group)[1]:
+                        self.add_ipset_member("rules_from_sg", member, rule["protocol"],
+                                              rule["destination_port"], source=source)
                 else:
                     logging.warning("Unhandled source type: %s", rule["source_type"])
 
@@ -126,11 +118,11 @@ class Generator(object):
         index = None
         self.addresses_v4 = self.get_addresses("ipv4")
         self.addresses_v6 = self.get_addresses("ipv6_public")
-        self.source_sets = {}
         self.output_file = open("output.sh", "w")
 
         self.create_ipset("rules_from_any", "hash:ip,port")
         self.create_ipset("rules_from_cidr", "hash:ip,port,net")
+        self.create_ipset("rules_from_sg", "hash:ip,port,ip")
         groups = security_groups.get_security_groups(self.etcd_client)
         for group_id, name in groups.items():
             self.process_security_group(group_id, name)
@@ -138,9 +130,8 @@ class Generator(object):
         self.output("ip6tables -A FORWARD -m set --match-set rules_from_any_inet6 dst,dst -j ACCEPT")
         self.output("iptables -A FORWARD -m set --match-set rules_from_cidr_inet dst,dst,src -j ACCEPT")
         self.output("ip6tables -A FORWARD -m set --match-set rules_from_cidr_inet6 dst,dst,src -j ACCEPT")
-        for source_group, destination_set in self.source_sets.items():
-            self.output("iptables -A FORWARD -m set --match-set {}_inet dst,dst --match-set source_{}_inet src -j ACCEPT".format(destination_set, self.short_group_id(source_group)))
-            self.output("ip6tables -A FORWARD -m set --match-set {}_inet6 dst,dst --match-set source_{}_inet6 src -j ACCEPT".format(destination_set, self.short_group_id(source_group)))
+        self.output("iptables -A FORWARD -m set --match-set rules_from_sg_inet dst,dst,src -j ACCEPT")
+        self.output("ip6tables -A FORWARD -m set --match-set rules_from_sg_inet6 dst,dst,src -j ACCEPT")
         return index
 
     def main(self):