#!/usr/bin/env python

#

# (C) 2014 by Ana Rey Botello <anarey@gmail.com>

#

# Based on iptables-test.py:

# (C) 2012 by Pablo Neira Ayuso <pablo@netfilter.org>"

#

# This program is free software; you can redistribute it and/or modify

# it under the terms of the GNU General Public License as published by

# the Free Software Foundation; either version 2 of the License, or

# (at your option) any later version.

#

# Thanks to the Outreach Program for Women (OPW) for sponsoring this test

# infrastructure.

from __future__ import print_function

import sys

import os

import argparse

import signal

import json

import traceback

import tempfile

TESTS_PATH = os.path.dirname(os.path.abspath(__file__))

sys.path.insert(0, os.path.join(TESTS_PATH, '../../py/'))

os.environ['TZ'] = 'UTC-2'

from nftables import Nftables

TESTS_DIRECTORY = ["any", "arp", "bridge", "inet", "ip", "ip6"]

LOGFILE = "/tmp/nftables-test.log"

log_file = None

table_list = []

chain_list = []

all_set = dict()

obj_list = []

signal_received = 0

class Colors:

    if sys.stdout.isatty():

        HEADER = '\033[95m'

        GREEN = '\033[92m'

        YELLOW = '\033[93m'

        RED = '\033[91m'

        ENDC = '\033[0m'

    else:

        HEADER = ''

        GREEN = ''

        YELLOW = ''

        RED = ''

        ENDC = ''

class Chain:

    """Class that represents a chain"""

    def __init__(self, name, config, lineno):

        self.name = name

        self.config = config

        self.lineno = lineno

    def __eq__(self, other):

        return self.__dict__ == other.__dict__

    def __str__(self):

        return "%s" % self.name

class Table:

    """Class that represents a table"""

    def __init__(self, family, name, chains):

        self.family = family

        self.name = name

        self.chains = chains

    def __eq__(self, other):

        return self.__dict__ == other.__dict__

    def __str__(self):

        return "%s %s" % (self.family, self.name)

class Set:

    """Class that represents a set"""

    def __init__(self, family, table, name, type, timeout, flags):

        self.family = family

        self.table = table

        self.name = name

        self.type = type

        self.timeout = timeout

        self.flags = flags

    def __eq__(self, other):

        return self.__dict__ == other.__dict__

class Obj:

    """Class that represents an object"""

    def __init__(self, table, family, name, type, spcf):

        self.table = table

        self.family = family

        self.name = name

        self.type = type

        self.spcf = spcf

    def __eq__(self, other):

        return self.__dict__ == other.__dict__

def print_msg(reason, errstr, filename=None, lineno=None, color=None):

    '''

    Prints a message with nice colors, indicating file and line number.

    '''

    color_errstr = "%s%s%s" % (color, errstr, Colors.ENDC)

    if filename and lineno:

        sys.stderr.write("%s: %s line %d: %s\n" %

                         (filename, color_errstr, lineno + 1, reason))

    else:

        sys.stderr.write("%s %s\n" % (color_errstr, reason))

    sys.stderr.flush() # So that the message stay in the right place.

def print_error(reason, filename=None, lineno=None):

    print_msg(reason, "ERROR:", filename, lineno, Colors.RED)

def print_warning(reason, filename=None, lineno=None):

    print_msg(reason, "WARNING:", filename, lineno, Colors.YELLOW)

def print_info(reason, filename=None, lineno=None):

    print_msg(reason, "INFO:", filename, lineno, Colors.GREEN)

def color_differences(rule, other, color):

    rlen = len(rule)

    olen = len(other)

    out = ""

    i = 0

    # find equal part at start

    for i in range(rlen):

        if i >= olen or rule[i] != other[i]:

            break

    if i > 0:

        out += rule[:i]

        rule = rule[i:]

        other = other[i:]

        rlen = len(rule)

        olen = len(other)

    # find equal part at end

    for i in range(1, rlen + 1):

        if i > olen or rule[rlen - i] != other[olen - i]:

            i -= 1

            break

    if rlen > i:

        out += color + rule[:rlen - i] + Colors.ENDC

        rule = rule[rlen - i:]

    out += rule

    return out

