'''
OFTests for functionality needed from the OLT.
'''
import logging
from __builtin__ import xrange
from oftest import config
import oftest.base_tests as base_tests
import oftest.packet as scapy

import ofp
import time
import copy

from oftest.testutils import *

onu_port = test_param_get("onu_port", 130)
onu_port2 = test_param_get("onu_port2", 130)
olt_port = test_param_get("olt_port", 258)


def double_vlan_udp_packet(pktlen=100,
                           eth_dst='00:01:02:03:04:05',
                           eth_src='00:06:07:08:09:0a',
                           dl_vlan_enable=False,
                           c_vlan_vid=0,
                           c_vlan_pcp=0,
                           s_vlan_vid=0,
                           s_vlan_pcp=0,
                           ip_src='192.168.0.1',
                           ip_dst='192.168.0.2',
                           ip_tos=0,
                           ip_ttl=64,
                           udp_sport=1234,
                           udp_dport=80,
                           ip_ihl=None,
                           ip_options=False,
                           eth_type=0x8100
                           ):
    """
    Return a double vlan tagged dataplane UDP packet
    Supports a few parameters:
    @param len Length of packet in bytes w/o CRC
    @param eth_dst Destination MAC
    @param eth_src Source MAC
    @param dl_vlan_enable True if the packet is with vlan, False otherwise
    @param c_vlan_vid CVLAN ID
    @param c_vlan_pcp CVLAN priority
    @param s_vlan_vid SVLAN ID
    @param s_vlan_pcp SVLAN priority
    @param ip_src IP source
    @param ip_dst IP destination
    @param ip_tos IP ToS
    @param ip_ttl IP TTL
    @param udp_dport UDP destination port
    @param udp_sport UDP source port

    Generates a simple UDP packet. Users shouldn't assume anything about
    this packet other than that it is a valid ethernet/IP/UDP frame.
    """

    if MINSIZE > pktlen:
        pktlen = MINSIZE

    # Note Dot1Q.id is really CFI
    if (dl_vlan_enable):
        pkt = scapy.Ether(dst=eth_dst, src=eth_src, type=eth_type) / \
              scapy.Dot1Q(prio=s_vlan_pcp, vlan=s_vlan_vid) / \
              scapy.Dot1Q(prio=c_vlan_pcp, vlan=c_vlan_vid) / \
              scapy.IP(src=ip_src, dst=ip_dst, tos=ip_tos, ttl=ip_ttl, ihl=ip_ihl) / \
              scapy.UDP(sport=udp_sport, dport=udp_dport)
    else:
        if not ip_options:
            pkt = scapy.Ether(dst=eth_dst, src=eth_src) / \
                  scapy.IP(src=ip_src, dst=ip_dst, tos=ip_tos, ttl=ip_ttl, ihl=ip_ihl) / \
                  scapy.UDP(sport=udp_sport, dport=udp_dport)

        else:
            pkt = scapy.Ether(dst=eth_dst, src=eth_src) / \
                  scapy.IP(src=ip_src, dst=ip_dst, tos=ip_tos, ttl=ip_ttl, ihl=ip_ihl, options=ip_options) / \
                  scapy.UDP(sport=udp_sport, dport=udp_dport)

    pkt = pkt / ("D" * (pktlen - len(pkt)))

    return pkt


def testPacketIn(self, match, parsed_pkt):
    delete_all_flows(self.controller)

    pkt = str(parsed_pkt)

    for of_port in config["port_map"]:
        m = copy.deepcopy(match)
        m.oxm_list.append(ofp.oxm.in_port(of_port))
        request = ofp.message.flow_add(
            table_id=test_param_get("table", 0),
            cookie=42,
            match=m,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[
                        ofp.action.output(
                            port=ofp.OFPP_CONTROLLER,
                            max_len=ofp.OFPCML_NO_BUFFER)])],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)
        logging.info("Inserting flow sending matching packets to controller")
        self.controller.message_send(request)
        do_barrier(self.controller)

    for of_port in config["port_map"]:
        logging.info("PacketInExact test, port %d", of_port)
        self.dataplane.send(of_port, pkt)
        verify_packet_in(self, pkt, of_port, ofp.OFPR_ACTION)
        verify_packets(self, pkt, [])


