Dzmitry Stremkouski | e353ce3 | 2018-08-30 17:22:32 +0200 | [diff] [blame] | 1 | # -*- coding: utf-8 -*- |
| 2 | from os import chmod,remove |
| 3 | from time import time |
| 4 | from subprocess import Popen,PIPE |
| 5 | |
| 6 | def get_tables(family="ipv4"): |
| 7 | |
| 8 | ''' List iptables tables |
| 9 | |
| 10 | :param family: iptables ip family version. type: str |
| 11 | |
| 12 | ''' |
| 13 | |
| 14 | if family == "ipv4": |
| 15 | cmd = 'iptables-save' |
| 16 | elif family == "ipv6": |
| 17 | cmd = 'ip6tables-save' |
| 18 | else: |
| 19 | return "Invalid ip family specified. Use either ipv4 or ipv6" |
| 20 | |
| 21 | tables = [] |
| 22 | try: |
| 23 | tables_list = Popen(cmd, shell=True, stdout=PIPE) |
| 24 | while True: |
| 25 | line = tables_list.stdout.readline() |
| 26 | if line != '': |
| 27 | if line[0] == "*": |
| 28 | tables.append(line.rstrip()[1:]) |
| 29 | else: |
| 30 | break |
| 31 | except: |
| 32 | return "Error getting list of tables" |
| 33 | |
| 34 | return tables |
| 35 | |
| 36 | def get_chains(family="ipv4", table="filter"): |
| 37 | |
| 38 | ''' List iptables chains |
| 39 | |
| 40 | :param family: iptables ip family version. type: str |
| 41 | :param table: Lookup chains for this table. type: str |
| 42 | |
| 43 | ''' |
| 44 | |
| 45 | if family == "ipv4": |
| 46 | cmd = 'iptables-save' |
| 47 | elif family == "ipv6": |
| 48 | cmd = 'ip6tables-save' |
| 49 | else: |
| 50 | return "Invalid ip family specified. Use either ipv4 or ipv6" |
| 51 | |
| 52 | cmd += ' -t ' + table |
| 53 | |
| 54 | chains = [] |
| 55 | try: |
| 56 | chains_list = Popen(cmd, shell=True, stdout=PIPE) |
| 57 | while True: |
| 58 | line = chains_list.stdout.readline() |
| 59 | if line != '': |
| 60 | if line[0] == ":": |
| 61 | chains.append(line.rstrip()[1:].split(' ')[0]) |
| 62 | else: |
| 63 | break |
| 64 | except: |
| 65 | return "Error getting list of chains" |
| 66 | |
| 67 | return chains |
| 68 | |
| 69 | def get_structure(family="ipv4"): |
| 70 | |
| 71 | ''' Get structure of all chains in all tables |
| 72 | |
| 73 | :param family: iptables ip family version. type: str |
| 74 | |
| 75 | ''' |
| 76 | |
| 77 | if family == "ipv4": |
| 78 | cmd = 'iptables-save' |
| 79 | elif family == "ipv6": |
| 80 | cmd = 'ip6tables-save' |
| 81 | else: |
| 82 | return "Invalid ip family specified. Use either ipv4 or ipv6" |
| 83 | |
| 84 | tables = [] |
| 85 | tables_list = Popen(cmd, shell=True, stdout=PIPE) |
| 86 | while True: |
| 87 | line = tables_list.stdout.readline() |
| 88 | if line != '': |
| 89 | line = line.rstrip().lstrip() |
| 90 | if line[0] == "*": |
| 91 | elem = {} |
| 92 | table_name = line[1:].split(' ')[0] |
| 93 | elem[table_name] = [] |
| 94 | if line[0] == ":": |
| 95 | elem[table_name].append(line[1:].split(' ')[0]) |
| 96 | if line == "COMMIT": |
| 97 | tables.append(elem) |
| 98 | else: |
| 99 | break |
| 100 | |
| 101 | return tables |
| 102 | |
| 103 | def run_script(script): |
| 104 | |
| 105 | ''' Execute local script |
| 106 | |
| 107 | :param script: script to be executed, storad localy: str |
| 108 | |
| 109 | ''' |
| 110 | |
| 111 | chmod(script, 0o700) |
| 112 | process = Popen([script],stdout=PIPE,stderr=PIPE) |
| 113 | process.wait() |
| 114 | code = process.returncode |
| 115 | remove(script) |
| 116 | return code |
| 117 | |
| 118 | |
| 119 | def flush_all(family="ipv4"): |
| 120 | |
| 121 | ''' Flush all chains in all tables |
| 122 | |
| 123 | :param family: iptables ip family version. type: str |
| 124 | |
| 125 | ''' |
| 126 | |
| 127 | if family == "ipv4": |
| 128 | cmd = 'iptables' |
| 129 | rmmod = 'iptable_' |
| 130 | elif family == "ipv6": |
| 131 | cmd = 'ip6tables' |
| 132 | rmmod = 'ip6table_' |
| 133 | else: |
| 134 | return "Invalid ip family specified. Use either ipv4 or ipv6" |
| 135 | |
| 136 | tables = get_structure(family) |
| 137 | |
| 138 | f_name = '/tmp/' + cmd + '-flush-' + str(time()).split('.')[0] + '.sh' |
| 139 | |
| 140 | with open(f_name, 'w') as f: |
| 141 | f.write('#!/bin/sh\n') |
| 142 | for table in tables: |
| 143 | for var in enumerate(table): |
| 144 | t_name = var[1] |
| 145 | for chain in table[t_name]: |
| 146 | f.write(cmd + ' -t ' + t_name + " -F " + chain + '\n') |
| 147 | if chain not in ['INPUT','FORWARD','OUTPUT','PREROUTING','POSTROUTING']: |
| 148 | f.write(cmd + ' -t ' + t_name + " -X " + chain + '\n') |
| 149 | f.write('rmmod ' + rmmod + t_name + '\n') |
| 150 | |
| 151 | return run_script(f_name) |
| 152 | |
| 153 | def set_policy_all(family="ipv4", policy="ACCEPT"): |
| 154 | |
| 155 | ''' Set policy for all chains in all tables |
| 156 | |
| 157 | :param family: iptables ip family version. type: str |
| 158 | :param policy: iptables chain policy. type: str |
| 159 | |
| 160 | ''' |
| 161 | |
| 162 | if family == "ipv4": |
| 163 | cmd = 'iptables' |
| 164 | elif family == "ipv6": |
| 165 | cmd = 'ip6tables' |
| 166 | else: |
| 167 | return "Invalid ip family specified. Use either ipv4 or ipv6" |
| 168 | |
| 169 | tables = get_structure(family) |
| 170 | |
| 171 | f_name = '/tmp/' + cmd + '-policy-' + str(time()).split('.')[0] + '.sh' |
| 172 | |
| 173 | with open(f_name, 'w') as f: |
| 174 | f.write('#!/bin/sh\n') |
| 175 | for table in tables: |
| 176 | for var in enumerate(table): |
| 177 | t_name = var[1] |
| 178 | for chain in table[t_name]: |
| 179 | f.write(cmd + ' -t ' + t_name + " -P " + chain + ' ' + policy + '\n') |
| 180 | |
| 181 | return run_script(f_name) |
| 182 | |
| 183 | def remove_stale_tables(config_file, family="ipv4"): |
| 184 | |
| 185 | ''' Remove tables which are not in config file |
| 186 | to prevet flushing all the tables |
| 187 | |
| 188 | :param family: iptables ip family version. type: str |
| 189 | :param config_file: iptables rules persistent config file. type: str |
| 190 | |
| 191 | ''' |
| 192 | |
| 193 | if family == "ipv4": |
| 194 | cmd = 'iptables' |
| 195 | rmmod = 'iptable_' |
| 196 | elif family == "ipv6": |
| 197 | cmd = 'ip6tables' |
| 198 | rmmod = 'ip6table_' |
| 199 | else: |
| 200 | return "Invalid ip family specified. Use either ipv4 or ipv6" |
| 201 | |
| 202 | runtime_tables = get_tables(family) |
| 203 | |
| 204 | config_tables = [] |
| 205 | for line in open(config_file, 'r'): |
| 206 | if line != '': |
| 207 | if line[0] == "*": |
| 208 | config_tables.append(line.rstrip()[1:]) |
| 209 | |
| 210 | runtime_tables.sort() |
| 211 | config_tables.sort() |
| 212 | diff = list(set(runtime_tables) - set(config_tables)) |
| 213 | |
| 214 | if diff != []: |
| 215 | tables = get_structure(family) |
| 216 | f_name = '/tmp/' + cmd + '-flush-' + str(time()).split('.')[0] + '.sh' |
| 217 | with open(f_name, 'w') as f: |
| 218 | f.write('#!/bin/sh\n') |
| 219 | for table in tables: |
| 220 | for var in enumerate(table): |
| 221 | t_name = var[1] |
| 222 | if t_name in diff: |
| 223 | for chain in table[t_name]: |
| 224 | f.write(cmd + ' -t ' + t_name + " -F " + chain + '\n') |
| 225 | if chain not in ['INPUT','FORWARD','OUTPUT','PREROUTING','POSTROUTING']: |
| 226 | f.write(cmd + ' -t ' + t_name + " -X " + chain + '\n') |
| 227 | f.write('rmmod ' + rmmod + t_name + '\n') |
| 228 | |
| 229 | return run_script(f_name) |
| 230 | else: |
| 231 | return |