def print_differences_warning(filename, lineno, rule1, rule2, cmd):

    colored_rule1 = color_differences(rule1, rule2, Colors.YELLOW)

    colored_rule2 = color_differences(rule2, rule1, Colors.YELLOW)

    reason = "'%s': '%s' mismatches '%s'" % (cmd, colored_rule1, colored_rule2)

    print_warning(reason, filename, lineno)

def print_differences_error(filename, lineno, cmd):

    reason = "'%s': Listing is broken." % cmd

    print_error(reason, filename, lineno)

def table_exist(table, filename, lineno):

    '''

    Exists a table.

    '''

    cmd = "list table %s" % table

    ret = execute_cmd(cmd, filename, lineno)

    return True if (ret == 0) else False

def table_flush(table, filename, lineno):

    '''

    Flush a table.

    '''

    cmd = "flush table %s" % table

    execute_cmd(cmd, filename, lineno)

    return cmd

def table_create(table, filename, lineno):

    '''

    Adds a table.

    '''

    # We check if table exists.

    if table_exist(table, filename, lineno):

        reason = "Table %s already exists" % table

        print_error(reason, filename, lineno)

        return -1

    table_list.append(table)

    # We add a new table

    cmd = "add table %s" % table

    ret = execute_cmd(cmd, filename, lineno)

    if ret != 0:

        reason = "Cannot " + cmd

        print_error(reason, filename, lineno)

        table_list.remove(table)

        return -1

    # We check if table was added correctly.

    if not table_exist(table, filename, lineno):

        table_list.remove(table)

        reason = "I have just added the table %s " \

                 "but it does not exist. Giving up!" % table

        print_error(reason, filename, lineno)

        return -1

    for table_chain in table.chains:

        chain = chain_get_by_name(table_chain)

        if chain is None:

            reason = "The chain %s requested by table %s " \

                     "does not exist." % (table_chain, table)

            print_error(reason, filename, lineno)

        else:

            chain_create(chain, table, filename)

    return 0

def table_delete(table, filename=None, lineno=None):

    '''

    Deletes a table.

    '''

    if not table_exist(table, filename, lineno):

        reason = "Table %s does not exist but I added it before." % table

        print_error(reason, filename, lineno)

        return -1

    cmd = "delete table %s" % table

    ret = execute_cmd(cmd, filename, lineno)

    if ret != 0:

        reason = "%s: I cannot delete table %s. Giving up!" % (cmd, table)

        print_error(reason, filename, lineno)

        return -1

    if table_exist(table, filename, lineno):

        reason = "I have just deleted the table %s " \

                 "but it still exists." % table

        print_error(reason, filename, lineno)

        return -1

    return 0

def chain_exist(chain, table, filename):

    '''

    Checks a chain

    '''

    cmd = "list chain %s %s" % (table, chain)

    ret = execute_cmd(cmd, filename, chain.lineno)

    return True if (ret == 0) else False

def chain_create(chain, table, filename):

    '''

    Adds a chain

    '''

    if chain_exist(chain, table, filename):

        reason = "This chain '%s' exists in %s. I cannot create " \

                 "two chains with same name." % (chain, table)

        print_error(reason, filename, chain.lineno)

        return -1

    cmd = "add chain %s %s" % (table, chain)

    if chain.config:

        cmd += " { %s; }" % chain.config

    ret = execute_cmd(cmd, filename, chain.lineno)

    if ret != 0:

        reason = "I cannot create the chain '%s'" % chain

        print_error(reason, filename, chain.lineno)

        return -1

    if not chain_exist(chain, table, filename):

        reason = "I have added the chain '%s' " \

                 "but it does not exist in %s" % (chain, table)

        print_error(reason, filename, chain.lineno)

        return -1

    return 0