def buildIgmp(payload):
    ether = scapy.Ether(src="00:01:02:03:04:05")
    ip = scapy.IP(src="1.2.3.4")
    payload.igmpize(ip, ether)
    pkt = ether / ip / payload
    if len(pkt) < 60:
        pad_len = 60 - len(pkt)
        pad = scapy.PAD()
        pad.load = '\x00' * pad_len
        pkt = pkt / pad
    return pkt


class EapolPacketIn(base_tests.SimpleDataPlane):
    """Verify packet-ins are sent for EAPOL packets """

    def runTest(self):
        logging.info("Running EAPOL Packet In test")

        match = ofp.match()
        match.oxm_list.append(ofp.oxm.eth_type(0x888e))
        # Use ethertype 0x888e and multicast destination MAC address
        pkt = simple_eth_packet(pktlen=60, eth_dst='01:00:5E:7F:FF:FF', eth_type=0x888e)

        testPacketIn(self, match, pkt)


class IGMPPacketIn(base_tests.SimpleDataPlane):
    """Verify packet-ins are sent for IGMP packets """

    def runTest(self):
        logging.info("Running IGMP Packet In test")

        match = ofp.match()
        match.oxm_list.append(ofp.oxm.eth_type(0x800))
        match.oxm_list.append(ofp.oxm.ip_proto(2))

        pkt = scapy.Ether(dst='01:00:5E:7F:FF:FF', src='00:00:00:00:00:01') / \
              scapy.IP(src='10.0.0.1', dst='10.0.0.2', ttl=60, tos=0, id=0, proto=2)

        pkt = pkt / ("0" * (100 - len(pkt)))

        testPacketIn(self, match, pkt)


class IGMPQueryPacketOut(base_tests.SimpleDataPlane):
    """Verify sending multicast membership queries down to onu_ports"""

    def runTest(self):
        logging.info("Running IGMP query packet out")

        igmp = scapy.IGMP(type=0x11, gaddr="224.0.0.1")
        pkt = buildIgmp(igmp)

        msg = ofp.message.packet_out()
        msg.in_port = ofp.OFPP_CONTROLLER
        msg.buffer_id = 0xffffffff
        msg.data = str(pkt)
        msg.actions = [ofp.action.output(
            port=onu_port,
            max_len=ofp.OFPCML_NO_BUFFER)]
        time.sleep(1)

        self.controller.message_send(msg)

        verify_no_errors(self.controller)

        verify_packet(self, pkt, onu_port)


class TestMeter(base_tests.SimpleDataPlane):
    def runTest(self):
        logging.info("Running Meter tests")
        dropMeterBand = ofp.meter_band.drop(rate=640)
        meter_mod = ofp.message.meter_mod(xid=1, command=ofp.OFPMC_ADD, meter_id=1, meters=[dropMeterBand])
        self.controller.message_send(meter_mod)

        time.sleep(1)

        verify_no_errors(self.controller)

        vlan_id = 201
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(onu_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | vlan_id))

        request = ofp.message.flow_add(
            table_id=test_param_get("table", 0),
            cookie=42,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[ofp.action.output(port=olt_port)]
                ),
                ofp.instruction.meter(meter_id=1)
            ],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)

        self.controller.message_send(request)
        time.sleep(1)
        verify_no_errors(self.controller)

        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(olt_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | vlan_id))

        request = ofp.message.flow_add(
            table_id=test_param_get("table", 0),
            cookie=43,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[ofp.action.output(port=onu_port)]
                ),
                ofp.instruction.meter(meter_id=1)
            ],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)

        self.controller.message_send(request)
        time.sleep(1)
        verify_no_errors(self.controller)
        do_barrier(self.controller)
        time.sleep(5)

        inPkt = simple_udp_packet(dl_vlan_enable=True, vlan_vid=vlan_id, vlan_pcp=0)
        # downstream
        # self.dataplane.send(olt_port, str(inPkt))
        # verify_packet(self, inPkt, onu_port)
        # upstream
        # for i in range(1,400):
        #	self.dataplane.send(onu_port, str(inPkt))
        # verify_packet(self, inPkt, olt_port)

        # clean up the test
        meter_mod = ofp.message.meter_mod(xid=2, command=ofp.OFPMC_DELETE, meter_id=1)
        self.controller.message_send(meter_mod)
        time.sleep(1)
        delete_all_flows(self.controller)
        verify_no_errors(self.controller)


