iptables

import os
import re


class Iptables(object):

    _iptables = 'iptables'

    @classmethod
    def get_chain(cls):
        output = os.popen('%s -n -v -L -t filter -x' % cls._iptables).read()
        # print repr(output)
        INPUT_lines = re.findall('Chain INPUT(?:[\s\S]+?)destination(?:[\s\S]+?)\n([\s\S]+?)Chain ', output)
        # FORWARD_lines = re.findall('Chain FORWARD(?:[\s\S]+?)destination(?:[\s\S]+?)\n([\s\S]+?)Chain ', output)
        OUTPUT_lines = re.findall('Chain OUTPUT(?:[\s\S]+?)destination(?:[\s\S]+?)\n([\s\S]+)', output)

        chain_INPUT = []
        for line in INPUT_lines[0].split('\n'):
            arr_line = line.split()
            if len(arr_line) != 10:
                continue
            chain = {'type': arr_line[2], 'port': arr_line[9][4:], 'bytes': arr_line[1]}
            chain_INPUT.append(chain)
            
        chain_OUTPUT = []
        if len(OUTPUT_lines) > 0:
            for line in OUTPUT_lines[0].split('\n'):
                arr_line = line.split()
                if len(arr_line) != 10:
                    continue
                chain = {'type': arr_line[2], 'port': arr_line[9][4:], 'bytes': arr_line[1]}
                chain_OUTPUT.append(chain)

        return chain_INPUT, chain_OUTPUT

    @classmethod
    def add_chain(cls, chain, type, port):
        if chain == 'INPUT':
            protstr = 'dport'
        else:
            protstr = 'sport'
        output = os.popen('%s -I %s -p %s --%s %s' % (cls._iptables, chain, type, protstr, port))


class Ip6tables(Iptables):

    _iptables = 'ip6tables'

    @classmethod
    def get_chain(cls):
        output = os.popen('%s -n -v -L -t filter -x' % cls._iptables).read()
        # print repr(output)
        INPUT_lines = re.findall('Chain INPUT(?:[\s\S]+?)destination(?:[\s\S]+?)\n([\s\S]+?)Chain ', output)
        # FORWARD_lines = re.findall('Chain FORWARD(?:[\s\S]+?)destination(?:[\s\S]+?)\n([\s\S]+?)Chain ', output)
        OUTPUT_lines = re.findall('Chain OUTPUT(?:[\s\S]+?)destination(?:[\s\S]+?)\n([\s\S]+)', output)

        chain_INPUT = []
        for line in INPUT_lines[0].split('\n'):
            arr_line = line.split()
            if len(arr_line) != 9:
                continue
            chain = {'type': arr_line[2], 'port': arr_line[8][4:], 'bytes': arr_line[1]}
            chain_INPUT.append(chain)
            
        chain_OUTPUT = []
        if len(OUTPUT_lines) > 0:
            for line in OUTPUT_lines[0].split('\n'):
                arr_line = line.split()
                if len(arr_line) != 9:
                    continue
                chain = {'type': arr_line[2], 'port': arr_line[8][4:], 'bytes': arr_line[1]}
                chain_OUTPUT.append(chain)

        return chain_INPUT, chain_OUTPUT

"""
chain_INPUT, chain_OUTPUT = Iptables.get_chain()
if len(chain_INPUT) == 0:
    Iptables.add_chain('INPUT', 'tcp', 80)
    Iptables.add_chain('INPUT', 'udp', 80)
    Iptables.add_chain('OUTPUT', 'tcp', 80)
    Iptables.add_chain('OUTPUT', 'udp', 80)
chain_INPUT, chain_OUTPUT = Ip6tables.get_chain()
print chain_INPUT, chain_OUTPUT
if len(chain_INPUT) == 0:
    Ip6tables.add_chain('INPUT', 'tcp', 80)
    Ip6tables.add_chain('INPUT', 'udp', 80)
    Ip6tables.add_chain('OUTPUT', 'tcp', 80)
    Ip6tables.add_chain('OUTPUT', 'udp', 80)
"""