def chain_delete(chain, table, filename=None, lineno=None):

    '''

    Flushes and deletes a chain.

    '''

    if not chain_exist(chain, table, filename):

        reason = "The chain %s does not exist in %s. " \

                 "I cannot delete it." % (chain, table)

        print_error(reason, filename, lineno)

        return -1

    cmd = "flush chain %s %s" % (table, chain)

    ret = execute_cmd(cmd, filename, lineno)

    if ret != 0:

        reason = "I cannot " + cmd

        print_error(reason, filename, lineno)

        return -1

    cmd = "delete chain %s %s" % (table, chain)

    ret = execute_cmd(cmd, filename, lineno)

    if ret != 0:

        reason = "I cannot " + cmd

        print_error(reason, filename, lineno)

        return -1

    if chain_exist(chain, table, filename):

        reason = "The chain %s exists in %s. " \

                 "I cannot delete this chain" % (chain, table)

        print_error(reason, filename, lineno)

        return -1

    return 0

def chain_get_by_name(name):

    for chain in chain_list:

        if chain.name == name:

            break

    else:

        chain = None

    return chain

def set_add(s, test_result, filename, lineno):

    '''

    Adds a set.

    '''

    if not table_list:

        reason = "Missing table to add rule"

        print_error(reason, filename, lineno)

        return -1

    for table in table_list:

        s.table = table.name

        s.family = table.family

        if _set_exist(s, filename, lineno):

            reason = "Set %s already exists in %s" % (s.name, table)

            print_error(reason, filename, lineno)

            return -1

        flags = s.flags

        if flags != "":

            flags = "flags %s; " % flags

        cmd = "add set %s %s { type %s;%s %s}" % (table, s.name, s.type, s.timeout, flags)

        ret = execute_cmd(cmd, filename, lineno)

        if (ret == 0 and test_result == "fail") or \

                (ret != 0 and test_result == "ok"):

            reason = "%s: I cannot add the set %s" % (cmd, s.name)

            print_error(reason, filename, lineno)

            return -1

        if not _set_exist(s, filename, lineno):

            reason = "I have just added the set %s to " \

                     "the table %s but it does not exist" % (s.name, table)

            print_error(reason, filename, lineno)

            return -1

    return 0

def set_add_elements(set_element, set_name, state, filename, lineno):

    '''

    Adds elements to the set.

    '''

    if not table_list:

        reason = "Missing table to add rules"

        print_error(reason, filename, lineno)

        return -1

    for table in table_list:

        # Check if set exists.

        if (not set_exist(set_name, table, filename, lineno) or

                    set_name not in all_set) and state == "ok":

            reason = "I cannot add an element to the set %s " \

                     "since it does not exist." % set_name

            print_error(reason, filename, lineno)

            return -1

        element = ", ".join(set_element)

        cmd = "add element %s %s { %s }" % (table, set_name, element)

        ret = execute_cmd(cmd, filename, lineno)

        if (state == "fail" and ret == 0) or (state == "ok" and ret != 0):

            test_state = "This rule should have failed."

            reason = cmd + ": " + test_state

            print_error(reason, filename, lineno)

            return -1

        # Add element into all_set.

        if ret == 0 and state == "ok":

            for e in set_element:

                all_set[set_name].add(e)

    return 0

def set_delete_elements(set_element, set_name, table, filename=None,

                        lineno=None):

    '''

    Deletes elements in a set.

    '''

    for element in set_element:

        cmd = "delete element %s %s { %s }" % (table, set_name, element)

        ret = execute_cmd(cmd, filename, lineno)

        if ret != 0:

            reason = "I cannot delete element %s " \

                     "from the set %s" % (element, set_name)

            print_error(reason, filename, lineno)

            return -1

    return 0

def set_delete(table, filename=None, lineno=None):

    '''

    Deletes set and its content.

    '''

    for set_name in all_set.keys():

        # Check if exists the set

        if not set_exist(set_name, table, filename, lineno):

            reason = "The set %s does not exist, " \

                     "I cannot delete it" % set_name

            print_error(reason, filename, lineno)

            return -1

        # We delete all elements in the set

        set_delete_elements(all_set[set_name], set_name, table, filename,

                            lineno)

        # We delete the set.

        cmd = "delete set %s %s" % (table, set_name)

        ret = execute_cmd(cmd, filename, lineno)

        # Check if the set still exists after I deleted it.

        if ret != 0 or set_exist(set_name, table, filename, lineno):

            reason = "Cannot remove the set " + set_name

            print_error(reason, filename, lineno)

            return -1

    return 0