class TestDuplicateMeter(base_tests.SimpleDataPlane):
    def runTest(self):
        logging.info("Running Duplicate Meter Test")
        dropMeterBand = ofp.meter_band.drop(rate=500)
        meter_mod = ofp.message.meter_mod(xid=1, command=ofp.OFPMC_ADD, meter_id=1, meters=[dropMeterBand])
        self.controller.message_send(meter_mod)
        self.controller.message_send(meter_mod)

        time.sleep(1)

        try:
            verify_no_errors(self.controller)
        except AssertionError as e:
            if (not e.message == "unexpected error type=12 code=1"):
                raise AssertionError("Incorrect error type: %s" % e.message)


class VlanTest(base_tests.SimpleDataPlane):
    """Verify the switch can push/pop a VLAN tag and forward out a port """

    def runTest(self):
        logging.info("Running push VLAN test")

        vlan_id = 200

        delete_all_flows(self.controller)

        # PUSH
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(onu_port))
        match.oxm_list.append(ofp.oxm.vlan_vid_masked(value=ofp.OFPVID_PRESENT, value_mask=ofp.OFPVID_PRESENT))
        match.oxm_list.append(ofp.oxm.vlan_pcp(value=0))

        request = ofp.message.flow_add(
            table_id=test_param_get("table", 0),
            cookie=42,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[
                        ofp.action.push_vlan(ethertype=0x8100),
                        ofp.action.set_field(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | vlan_id)),
                        ofp.action.set_field(ofp.oxm.vlan_pcp(0)),
                        ofp.action.output(port=olt_port)])],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)

        logging.info("Inserting flow tagging upstream")
        self.controller.message_send(request)

        # POP
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(olt_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | vlan_id))

        request = ofp.message.flow_add(
            table_id=test_param_get("table", 0),
            cookie=43,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[
                        ofp.action.pop_vlan(),
                        ofp.action.output(port=onu_port)])],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)

        logging.info("Inserting flow tagging downstream")
        self.controller.message_send(request)
        do_barrier(self.controller)
        time.sleep(5)

        inPkt = simple_udp_packet(dl_vlan_enable=True, vlan_vid=0, vlan_pcp=0)
        outPkt = simple_udp_packet(pktlen=100, dl_vlan_enable=True,
                                   vlan_vid=vlan_id, vlan_pcp=0, dl_vlan_cfi=0)

        # Send untagged packet in the ONU port and expect tagged packet out the OLT port
        self.dataplane.send(onu_port, str(inPkt))
        verify_packet(self, outPkt, olt_port)

        # Send untagged packet in the OLT port and expect no packets to come out
        self.dataplane.send(olt_port, str(inPkt))
        verify_packets(self, outPkt, [])

        inPkt = simple_udp_packet(pktlen=104, dl_vlan_enable=True,
                                  vlan_vid=vlan_id, vlan_pcp=0, dl_vlan_cfi=0)
        outPkt = simple_udp_packet(pktlen=104, dl_vlan_enable=True, vlan_vid=0, vlan_pcp=0)

        # Send tagged packet in the OLT port and expect untagged packet out the OLT port
        self.dataplane.send(olt_port, str(inPkt))
        verify_packet(self, outPkt, onu_port)

        # Send tagged packet in the ONU port and expect no packets to come out
        self.dataplane.send(onu_port, str(inPkt))
        verify_packets(self, outPkt, [])


def createAllGroupAdd(group_id, ports=[]):
    buckets = []

    for portNum in ports:
        buckets.append(ofp.common.bucket(watch_port=ofp.OFPP_ANY, watch_group=ofp.OFPG_ANY,
                                         actions=[ofp.action.pop_vlan(), ofp.action.output(port=portNum)]))

    group_add = ofp.message.group_add(group_type=ofp.OFPGT_ALL, group_id=group_id, buckets=buckets)

    return group_add


