#!/usr/bin/env python
# -*- coding: utf-8 -*-


from copy import deepcopy as copy
from ansible.module_utils.basic import *
import difflib
import commands

EXAMPLES = '''
'''

BINS = dict(
    ipv4='iptables',
    ipv6='ip6tables',
)

params = {
        'protocol': '-p',
        'source': '-s',
        'destination': '-d',
        'match': '-m',
        'jump': '-j',
        'goto': '-g',
        'in_interface': '-i',
        'out_interface': '-o',
        'fragment': '-f',
        'set_counters': '-c',
        'source_port': '--sport',
        'destination_port': '--dport',
        'to_ports': '--to-ports',
        }


def get_params(key):
    if key in params.keys():
        return params[key]
    elif key.startswith('-'):
        return key
    elif ':' in key:
        key, value = key.split(':', 1)
        return get_params(key.strip()) + ' ' + value.strip()
    else:
        return '--' + key


class iptables_Rule(object):
    def __init__(self, nat=None, filters=None, raw=None, mangle=None):
        self.nat_table = nat
        self.filter_table = filters
        self.raw_tables = raw
        self.mangle_tables = mangle

    def dumps_config(self):
        rules = [[]]

        def parse_tables(table, data, default_chain):
            if data:
                rules[0].append("*" + table.strip())
                for chain in data.keys():
                    if chain in default_chain:
                        rules[0].append(":" + chain + " ACCEPT [0:0]")
                    else:
                        rules[0].append(":" + chain + " - [0:0]")

                for chain in data.keys():
                    rule = [parse_rule(x) for x in data[chain]]
                    if rule:
                        chain_rules = reduce(lambda x, y: x + y, rule)
                        rules[0] += ['-A ' + chain + ' ' + x for x in chain_rules]
                rules[0].append("COMMIT")

        parse_tables('nat', self.nat_table, ("PREROUTING", "POSTROUTING", "OUTPUT"))
        parse_tables('filter', self.filter_table, ("INPUT", "FORWARD", "OUTPUT"))
        parse_tables('raw', self.raw_tables, ("PREROUTING", "OUTPUT"))
        parse_tables('mangle', self.mangle_tables, ("PREROUTING", "INPUT", "FORWARD", "OUTPUT", "POSTROUTING"))
        return rules[0]

    def __str__(self):
        return '\n'.join(self.dumps_config()) + '\n'


def parse_rule(rule):
    rule = copy(rule)
    result = ['']
    if not isinstance(rule, (tuple, list, dict)):
        return [str(rule)]

    if isinstance(rule, (tuple, list)):
        value = [parse_rule(x) for x in rule]
        return reduce(lambda x, y: x + y, value)

    for key in rule.keys():
        if isinstance(rule[key], (list, tuple, dict)):
            value = parse_rule(rule[key])
            r = copy(result)
            result = [' '.join((x, get_params(key), y)) for x in r for y in value]
        else:
            result = [x + ' ' + ' '.join((get_params(key), str(rule[key]))) for x in result]

    return result


def main():
    module = AnsibleModule(
        argument_spec={
            "filter": dict(required=False, default=None),
            "nat": dict(required=False, default=None),
            "raw": dict(required=False, default=None),
            "mangle": dict(required=False, default=None),
            "load": dict(required=False, default=True),
            "ip_version": dict(required=False, default='ipv4', choices=['ipv4', 'ipv6'])
        },
        supports_check_mode=True
    )
    iptables = iptables_Rule(
            nat=module.params['nat'],
            filters=module.params['filter'],
            raw=module.params['raw'],
            mangle=module.params['mangle'])

    iptables_file_path = "/tmp/.%s" % BINS[module.params['ip_version']]
    try:
        with open(iptables_file_path) as f:
            old_rules = map(lambda x: x.strip(), f.readlines())
    except:
        old_rules = []

    diff = list(difflib.context_diff(old_rules, iptables.dumps_config()))
    result = {}

    if diff:
        result['changed'] = True
    else:
        result['changed'] = False

    result['failed'] = False

    if not module.check_mode:
        with open(iptables_file_path + ".new", 'w+') as f:
            f.write(str(iptables))

    if not module.check_mode and module.params['load']:
        (status, output) = commands.getstatusoutput(
                "cat %s | %s-restore" % (iptables_file_path + ".new", BINS[module.params['ip_version']]))

        if status == 0:
            commands.getstatusoutput("mv -f %s %s" % (iptables_file_path + ".new", iptables_file_path))
        else:
            result['failed'] = True
            result['msg'] = output

    module.exit_json(**result)


if __name__ == "__main__":
    main()