def set_exist(set_name, table, filename, lineno):

    '''

    Check if the set exists.

    '''

    cmd = "list set %s %s" % (table, set_name)

    ret = execute_cmd(cmd, filename, lineno)

    return True if (ret == 0) else False

def _set_exist(s, filename, lineno):

    '''

    Check if the set exists.

    '''

    cmd = "list set %s %s %s" % (s.family, s.table, s.name)

    ret = execute_cmd(cmd, filename, lineno)

    return True if (ret == 0) else False

def set_check_element(rule1, rule2):

    '''

    Check if element exists in anonymous sets.

    '''

    pos1 = rule1.find("{")

    pos2 = rule2.find("{")

    if (rule1[:pos1] != rule2[:pos2]):

        return False

    end1 = rule1.find("}")

    end2 = rule2.find("}")

    if (pos1 != -1) and (pos2 != -1) and (end1 != -1) and (end2 != -1):

        list1 = (rule1[pos1 + 1:end1].replace(" ", "")).split(",")

        list2 = (rule2[pos2 + 1:end2].replace(" ", "")).split(",")

        list1.sort()

        list2.sort()

        if list1 != list2:

            return False

        return rule1[end1:] == rule2[end2:]

    return False

def obj_add(o, test_result, filename, lineno):

    '''

    Adds an object.

    '''

    if not table_list:

        reason = "Missing table to add rule"

        print_error(reason, filename, lineno)

        return -1

    for table in table_list:

        o.table = table.name

        o.family = table.family

        obj_handle = o.type + " " + o.name

        if _obj_exist(o, filename, lineno):

            reason = "The %s already exists in %s" % (obj_handle, table)

            print_error(reason, filename, lineno)

            return -1

        cmd = "add %s %s %s %s" % (o.type, table, o.name, o.spcf)

        ret = execute_cmd(cmd, filename, lineno)

        if (ret == 0 and test_result == "fail") or \

                (ret != 0 and test_result == "ok"):

            reason = "%s: I cannot add the %s" % (cmd, obj_handle)

            print_error(reason, filename, lineno)

            return -1

        exist = _obj_exist(o, filename, lineno)

        if exist:

            if test_result == "ok":

                 return 0

            reason = "I added the %s to the table %s " \

                     "but it should have failed" % (obj_handle, table)

            print_error(reason, filename, lineno)

            return -1

        if test_result == "fail":

            return 0

        reason = "I have just added the %s to " \

                 "the table %s but it does not exist" % (obj_handle, table)

        print_error(reason, filename, lineno)

        return -1

def obj_delete(table, filename=None, lineno=None):

    '''

    Deletes object.

    '''

    for o in obj_list:

        obj_handle = o.type + " " + o.name

        # Check if exists the obj

        if not obj_exist(o, table, filename, lineno):

            reason = "The %s does not exist, I cannot delete it" % obj_handle

            print_error(reason, filename, lineno)

            return -1

        # We delete the object.

        cmd = "delete %s %s %s" % (o.type, table, o.name)

        ret = execute_cmd(cmd, filename, lineno)

        # Check if the object still exists after I deleted it.

        if ret != 0 or obj_exist(o, table, filename, lineno):

            reason = "Cannot remove the " + obj_handle

            print_error(reason, filename, lineno)

            return -1

    return 0

def obj_exist(o, table, filename, lineno):

    '''

    Check if the object exists.

    '''

    cmd = "list %s %s %s" % (o.type, table, o.name)

    ret = execute_cmd(cmd, filename, lineno)

    return True if (ret == 0) else False

def _obj_exist(o, filename, lineno):

    '''

    Check if the object exists.

    '''

    cmd = "list %s %s %s %s" % (o.type, o.family, o.table, o.name)

    ret = execute_cmd(cmd, filename, lineno)

    return True if (ret == 0) else False