def createAllGroupMod(group_id, ports=[]):
    buckets = []

    for portNum in ports:
        buckets.append(ofp.common.bucket(watch_port=ofp.OFPP_ANY, watch_group=ofp.OFPG_ANY,
                                         actions=[ofp.action.output(port=portNum)]))

    group_mod = ofp.message.group_mod(command=ofp.OFPGC_MODIFY, group_type=ofp.OFPGT_ALL, group_id=group_id,
                                      buckets=buckets)

    return group_mod


class TestGroupAdd(base_tests.SimpleDataPlane):
    def runTest(self):
        logging.info("Running Group tests")
        delete_all_flows(self.controller)
        delete_all_groups(self.controller)

        test_group_id = 1

        # output to two ONU
        group_add = createAllGroupAdd(test_group_id, ports=[onu_port, onu_port2])

        self.controller.message_send(group_add)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # Remove the group and then readd it.
        group_delete = ofp.message.group_delete(group_id=test_group_id)
        self.controller.message_send(group_delete)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        group_add = createAllGroupAdd(test_group_id, [onu_port, onu_port2])
        self.controller.message_send(group_add)
        do_barrier(self.controller)
        verify_no_errors(self.controller)


        # clean up the test
        group_delete = ofp.message.group_delete(group_id=test_group_id)
        self.controller.message_send(group_delete)

        do_barrier(self.controller)
        verify_no_errors(self.controller)


class TestGroupMod(base_tests.SimpleDataPlane):
    def runTest(self):
        logging.info("Running Group tests")
        delete_all_flows(self.controller)
        delete_all_groups(self.controller)

        test_group_id = 1

        group_add = createAllGroupAdd(test_group_id, [onu_port, onu_port2])

        self.controller.message_send(group_add)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # Modifying the group
        group_mod = createAllGroupMod(test_group_id, [onu_port2])

        self.controller.message_send(group_mod)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # Add a bucket into the group
        group_mod = createAllGroupMod(test_group_id, [onu_port, onu_port2])
        self.controller.message_send(group_mod)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # Modifying a non-existing group
        group_mod = createAllGroupMod(777, [onu_port2])

        self.controller.message_send(group_mod)
        do_barrier(self.controller)
        errorExperienced = 0
        try:
            verify_no_errors(self.controller)
        except AssertionError as e:
            errorExperienced = 1
            if (not (e.message == "unexpected error type=6 code=8")):
                raise AssertionError("Incorrect error type: %s" % e.message)
        if not errorExperienced:
            raise AssertionError("An error message is expected, but not shown.")


        # clean up the test
        group_delete = ofp.message.group_delete(xid=2, group_id=test_group_id)
        self.controller.message_send(group_delete)

        do_barrier(self.controller)
        verify_no_errors(self.controller)


class TestDuplicateGroup(base_tests.SimpleDataPlane):
    def runTest(self):
        logging.info("Running Group tests")
        delete_all_flows(self.controller)
        delete_all_groups(self.controller)

        test_group_id = 1
        group_add = createAllGroupAdd(test_group_id, ports=[onu_port])

        self.controller.message_send(group_add)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # Add the same group id
        duplicate_group_fail = 0

        self.controller.message_send(group_add)
        do_barrier(self.controller)
        try:
            verify_no_errors(self.controller)
        except AssertionError as e:
            duplicate_group_fail = 1
            if (not e.message == "unexpected error type=6 code=0"):
                raise AssertionError("Incorrect error type: %s" % e.message)
        if not duplicate_group_fail:
            raise AssertionError("Adding duplicate groups didn't raise an error.")

        # clean up the test
        group_delete = ofp.message.group_delete(xid=2, group_id=test_group_id)
        self.controller.message_send(group_delete)

        do_barrier(self.controller)
        verify_no_errors(self.controller)


