# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright 2010 United States Government as represented by the
# Administrator of the National Aeronautics and Space Administration.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.
"""Unit Tests for network code."""

from nova.network import linux_net
from nova import test


class IptablesManagerTestCase(test.TestCase):
    sample_filter = ['#Generated by iptables-save on Fri Feb 18 15:17:05 2011',
                     '*filter',
                     ':INPUT ACCEPT [2223527:305688874]',
                     ':FORWARD ACCEPT [0:0]',
                     ':OUTPUT ACCEPT [2172501:140856656]',
                     ':nova-compute-FORWARD - [0:0]',
                     ':nova-compute-INPUT - [0:0]',
                     ':nova-compute-local - [0:0]',
                     ':nova-compute-OUTPUT - [0:0]',
                     ':nova-filter-top - [0:0]',
                     '-A FORWARD -j nova-filter-top ',
                     '-A OUTPUT -j nova-filter-top ',
                     '-A nova-filter-top -j nova-compute-local ',
                     '-A INPUT -j nova-compute-INPUT ',
                     '-A OUTPUT -j nova-compute-OUTPUT ',
                     '-A FORWARD -j nova-compute-FORWARD ',
                     '-A INPUT -i virbr0 -p udp -m udp --dport 53 -j ACCEPT ',
                     '-A INPUT -i virbr0 -p tcp -m tcp --dport 53 -j ACCEPT ',
                     '-A INPUT -i virbr0 -p udp -m udp --dport 67 -j ACCEPT ',
                     '-A INPUT -i virbr0 -p tcp -m tcp --dport 67 -j ACCEPT ',
                     '-A FORWARD -s 192.168.122.0/24 -i virbr0 -j ACCEPT ',
                     '-A FORWARD -i virbr0 -o virbr0 -j ACCEPT ',
                     '-A FORWARD -o virbr0 -j REJECT --reject-with '
                     'icmp-port-unreachable ',
                     '-A FORWARD -i virbr0 -j REJECT --reject-with '
                     'icmp-port-unreachable ',
                     'COMMIT',
                     '# Completed on Fri Feb 18 15:17:05 2011']

    sample_nat = ['# Generated by iptables-save on Fri Feb 18 15:17:05 2011',
                  '*nat',
                  ':PREROUTING ACCEPT [3936:762355]',
                  ':INPUT ACCEPT [2447:225266]',
                  ':OUTPUT ACCEPT [63491:4191863]',
                  ':POSTROUTING ACCEPT [63112:4108641]',
                  ':nova-compute-OUTPUT - [0:0]',
                  ':nova-compute-floating-ip-snat - [0:0]',
                  ':nova-compute-SNATTING - [0:0]',
                  ':nova-compute-PREROUTING - [0:0]',
                  ':nova-compute-POSTROUTING - [0:0]',
                  ':nova-postrouting-bottom - [0:0]',
                  '-A PREROUTING -j nova-compute-PREROUTING ',
                  '-A OUTPUT -j nova-compute-OUTPUT ',
                  '-A POSTROUTING -j nova-compute-POSTROUTING ',
                  '-A POSTROUTING -j nova-postrouting-bottom ',
                  '-A nova-postrouting-bottom -j nova-compute-SNATTING ',
                  '-A nova-compute-SNATTING -j nova-compute-floating-ip-snat ',
                  'COMMIT',
                  '# Completed on Fri Feb 18 15:17:05 2011']

    def setUp(self):
        super(IptablesManagerTestCase, self).setUp()
        self.manager = linux_net.IptablesManager()

    def test_filter_rules_are_wrapped(self):
        current_lines = self.sample_filter

        table = self.manager.ipv4['filter']
        table.add_rule('FORWARD', '-s 1.2.3.4/5 -j DROP')
        new_lines = self.manager._modify_rules(current_lines, table)
        self.assertTrue('-A runner.py-FORWARD '
                        '-s 1.2.3.4/5 -j DROP' in new_lines)

        table.remove_rule('FORWARD', '-s 1.2.3.4/5 -j DROP')
        new_lines = self.manager._modify_rules(current_lines, table)
        self.assertTrue('-A runner.py-FORWARD '
                        '-s 1.2.3.4/5 -j DROP' not in new_lines)

    def test_nat_rules(self):
        current_lines = self.sample_nat
        new_lines = self.manager._modify_rules(current_lines,
                                               self.manager.ipv4['nat'])

        for line in [':nova-compute-OUTPUT - [0:0]',
                     ':nova-compute-floating-ip-snat - [0:0]',
                     ':nova-compute-SNATTING - [0:0]',
                     ':nova-compute-PREROUTING - [0:0]',
                     ':nova-compute-POSTROUTING - [0:0]']:
            self.assertTrue(line in new_lines, "One of nova-compute's chains "
                                               "went missing.")

        seen_lines = set()
        for line in new_lines:
            line = line.strip()
            self.assertTrue(line not in seen_lines,
                            "Duplicate line: %s" % line)
            seen_lines.add(line)

        last_postrouting_line = ''

        for line in new_lines:
            if line.startswith('-A POSTROUTING'):
                last_postrouting_line = line

        self.assertTrue('-j nova-postrouting-bottom' in last_postrouting_line,
                        "Last POSTROUTING rule does not jump to "
                        "nova-postouting-bottom: %s" % last_postrouting_line)

        for chain in ['POSTROUTING', 'PREROUTING', 'OUTPUT']:
            self.assertTrue('-A %s -j runner.py-%s' %
                            (chain, chain) in new_lines,
                            "Built-in chain %s not wrapped" % (chain,))

    def test_filter_rules(self):
        current_lines = self.sample_filter
        new_lines = self.manager._modify_rules(current_lines,
                                               self.manager.ipv4['filter'])

        for line in [':nova-compute-FORWARD - [0:0]',
                     ':nova-compute-INPUT - [0:0]',
                     ':nova-compute-local - [0:0]',
                     ':nova-compute-OUTPUT - [0:0]']:
            self.assertTrue(line in new_lines, "One of nova-compute's chains"
                                               " went missing.")

        seen_lines = set()
        for line in new_lines:
            line = line.strip()
            self.assertTrue(line not in seen_lines,
                            "Duplicate line: %s" % line)
            seen_lines.add(line)

        for chain in ['FORWARD', 'OUTPUT']:
            for line in new_lines:
                if line.startswith('-A %s' % chain):
                    self.assertTrue('-j nova-filter-top' in line,
                                    "First %s rule does not "
                                    "jump to nova-filter-top" % chain)
                    break

        self.assertTrue('-A nova-filter-top '
                        '-j runner.py-local' in new_lines,
                        "nova-filter-top does not jump to wrapped local chain")

        for chain in ['INPUT', 'OUTPUT', 'FORWARD']:
            self.assertTrue('-A %s -j runner.py-%s' %
                            (chain, chain) in new_lines,
                            "Built-in chain %s not wrapped" % (chain,))