def output_clean(pre_output, chain):

    pos_chain = pre_output.find(chain.name)

    if pos_chain == -1:

        return ""

    output_intermediate = pre_output[pos_chain:]

    brace_start = output_intermediate.find("{")

    brace_end = output_intermediate.find("}")

    pre_rule = output_intermediate[brace_start:brace_end]

    if pre_rule[1:].find("{") > -1:  # this rule has a set.

        set = pre_rule[1:].replace("\t", "").replace("\n", "").strip()

        set = set.split(";")[2].strip() + "}"

        remainder = output_clean(chain.name + " {;;" + output_intermediate[brace_end+1:], chain)

        if len(remainder) <= 0:

            return set

        return set + " " + remainder

    else:

        rule = pre_rule.split(";")[2].replace("\t", "").replace("\n", "").\

            strip()

    if len(rule) < 0:

        return ""

    return rule

def payload_check_elems_to_set(elems):

    newset = set()

    for n, line in enumerate(elems.split('[end]')):

        e = line.strip()

        if e in newset:

            print_error("duplicate", e, n)

            return newset

        newset.add(e)

    return newset

def payload_check_set_elems(want, got):

    if want.find('element') < 0 or want.find('[end]') < 0:

        return 0

    if got.find('element') < 0 or got.find('[end]') < 0:

        return 0

    set_want = payload_check_elems_to_set(want)

    set_got = payload_check_elems_to_set(got)

    return set_want == set_got

def payload_check(payload_buffer, file, cmd):

    file.seek(0, 0)

    i = 0

    if not payload_buffer:

        return False

    for lineno, want_line in enumerate(payload_buffer):

        line = file.readline()

        if want_line == line:

            i += 1

            continue

        if want_line.find('[') < 0 and line.find('[') < 0:

            continue

        if want_line.find(']') < 0 and line.find(']') < 0:

            continue

        if payload_check_set_elems(want_line, line):

            continue

        print_differences_warning(file.name, lineno, want_line.strip(),

                                  line.strip(), cmd)

        return 0

    return i > 0

def json_dump_normalize(json_string, human_readable = False):

    json_obj = json.loads(json_string)

    if human_readable:

        return json.dumps(json_obj, sort_keys = True,

                          indent = 4, separators = (',', ': '))

    else:

        return json.dumps(json_obj, sort_keys = True)

def json_validate(json_string):

    json_obj = json.loads(json_string)

    try:

        nftables.json_validate(json_obj)

    except Exception:

        print_error("schema validation failed for input '%s'" % json_string)

        print_error(traceback.format_exc())