class TestGroupAndFlow(base_tests.SimpleDataPlane):
    def runTest(self):
        logging.info("Running Group tests")
        delete_all_flows(self.controller)
        delete_all_groups(self.controller)

        # Create a group
        test_group_id = 1
        group_add = createAllGroupAdd(test_group_id, ports=[onu_port])

        self.controller.message_send(group_add)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # Create a flow rule matching olt port and vlan id
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(olt_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | 201))

        flow_pointing_to_group = ofp.message.flow_add(
            table_id=test_param_get("table", 0),
            cookie=43,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[ofp.action.group(group_id=test_group_id)])],
            buffer_id=ofp.OFP_NO_BUFFER, priority=1000)

        self.controller.message_send(flow_pointing_to_group)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # After letting a flow rule point to the group, test we can do group_mod
        group_mod = createAllGroupMod(test_group_id, ports=[onu_port2])
        self.controller.message_send(group_mod)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # Test we can remove flows and then remove group
        flow_delete = ofp.message.flow_delete(table_id=test_param_get("table", 0))
        self.controller.message_send(flow_delete)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        group_delete = ofp.message.group_delete(xid=3, group_id=test_group_id)
        self.controller.message_send(group_delete)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # Add the group and flow back, test it we can first remove group and then remove the flow. 
        '''group_add = createAllGroupAdd(test_group_id, ports=[onu_port])
        self.controller.message_send(group_add)

        self.controller.message_send(flow_pointing_to_group)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        group_delete = ofp.message.group_delete(xid = 4, group_id = test_group_id)
        self.controller.message_send(group_delete)'''

        self.controller.message_send(flow_delete)
        do_barrier(self.controller)
        verify_no_errors(self.controller)


class TestGroupForwarding(base_tests.SimpleDataPlane):
    def runTest(self):
        logging.info("Running Group datapath forwarding tests")
        delete_all_flows(self.controller)
        delete_all_groups(self.controller)

        vlan_id = 201

        # Create a group
        test_group_id = 1
        group_add = createAllGroupAdd(test_group_id, [onu_port])

        self.controller.message_send(group_add)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # Create a flow rule matching olt port and vlan id
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(olt_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | vlan_id))

        request = ofp.message.flow_add(
            table_id=test_param_get("table", 0),
            cookie=43,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[ofp.action.group(group_id=test_group_id)]),
            ],
            buffer_id=ofp.OFP_NO_BUFFER, priority=1000)

        self.controller.message_send(request)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # It takes some time for flows to propagate down to the data plane
        time.sleep(10)

        inPkt = simple_udp_packet(pktlen=104, dl_vlan_enable=True,
                                  vlan_vid=vlan_id, vlan_pcp=0, dl_vlan_cfi=0, eth_dst="01:01:11:12:11:12")
        outPkt = inPkt
        self.dataplane.send(olt_port, str(inPkt))
        verify_packet(self, outPkt, onu_port)


        # Now put 2 ONU ports in the group and test that the input packet is 
        # duplicated out both ports
        group_mod = createAllGroupMod(test_group_id, ports=[onu_port, onu_port2])
        self.controller.message_send(group_mod)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # It takes some time for flows to propagate down to the data plane
        time.sleep(10)

        self.dataplane.send(olt_port, str(inPkt))
        verify_packet(self, outPkt, onu_port2)
        verify_packet(self, outPkt, onu_port)
        # verify_packets(self, outPkt, [onu_port,onu_port2])

        # clean up the test
        request = ofp.message.flow_delete(table_id=test_param_get("table", 0))
        self.controller.message_send(request)
        group_delete = ofp.message.group_delete(xid=2, group_id=test_group_id)
        self.controller.message_send(group_delete)

        do_barrier(self.controller)
        verify_no_errors(self.controller)


