python 自动添加iptables规则脚本

#!/usr/bin/env python3
# -*- coding:utf8 -*-
# Description: test
import re
import json
import argparse
import subprocess


class AddIptablesRule:

    @classmethod
    def parameters(cls):
        """
        传递参数
        :return:
        """
        parser = argparse.ArgumentParser()
        parser.add_argument("--port", "-port", help="开放端口: port1|port1:port2:port3", required=True)
        parser.add_argument("--ip", "-ip", help="开放网段: 127.0.0.1|127.0.0.1/8")
        parser.add_argument("--action", "-action", help="执行动作:ACCEPT|REJECT",required=True)
        iptables = parser.parse_args()
        return iptables

    @staticmethod
    def system_command(command):
        """
        调用linux系统命令
        :param command:
        :return:
        """
        shell = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
        stdout, stderr = shell.communicate()
        try:
            return stdout.decode("utf8"), stderr.decode("utf8"), shell.returncode
        except Exception as e:
            return stdout.decode("gbk"), stderr.decode("gbk"), shell.returncode

    def centos_version(self):
        """
        centos版本 6/7
        :return:
        """
        stdout, stderr, return_code = self.system_command('cat /etc/redhat-release')
        centos_version = re.findall('(\d+)',stdout)[0]
        return int(centos_version)  # 6/7

    def add_iptables_rule(self):
        """
        centos6 适用iptables
        :return:
        """

        ip_address_list = self.parameters().ip
        port_list = str(self.parameters().port)
        action = self.parameters().action

        # 允许或者禁止所有端口
        for port in port_list.split(','):
            if action in ['ACCEPT']:
                drop_sport = "iptables -I INPUT -p TCP --sport {} -j DROP".format(port)
                drop_dport = "iptables -I INPUT -p TCP --dport {} -j DROP".format(port)
                self.system_command(drop_sport)
                self.system_command(drop_dport)
            else:
                accept_sport = "iptables -I INPUT -p TCP --sport {} -j ACCEPT".format(port)
                accept_dport = "iptables -I INPUT -p TCP --dport {} -j ACCEPT".format(port)
                self.system_command(accept_sport)
                self.system_command(accept_dport)

        # 允许或者禁止指定端口
        ipaddress = []
        ports = []
        status_code = None
        if ip_address_list:
            for ip_address in ip_address_list.split(','):
                for port in port_list.split(','):
                    if action in ['ACCEPT']:
                        for in_out_port in ["--sport", "--dport"]:
                            accept_port = ["iptables -I INPUT -p TCP"]
                            if ip_address:
                                accept_port.append("-s {}".format(ip_address))
                            if port:
                                accept_port.append("{} {}".format(in_out_port, port))
                            accept_port.append('-j {}'.format(action))
                            iptables_rule = ' '.join(accept_port)
                            status_code = iptables_rule
                            self.system_command(iptables_rule)
                            if port not in ports:
                                ports.append(port)
                            if ip_address not in ipaddress:
                                ipaddress.append(ip_address)

                    else:
                        for in_out_port in ["--sport", "--dport"]:
                            drop_port = ["iptables -I INPUT -p TCP"]
                            if ip_address:
                                drop_port.append("-s {}".format(ip_address))
                            if port:
                                drop_port.append("{} {}".format(in_out_port,port))
                            drop_port.append('-j {}'.format(action))
                            iptables_rule = ' '.join(drop_port)
                            status_code = iptables_rule
                            self.system_command(iptables_rule)
                            if port not in ports:
                                ports.append(port)
                            if ip_address not in ipaddress:
                                ipaddress.append(ip_address)

        status = 'success' if status_code else 'failed'
        return {'action': action, 'ipaddress': ipaddress, 'ports': ports, 'status': status}

    def add_firewall_rule(self):
        """
        centos7适用firewall
        :return:
        """
        ip_address_list = str(self.parameters().ip)
        port_list = str(self.parameters().port)
        action = self.parameters().action

        # 允许或者禁止所有端口
        for port in port_list.split(','):
            if action in ['ACCEPT']:
                drop_port = "firewall-cmd --zone=public --remove-port={}/tcp --permanent".format(port)
                self.system_command(drop_port)

            else:
                accept_port = "firewall-cmd --zone=public --add-port={}/tcp --permanent".format(port)
                self.system_command(accept_port)

        # 允许或者禁止指定端口
        status_code = None
        ipaddress = []
        ports = []
        for ip_address in ip_address_list.split(','):
            for port in port_list.split(','):
                # 允许或者禁止指定端口
                if action in ['ACCEPT']:

                    add_firewall_rule = ['firewall-cmd --permanent --add-rich-rule="rule family="ipv4"']
                    remove_firewall_rule = ['firewall-cmd --permanent --remove-rich-rule="rule family="ipv4"']
                    if ip_address:
                        add_firewall_rule.append('source address="{}"'.format(ip_address))
                        remove_firewall_rule.append('source address="{}"'.format(ip_address))
                    add_firewall_rule.append('port protocol="tcp" port="{}" accept"'.format(port))
                    remove_firewall_rule.append('port protocol="tcp" port="{}" reject"'.format(port))

                    add_firewall_rule = ' '.join(add_firewall_rule)
                    remove_firewall_rule = ' '.join(remove_firewall_rule)
                    status_code = add_firewall_rule
                    self.system_command(remove_firewall_rule)
                    self.system_command(add_firewall_rule)
                    if port not in ports:
                        ports.append(port)
                    if ip_address not in ipaddress:
                        ipaddress.append(ip_address)

                else:
                    add_firewall_rule = ['firewall-cmd --permanent --add-rich-rule="rule family="ipv4"']
                    remove_firewall_rule = ['firewall-cmd --permanent --remove-rich-rule="rule family="ipv4"']
                    if ip_address:
                        add_firewall_rule.append('source address="{}"'.format(ip_address))
                        remove_firewall_rule.append('source address="{}"'.format(ip_address))
                    add_firewall_rule.append('port protocol="tcp" port="{}" reject"'.format(port))
                    remove_firewall_rule.append('port protocol="tcp" port="{}" accept"'.format(port))

                    add_firewall_rule = ' '.join(add_firewall_rule)
                    remove_firewall_rule = ' '.join(remove_firewall_rule)
                    status_code = add_firewall_rule
                    self.system_command(remove_firewall_rule)
                    self.system_command(add_firewall_rule)
                    if port not in ports:
                        ports.append(port)
                    if ip_address not in ipaddress:
                        ipaddress.append(ip_address)

        self.system_command('firewall-cmd --reload ')
        status = 'success' if status_code else 'failed'
        return {'action': action, 'ipaddress': ipaddress, 'ports': ports, 'status': status}

    def select_version(self):
        """
        选择centos版本
        :return:
        """
        if str(self.centos_version()).startswith('7'):
            print(json.dumps(self.add_firewall_rule()) )
            return json.dumps(self.add_firewall_rule())
        else:
            print(json.dumps(self.add_iptables_rule()) )
            return json.dumps(self.add_iptables_rule())


if __name__ == '__main__':
    a = AddIptablesRule()
    a.select_version()

使用方法

python3 add_iptables_rule.py --port '8088,9090' --ip '152.0.0.1,127.0.0.1' --action ACCEPT