def rule_add(rule, filename, lineno, force_all_family_option, filename_path):

    '''

    Adds a rule

    '''

    # TODO Check if a rule is added correctly.

    ret = warning = error = unit_tests = 0

    if not table_list or not chain_list:

        reason = "Missing table or chain to add rule."

        print_error(reason, filename, lineno)

        return [-1, warning, error, unit_tests]

    if rule[1].strip() == "ok":

        payload_expected = None

        try:

            payload_log = open("%s.payload" % filename_path)

            payload_expected = payload_find_expected(payload_log, rule[0])

        except:

            payload_log = None

        if enable_json_option:

            try:

                json_log = open("%s.json" % filename_path)

                json_input = json_find_expected(json_log, rule[0])

            except:

                json_input = None

            if not json_input:

                print_error("did not find JSON equivalent for rule '%s'"

                            % rule[0])

            else:

                try:

                    json_input = json_dump_normalize(json_input)

                except ValueError:

                    reason = "Invalid JSON syntax in rule: %s" % json_input

                    print_error(reason)

                    return [-1, warning, error, unit_tests]

            try:

                json_log = open("%s.json.output" % filename_path)

                json_expected = json_find_expected(json_log, rule[0])

            except:

                # will use json_input for comparison

                json_expected = None

            if json_expected:

                try:

                    json_expected = json_dump_normalize(json_expected)

                except ValueError:

                    reason = "Invalid JSON syntax in expected output: %s" % json_expected

                    print_error(reason)

                    return [-1, warning, error, unit_tests]

    for table in table_list:

        if rule[1].strip() == "ok":

            table_payload_expected = None

            try:

                payload_log = open("%s.payload.%s" % (filename_path, table.family))

                table_payload_expected = payload_find_expected(payload_log, rule[0])

            except:

                if not payload_log:

                    print_error("did not find any payload information",

                                filename_path)

                elif not payload_expected:

                    print_error("did not find payload information for "

                                "rule '%s'" % rule[0], payload_log.name, 1)

            if not table_payload_expected:

                table_payload_expected = payload_expected

        for table_chain in table.chains:

            chain = chain_get_by_name(table_chain)

            unit_tests += 1

            table_flush(table, filename, lineno)

            payload_log = tempfile.TemporaryFile(mode="w+")

            # Add rule and check return code

            cmd = "add rule %s %s %s" % (table, chain, rule[0])

            ret = execute_cmd(cmd, filename, lineno, payload_log, debug="netlink")

            state = rule[1].rstrip()

            if (ret in [0,134] and state == "fail") or (ret != 0 and state == "ok"):

                if state == "fail":

                    test_state = "This rule should have failed."

                else:

                    test_state = "This rule should not have failed."

                reason = cmd + ": " + test_state

                print_error(reason, filename, lineno)

                ret = -1

                error += 1

                if not force_all_family_option:

                    return [ret, warning, error, unit_tests]

            if state == "fail" and ret != 0:

                ret = 0

                continue

            if ret != 0:

                continue

            # Check for matching payload

            if state == "ok" and not payload_check(table_payload_expected,

                                                   payload_log, cmd):

                error += 1

                gotf = open("%s.payload.got" % filename_path, 'a')

                payload_log.seek(0, 0)

                gotf.write("# %s\n" % rule[0])

                while True:

                    line = payload_log.readline()

                    if line == "":

                        break

                    gotf.write(line)

                gotf.close()

                print_warning("Wrote payload for rule %s" % rule[0],

                              gotf.name, 1)

            # Check for matching ruleset listing

            numeric_proto_old = nftables.set_numeric_proto_output(True)

            stateless_old = nftables.set_stateless_output(True)

            list_cmd = 'list table %s' % table

            rc, pre_output, err = nftables.cmd(list_cmd)

            nftables.set_numeric_proto_output(numeric_proto_old)

            nftables.set_stateless_output(stateless_old)

            output = pre_output.split(";")

            if len(output) < 2:

                reason = cmd + ": Listing is broken."

                print_error(reason, filename, lineno)

                ret = -1

                error += 1

                if not force_all_family_option:

                    return [ret, warning, error, unit_tests]

                continue

            rule_output = output_clean(pre_output, chain)

            retest_output = False

            if len(rule) == 3:

                teoric_exit = rule[2]

                retest_output = True

            else:

                teoric_exit = rule[0]

            if rule_output.rstrip() != teoric_exit.rstrip():

                if rule[0].find("{") != -1:  # anonymous sets

                    if not set_check_element(teoric_exit.rstrip(),

                                         rule_output.rstrip()):

                        warning += 1

                        retest_output = True

                        print_differences_warning(filename, lineno,

                                                  teoric_exit.rstrip(),

                                                  rule_output, cmd)

                        if not force_all_family_option:

                            return [ret, warning, error, unit_tests]

                else:

                    if len(rule_output) <= 0:

                        error += 1

                        print_differences_error(filename, lineno, cmd)

                        if not force_all_family_option:

                            return [ret, warning, error, unit_tests]

                    warning += 1

                    retest_output = True

                    print_differences_warning(filename, lineno,

                                              teoric_exit.rstrip(),

                                              rule_output, cmd)

                    if not force_all_family_option:

                        return [ret, warning, error, unit_tests]

            if retest_output:

                table_flush(table, filename, lineno)

                # Add rule and check return code

                cmd = "add rule %s %s %s" % (table, chain, rule_output.rstrip())

                ret = execute_cmd(cmd, filename, lineno, payload_log, debug="netlink")

                if ret != 0:

                    test_state = "Replaying rule failed."

                    reason = cmd + ": " + test_state

                    print_warning(reason, filename, lineno)