class TestGroupModForwarding(base_tests.SimpleDataPlane):
    def runTest(self):
        logging.info("Running datapath forwarding tests for group mod")
        delete_all_flows(self.controller)
        delete_all_groups(self.controller)

        vlan_id = 201

        # Create a group
        test_group_id = 1
        group_add = createAllGroupAdd(test_group_id, [onu_port, onu_port2])

        self.controller.message_send(group_add)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # Create a flow rule matching olt port and vlan id
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(olt_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | vlan_id))

        request = ofp.message.flow_add(
            table_id=test_param_get("table", 0),
            cookie=43,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[ofp.action.group(group_id=test_group_id)]),
            ],
            buffer_id=ofp.OFP_NO_BUFFER, priority=1000)

        self.controller.message_send(request)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # It takes some time for flows to propagate down to the data plane
        time.sleep(10)

        inPkt = simple_udp_packet(pktlen=104, dl_vlan_enable=True,
                                  vlan_vid=vlan_id, vlan_pcp=0, dl_vlan_cfi=0, eth_dst="01:01:11:12:11:12")
        outPkt = inPkt
        self.dataplane.send(olt_port, str(inPkt))
        verify_packet(self, outPkt, onu_port)
        verify_packet(self, outPkt, onu_port2)

        # Now remove onu port 1 from the group and test that the input packet is no longer forwarded to onu port 1
        group_mod = createAllGroupMod(test_group_id, ports=[onu_port2])
        self.controller.message_send(group_mod)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        time.sleep(10)
        self.dataplane.send(olt_port, str(inPkt))
        verify_no_packet(self, outPkt, onu_port)
        verify_packet(self, outPkt, onu_port2)

        # Now remove all ports in the group and test that the input packet is no longer forwarded to any port 
        group_mod = createAllGroupMod(test_group_id, ports=[])
        self.controller.message_send(group_mod)
        do_barrier(self.controller)
        verify_no_errors(self.controller)
        time.sleep(10)
        self.dataplane.send(olt_port, str(inPkt))
        verify_packets(self, outPkt, [])


class TransparentVlanTest(base_tests.SimpleDataPlane):
    def runTest(self):
        logging.info("Running transparent vlan tests")
        delete_all_flows(self.controller)

        vlan_id = 201
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(onu_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | vlan_id))

        request = ofp.message.flow_add(
            table_id=test_param_get("table", 0),
            cookie=42,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[
                        ofp.action.output(port=olt_port)]),
            ],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)

        self.controller.message_send(request)

        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(olt_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | vlan_id))

        request = ofp.message.flow_add(
            table_id=test_param_get("table", 0),
            cookie=43,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[
                        ofp.action.output(port=onu_port)]),
            ],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)

        self.controller.message_send(request)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # It takes some time for flows to propagate down to the data plane
        time.sleep(2)

        inPkt = simple_udp_packet(dl_vlan_enable=True, vlan_vid=vlan_id, vlan_pcp=0)
        # upstream
        self.dataplane.send(onu_port, str(inPkt))
        verify_packet(self, inPkt, olt_port)
        # downstream
        self.dataplane.send(olt_port, str(inPkt))
        verify_packet(self, inPkt, onu_port)

        # clean up the test
        delete_all_flows(self.controller)
        do_barrier(self.controller)
        verify_no_errors(self.controller)


