blob: 90d421e416e4d0f8571ebf311a2e9367bef88c3d [file] [log] [blame]
Dzmitry Stremkouskie353ce32018-08-30 17:22:32 +02001# -*- coding: utf-8 -*-
2from os import chmod,remove
3from time import time
4from subprocess import Popen,PIPE
5
6def 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
36def 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
69def 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
103def 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
119def 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
153def 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
183def 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