#!/usr/bin/env python

from mininet.node import Host
from mininet.topo import Topo


class VLANHost(Host):
    """Host connected to VLAN interface"""

    def config(self, vlan=100, **params):
        """Configure VLANHost according to (optional) parameters
        vlan: VLAN ID for default interface"""

        hostid = super(Host, self).config(**params)

        intf = self.defaultIntf()
        # remove IP from default, "physical" interface
        self.cmd("ifconfig %s inet 0" % intf)
        # create VLAN interface
        self.cmd("vconfig add %s %d" % (intf, vlan))
        # assign the host's IP to the VLAN interface
        self.cmd("ifconfig %s.%d inet %s" % (intf, vlan, params["ip"]))
        # update the intf name and host's intf map
        new_name = "%s.%d" % (intf, vlan)
        # update the (Mininet) interface to refer to VLAN interface name
        intf.name = new_name
        # add VLAN interface to host's name to intf map
        self.nameToIntf[new_name] = intf
        return hostid


class VlanTopo(Topo):

    """Simple topology example."""

    def __init__(self):
        """Create custom topo."""

        # Initialize topology
        Topo.__init__(self)

        # Add hosts and switches
        host1 = self.addHost("h1", cls=VLANHost, vlan=200)
        host2 = self.addHost("h2", cls=VLANHost, vlan=300)
        host3 = self.addHost("h3", cls=VLANHost, vlan=200)
        host4 = self.addHost("h4", cls=VLANHost, vlan=300)
        host5 = self.addHost("h5", cls=VLANHost, vlan=200)
        host6 = self.addHost("h6", cls=VLANHost, vlan=300)

        s1 = self.addSwitch("s1")
        s2 = self.addSwitch("s2")
        s3 = self.addSwitch("s3")

        self.addLink(s1, s2)
        self.addLink(s2, host1)
        self.addLink(s2, host2)
        self.addLink(s2, host3)
        self.addLink(s1, s3)
        self.addLink(s3, host4)
        self.addLink(s3, host5)
        self.addLink(s3, host6)


topos = {"vlantopo": (lambda: VlanTopo())}