class DoubleVlanTest(base_tests.SimpleDataPlane):
    def runTest(self):
        logging.info("Running double vlan tests")
        delete_all_flows(self.controller)

        c_vlan_id = 100
        s_vlan_id = 102
        # upstream flow rule
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(onu_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(value=ofp.OFPVID_PRESENT))
        match.oxm_list.append(ofp.oxm.vlan_pcp(value=0))

        # push inner vlan (c-vlan) for upstream
        request = ofp.message.flow_add(
            table_id=test_param_get("table", 0),
            cookie=42,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[
                        ofp.action.set_field(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | c_vlan_id))]),
                ofp.instruction.goto_table(1)],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)

        self.controller.message_send(request)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # push outer vlan (s-vlan) for upstream
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(onu_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | c_vlan_id))
        match.oxm_list.append(ofp.oxm.vlan_pcp(0))

        request = ofp.message.flow_add(
            table_id=test_param_get("table", 1),
            cookie=43,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[
                        ofp.action.push_vlan(ethertype=0x8100),
                        ofp.action.set_field(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | s_vlan_id)),
                        ofp.action.set_field(ofp.oxm.vlan_pcp(0)),
                        ofp.action.output(port=olt_port)]),
            ],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)

        self.controller.message_send(request)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # strip outer vlan (s-vlan) for downstream
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(olt_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | s_vlan_id))
        match.oxm_list.append(ofp.oxm.vlan_pcp(0))
        request = ofp.message.flow_add(
            table_id=test_param_get("table", 0),
            cookie=44,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[ofp.action.pop_vlan()]),
                ofp.instruction.goto_table(1)],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)

        self.controller.message_send(request)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # rewrite inner vlan (c-vlan) to default (0) for downstream
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(olt_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | c_vlan_id))
        match.oxm_list.append(ofp.oxm.vlan_pcp(0))

        request = ofp.message.flow_add(
            table_id=test_param_get("table", 1),
            cookie=45,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[
                        ofp.action.pop_vlan(),
                        ofp.action.output(port=onu_port)])
            ],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)

        self.controller.message_send(request)
        do_barrier(self.controller)
        verify_no_errors(self.controller)
        # It takes some time for flows to propagate down to the data plane

        time.sleep(10)
        incorrectTagPkt = simple_udp_packet(pktlen=100, dl_vlan_enable=True, vlan_vid=100, vlan_pcp=0)
        untaggedPkt = simple_udp_packet(pktlen=100, dl_vlan_enable=True, vlan_vid=0, vlan_pcp=0)

        '''
		FIXME: hack because OLT does _not_ tag packets correctly
			it ignores the eth_type in the push_vlan action
			and always slaps on 0x8100 even though the 
			down stream flow must be 0x88a8.
	    '''
        doubleTaggedPkt = double_vlan_udp_packet(pktlen=104, dl_vlan_enable=True,
                                                 c_vlan_vid=c_vlan_id,
                                                 s_vlan_vid=s_vlan_id,
                                                 c_vlan_pcp=0, s_vlan_pcp=0, eth_type=0x88a8)

        upstreamDoubleTaggedPkt = double_vlan_udp_packet(pktlen=104, dl_vlan_enable=True,
                                                         c_vlan_vid=c_vlan_id,
                                                         s_vlan_vid=s_vlan_id,
                                                         c_vlan_pcp=0, s_vlan_pcp=0)

        # test upstream untagged packet got double tag at OLT
        logging.info(str(doubleTaggedPkt))

        self.dataplane.send(onu_port, str(untaggedPkt))
        verify_packet(self, upstreamDoubleTaggedPkt, olt_port)

        # test downstream doubletagged packet got untagged at ONU
        self.dataplane.send(olt_port, str(doubleTaggedPkt))
        verify_packet(self, untaggedPkt, onu_port)

        # test upstream doubletagged packet got dropped
        self.dataplane.send(onu_port, str(doubleTaggedPkt))
        verify_no_packet(self, upstreamDoubleTaggedPkt, olt_port)

        # test downstream untagged packet got dropped at ONU
        self.dataplane.send(olt_port, str(untaggedPkt))
        verify_no_packet(self, untaggedPkt, onu_port)

        # test upstream icorrectly tagged packet; should get dropped
        self.dataplane.send(onu_port, str(incorrectTagPkt))

        verify_no_packet(self, upstreamDoubleTaggedPkt, olt_port)

        time.sleep(2)

        # clean up the test
        delete_all_flows(self.controller)
        do_barrier(self.controller)
        verify_no_errors(self.controller)


class TestCyclingDoubleVlan(base_tests.SimpleDataPlane):
    """Cycle through vlans and test traffic flow"""

    def runTest(self):
        logging.info("Running cycling vlan test")
        delete_all_flows(self.controller)
        self.cookie = 42

        self.vlan_pairs = []

        for stag in xrange(2, 4000, 100):
            for ctag in xrange(2, 4000, 100):
                self.installDoubleTaggingRule(stag, ctag)
                self.vlan_pairs.append((stag, ctag))
                time.sleep(1)
                self.testPacketFlow()

    def testPacketFlow(self):
        for (s_vlan_id, c_vlan_id) in self.vlan_pairs:
            incorrectTagPkt = simple_udp_packet(pktlen=100, dl_vlan_enable=True, vlan_vid=100, vlan_pcp=0)
            untaggedPkt = simple_udp_packet(pktlen=100, dl_vlan_enable=True, vlan_vid=0, vlan_pcp=0)

            upstreamDoubleTaggedPkt = double_vlan_udp_packet(pktlen=104, dl_vlan_enable=True,
                                                             c_vlan_vid=c_vlan_id,
                                                             s_vlan_vid=s_vlan_id,
                                                             c_vlan_pcp=0, s_vlan_pcp=0)

            logging.info("Testing s-tag %d, c-tag %d" % (s_vlan_id, c_vlan_id))

            self.dataplane.send(onu_port, str(untaggedPkt))
            verify_packet(self, upstreamDoubleTaggedPkt, olt_port)

            # test downstream doubletagged packet got untagged at ONU
            self.dataplane.send(olt_port, str(upstreamDoubleTaggedPkt))
            verify_packet(self, untaggedPkt, onu_port)

            # test upstream doubletagged packet got dropped
            self.dataplane.send(onu_port, str(upstreamDoubleTaggedPkt))
            verify_no_packet(self, upstreamDoubleTaggedPkt, olt_port)

            # test downstream untagged packet got dropped at ONU
            self.dataplane.send(olt_port, str(untaggedPkt))
            verify_no_packet(self, untaggedPkt, onu_port)

            # test upstream icorrectly tagged packet; should get dropped
            self.dataplane.send(onu_port, str(incorrectTagPkt))
            verify_no_packet(self, upstreamDoubleTaggedPkt, olt_port)

    def installDoubleTaggingRule(self, s_vlan_id, c_vlan_id):
        # upstream flow rule
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(onu_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(value=ofp.OFPVID_PRESENT))
        match.oxm_list.append(ofp.oxm.vlan_pcp(value=0))
        self.cookie += 1

        # push inner vlan (c-vlan) for upstream
        request = ofp.message.flow_add(
            table_id=test_param_get("table", 0),
            cookie=self.cookie,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[
                        ofp.action.set_field(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | c_vlan_id))]),
                ofp.instruction.goto_table(1)],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)

        self.controller.message_send(request)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # push outer vlan (s-vlan) for upstream
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(onu_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | c_vlan_id))
        match.oxm_list.append(ofp.oxm.vlan_pcp(0))
        self.cookie += 1

        request = ofp.message.flow_add(
            table_id=test_param_get("table", 1),
            cookie=self.cookie,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[
                        ofp.action.push_vlan(ethertype=0x8100),
                        ofp.action.set_field(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | s_vlan_id)),
                        ofp.action.set_field(ofp.oxm.vlan_pcp(0)),
                        ofp.action.output(port=olt_port)]),
            ],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)

        self.controller.message_send(request)
        do_barrier(self.controller)
        verify_no_errors(self.controller)
        self.cookie += 1

        # strip outer vlan (s-vlan) for downstream
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(olt_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | s_vlan_id))
        match.oxm_list.append(ofp.oxm.vlan_pcp(0))
        request = ofp.message.flow_add(
            table_id=test_param_get("table", 0),
            cookie=self.cookie,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[ofp.action.pop_vlan()]),
                ofp.instruction.goto_table(1)],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)

        self.controller.message_send(request)
        do_barrier(self.controller)
        verify_no_errors(self.controller)

        # rewrite inner vlan (c-vlan) to default (0) for downstream
        match = ofp.match()
        match.oxm_list.append(ofp.oxm.in_port(olt_port))
        match.oxm_list.append(ofp.oxm.vlan_vid(ofp.OFPVID_PRESENT | c_vlan_id))
        match.oxm_list.append(ofp.oxm.vlan_pcp(0))
        self.cookie += 1

        request = ofp.message.flow_add(
            table_id=test_param_get("table", 1),
            cookie=self.cookie,
            match=match,
            instructions=[
                ofp.instruction.apply_actions(
                    actions=[
                        ofp.action.pop_vlan(),
                        ofp.action.output(port=onu_port)])
            ],
            buffer_id=ofp.OFP_NO_BUFFER,
            priority=1000)

        self.controller.message_send(request)
        do_barrier(self.controller)
        verify_no_errors(self.controller)