#!/usr/bin/python
import os.path
import sys
import pydhcplib
import pydhcplib.dhcp_network
from pydhcplib.dhcp_packet import *
from pydhcplib.type_hw_addr import hwmac
from pydhcplib.type_ipv4 import ipv4
from pydhcplib.type_strlist import strlist
import socket
import IN
from Queue import Queue
from threading import Thread
from subprocess import PIPE, Popen
import netifaces as ni
sys.path.append('/usr/lib/xen-default/lib/python/')
from xen.lowlevel import xs

import syslog as s

import time
from invirt import database
from invirt.config import structs as config

dhcp_options = {'domain_name_server': ','.join(config.dhcp.dns),
                'ip_address_lease_time': config.dhcp.get('leasetime', 60*60*24)}

class Interfaces(object):
    @staticmethod
    def primary_ip(name):
        """primary_ip returns an interface's primary IP address.

        This is the first IPv4 address returned by "ip addr show $name"
        """
        # TODO(quentin): The netifaces module is a pile of crappy C.
        # Figure out a way to do this in native Python.
        return ni.ifaddresses(name)[ni.AF_INET][0]['addr']

    @staticmethod
    def exists(name):
        """exists checks if an interface exists.

        Args:
        name: Interface name
        """
        return os.path.exists("/sys/class/net/"+name)


class DhcpBackend:
    def __init__(self, queue):
        database.connect()
        self.queue = queue
        self.main_ip = Interfaces.primary_ip(config.xen.iface)
    def add_route_and_arp(self, ip, intf, gateway):
        try:
            p = Popen(['ip', 'route', 'add', ip, 'dev', intf, 'src', self.main_ip, 'metric', '2' if intf.startswith('vif') else '1'], stdout=PIPE, stderr=PIPE)
            (out, err) = p.communicate()
            if p.returncode == 0:
                s.syslog(s.LOG_INFO, "Added route for IP %s to interface %s" % (ip, intf))
                self.queue.put((ip, gateway))
            sys.stderr.write(err)
            sys.stdout.write(out)
        except Exception as e:
            s.syslog(s.LOG_ERR, "Could not add route for IP %s: %s" % (ip, e))
    def findNIC(self, mac):
        database.clear_cache()
        return database.NIC.query.filter_by(mac_addr=mac).first()
    def find_interface(self, packet):
        chaddr = hwmac(packet.GetHardwareAddress())
        nic = self.findNIC(str(chaddr))
        return self.find_interface_by_nic(nic)
    def find_interface_by_nic(self, nic):
        if nic is None or nic.ip is None:
            return None
        ipstr = ''.join(reversed(['%02X' % i for i in ipv4(nic.ip.encode("utf-8")).list()]))
        for line in open('/proc/net/route'):
            parts = line.split()
            if parts[1] == ipstr:
                s.syslog(s.LOG_DEBUG, "find_interface found "+str(nic.ip)+" on "+parts[0])
                return parts[0]
        # Either the machine isn't running, or the route is missing.  We can
        # fix the latter.
        try:
            xsc = xs.xs()
            domid = xsc.read('', '/vm/%s/device/vif/0/frontend-id' % (nic.machine.uuid))
            # If we didn't find the domid, the machine is either off or the
            # UUID in xenstore isn't right.  Try slightly harder.
            if not domid:
                for uuid in xsc.ls('', '/vm'):
                    if xsc.read('', '/vm/%s/name' % (uuid)) == 'd_' + nic.machine.name:
                        domid = xsc.read('', '/vm/%s/device/vif/0/frontend-id' % (uuid))
            if not domid:
                xsc.close()
                return None
            for vifnum in xsc.ls('', '/local/domain/0/backend/vif/%s' % (domid)):
                if xsc.read('', '/local/domain/0/backend/vif/%s/%s/mac' % (domid, vifnum)) == nic.mac_addr:
                    # Prefer the tap if it exists; paravirtualized HVMs will
                    # have already unplugged it, so if it's there, it's the one
                    # in use.
                    for viftype in ('tap', 'vif'):
                        vif = '%s%s.%s' % (viftype, domid, vifnum)
                        if Interfaces.exists(vif):
                            self.add_route_and_arp(nic.ip, vif, nic.gateway)
                            xsc.close()
                            return vif
            xsc.close()
        except Exception as e:
            try:
                xsc.close()
            except Exception as e2:
                s.syslog(s.LOG_ERR, "Could not close connection to xenstore: %s" % (e2))
            s.syslog(s.LOG_ERR, "Could not find interface and add missing route: %s" % (e))
        return None
                            
    def getParameters(self, **extra):
        all_options=dict(dhcp_options)
        all_options.update(extra)
        options = {}
        for parameter, value in all_options.iteritems():
            if value is None:
                continue
            option_type = DhcpOptionsTypes[DhcpOptions[parameter]]

            if option_type == "ipv4" :
                # this is a single ip address
                options[parameter] = map(int,value.split("."))
            elif option_type == "ipv4+" :
                # this is multiple ip address
                iplist = value.split(",")
                opt = []
                for single in iplist :
                    opt.extend(ipv4(single).list())
                options[parameter] = opt
            elif option_type == "32-bits" :
                # This is probably a number...
                digit = int(value)
                options[parameter] = [digit>>24&0xFF,(digit>>16)&0xFF,(digit>>8)&0xFF,digit&0xFF]
            elif option_type == "16-bits" :
                digit = int(value)
                options[parameter] = [(digit>>8)&0xFF,digit&0xFF]

            elif option_type == "char" :
                digit = int(value)
                options[parameter] = [digit&0xFF]

            elif option_type == "bool" :
                if value=="False" or value=="false" or value==0 :
                    options[parameter] = [0]
                else : options[parameter] = [1]
                    
            elif option_type == "string" :
                options[parameter] = strlist(value).list()
            
            elif option_type == "RFC3397" :
                parsed_value = ""
                for item in value:
                    components = item.split('.')
                    item_fmt = "".join(chr(len(elt)) + elt for elt in components) + "\x00"
                    parsed_value += item_fmt
                
                options[parameter] = strlist(parsed_value).list()
            
            else :
                options[parameter] = strlist(value).list()
        return options

    def Discover(self, packet):
        s.syslog(s.LOG_DEBUG, "dhcp_backend : Discover ")
        chaddr = hwmac(packet.GetHardwareAddress())
        nic = self.findNIC(str(chaddr))
        if nic is None or nic.machine is None:
            return False
        ip = nic.ip.encode("utf-8")
        if ip is None:  #Deactivated?
            return False

        options = {}
        options['subnet_mask'] = nic.netmask.encode("utf-8")
        options['router'] = nic.gateway.encode("utf-8")
        if nic.hostname and '.' in nic.hostname:
            options['host_name'], options['domain_name'] = nic.hostname.encode('utf-8').split('.', 1)
        elif nic.machine.name:
            options['host_name'] = nic.machine.name.encode('utf-8')
            options['domain_name'] = config.dns.domains[0]
        else:
            hostname = None
        if DhcpOptions['domain_search'] in packet.GetOption('parameter_request_list'):
            options['host_name'] += '.' + options['domain_name']
            del options['domain_name']
            options['domain_search'] = [config.dhcp.search_domain]
        ip = ipv4(ip)
        s.syslog(s.LOG_DEBUG,"dhcp_backend : Discover result = "+str(ip))
        packet_parameters = self.getParameters(**options)

        # FIXME: Other offer parameters go here
        packet_parameters["yiaddr"] = ip.list()

        packet.SetMultipleOptions(packet_parameters)
        return True
        
    def Request(self, packet):
        s.syslog(s.LOG_DEBUG, "dhcp_backend : Request")
        
        discover = self.Discover(packet)
        
        chaddr = hwmac(packet.GetHardwareAddress())
        request = packet.GetOption("request_ip_address")
        if not request:
            request = packet.GetOption("ciaddr")
        yiaddr = packet.GetOption("yiaddr")

        if not discover:
            s.syslog(s.LOG_INFO,"Unknown MAC address: "+str(chaddr))
            return False
        
        if yiaddr!="0.0.0.0" and yiaddr == request :
            s.syslog(s.LOG_INFO,"Ack ip "+str(yiaddr)+" for "+str(chaddr))
            n = self.findNIC(str(chaddr))
            intf = self.find_interface_by_nic(n)
            s.syslog(s.LOG_ERR, "Interface is %s" % (intf))
            # Don't perform "other" actions if the machine isn't running
            other_action = n.other_action if n.other_action and intf else ''
            if other_action in ('renumber', 'renumber_dhcp'):
                (n.ip, n.netmask, n.gateway, n.other_ip, n.other_netmask,
                 n.other_gateway) = (
                 n.other_ip, n.other_netmask, n.other_gateway, n.ip,
                 n.netmask, n.gateway)
                other_action = n.other_action = 'dnat'
                database.session.add(n)
                database.session.flush()
            if other_action == 'dnat':
                # If the machine was booted in 'dnat' mode, then both
                # routes were already added by the invirt-database script.
                # If the machine was already on and has just been set to
                # 'dnat' mode, we need to add the route for the 'other' IP.
                # If the machine has just been 'renumbered' by us above,
                # the IPs will be swapped and only the route for the main
                # IP needs to be added.  Just try adding both of them, and
                # arp for whichever of them turns out to be new.
                for parms in [(n.ip, n.gateway), (n.other_ip, n.other_gateway)]:
                    self.add_route_and_arp(parms[0], intf, parms[1])
                try:
                    # iptables will let you add the same rule again and again;
                    # let's not do that.
                    p = Popen(['iptables', '-t', 'nat', '-C', 'PREROUTING', '-d', n.other_ip, '-j', 'DNAT', '--to-destination', n.ip], stdout=PIPE, stderr=PIPE)
                    (out, err) = p.communicate()
                    sys.stderr.write(err)
                    sys.stdout.write(out)
                    if p.returncode != 0:
                        p2 = Popen(['iptables', '-t', 'nat', '-A', 'PREROUTING', '-d', n.other_ip, '-j', 'DNAT', '--to-destination', n.ip], stdout=PIPE, stderr=PIPE)
                        (out, err) = p2.communicate()
                        sys.stderr.write(err)
                        sys.stdout.write(out)
                        if p2.returncode == 0:
                            s.syslog(s.LOG_INFO, "Added DNAT for IP %s to %s" % (n.other_ip, n.ip))
                        else:
                            s.syslog(s.LOG_ERR, "Could not add DNAT for IP %s to %s" % (n.other_ip, n.ip))
                except Exception as e:
                    s.syslog(s.LOG_ERR, "Could not check and/or add DNAT for IP %s to %s: %s" % (n.other_ip, n.ip, e))
            if other_action == 'remove':
                try:
                    p = Popen(['ip', 'route', 'del', n.other_ip, 'dev', intf], stdout=PIPE, stderr=PIPE)
                    (out, err) = p.communicate()
                    sys.stderr.write(err)
                    sys.stderr.write(out)
                    if p.returncode == 0:
                        s.syslog(s.LOG_INFO, "Removed route for IP %s" % (n.other_ip))
                    else:
                        s.syslog(s.LOG_ERR, "Could not remove route for IP %s" % (n.other_ip))
                except Exception as e:
                    s.syslog(s.LOG_ERR, "Could not run ip to remove route for IP %s: %s" % (n.other_ip, e))
                try:
                    p = Popen(['iptables', '-t', 'nat', '-D', 'PREROUTING', '-d', n.other_ip, '-j', 'DNAT', '--to-destination', n.ip], stdout=PIPE, stderr=PIPE)
                    (out, err) = p.communicate()
                    sys.stderr.write(err)
                    sys.stdout.write(out)
                    if p.returncode == 0:
                        s.syslog(s.LOG_INFO, "Removed DNAT for IP %s" % (n.other_ip))
                    else:
                        s.syslog(s.LOG_ERR, "Could not remove DNAT for IP %s" % (n.other_ip))
                except Exception as e:
                    s.syslog(s.LOG_ERR, "Could not run iptables to remove DNAT for IP %s: %s" % (n.other_ip, e))
                n.other_ip = n.other_netmask = n.other_gateway = n.other_action = None
                database.session.add(n)
                database.session.flush()
            # We went through the DISCOVER codepath already to populate some
            # of the packet's parameters.  If we renumbered the VM just above,
            # the packet is set to offer them what they asked for - the old
            # address.  So, we'll send them a DHCPNACK and they'll come right
            # back and be offered the new address.  The code above won't be
            # able to add duplicate routes, won't insert a duplicate DNAT,
            # and won't ARP again because the routes will exist, so this won't
            # incur much extra work.
            if request != map(int, n.ip.split('.')):
                return False
            return True
        else:
            s.syslog(s.LOG_INFO,"Requested ip "+str(request)+" not available for "+str(chaddr))
        return False

    def Decline(self, packet):
        pass
    def Release(self, packet):
        pass
    

class DhcpServer(pydhcplib.dhcp_network.DhcpServer):
    def __init__(self, backend, options = {'client_listenport':68,'server_listenport':67}):
        pydhcplib.dhcp_network.DhcpServer.__init__(self,"0.0.0.0",options["client_listen_port"],options["server_listen_port"],)
        self.backend = backend
        s.syslog(s.LOG_DEBUG, "__init__ DhcpServer")

    def SendDhcpPacketTo(self, To, packet):
        intf = self.backend.find_interface(packet)
        if intf:
            self.dhcp_socket.setsockopt(socket.SOL_SOCKET, IN.SO_BINDTODEVICE, intf)
            ret = self.dhcp_socket.sendto(packet.EncodePacket(), (To,self.emit_port))
            self.dhcp_socket.setsockopt(socket.SOL_SOCKET, IN.SO_BINDTODEVICE, '')
            return ret
        else:
            return self.dhcp_socket.sendto(packet.EncodePacket(),(To,self.emit_port))

    def SendPacket(self, packet):
        """Encode and send the packet."""
        
        giaddr = packet.GetOption('giaddr')

        # in all case, if giaddr is set, send packet to relay_agent
        # network address defines by giaddr
        if giaddr!=[0,0,0,0] :
            agent_ip = ".".join(map(str,giaddr))
            self.SendDhcpPacketTo(agent_ip,packet)
            s.syslog(s.LOG_DEBUG, "SendPacket to agent : "+agent_ip)

        # FIXME: This shouldn't broadcast if it has an IP address to send
        # it to instead. See RFC2131 part 4.1 for full details
        else :
            s.syslog(s.LOG_DEBUG, "No agent, broadcast packet.")
            self.SendDhcpPacketTo("255.255.255.255",packet)
            

    def HandleDhcpDiscover(self, packet):
        """Build and send DHCPOFFER packet in response to DHCPDISCOVER
        packet."""

        logmsg = "Get DHCPDISCOVER packet from " + hwmac(packet.GetHardwareAddress()).str()

        s.syslog(s.LOG_INFO, logmsg)
        offer = DhcpPacket()
        offer.CreateDhcpOfferPacketFrom(packet)
        
        if self.backend.Discover(offer):
            self.SendPacket(offer)
        # FIXME : what if false ?


    def HandleDhcpRequest(self, packet):
        """Build and send DHCPACK or DHCPNACK packet in response to
        DHCPREQUEST packet. 4 types of DHCPREQUEST exists."""

        ip = packet.GetOption("request_ip_address")
        sid = packet.GetOption("server_identifier")
        ciaddr = packet.GetOption("ciaddr")
        #packet.PrintHeaders()
        #packet.PrintOptions()

        if sid != [0,0,0,0] and ciaddr == [0,0,0,0] :
            s.syslog(s.LOG_INFO, "Get DHCPREQUEST_SELECTING_STATE packet")

        elif sid == [0,0,0,0] and ciaddr == [0,0,0,0] and ip :
            s.syslog(s.LOG_INFO, "Get DHCPREQUEST_INITREBOOT_STATE packet")

        elif sid == [0,0,0,0] and ciaddr != [0,0,0,0] and not ip :
            s.syslog(s.LOG_INFO,"Get DHCPREQUEST_INITREBOOT_STATE packet")

        else : s.syslog(s.LOG_INFO,"Get DHCPREQUEST_UNKNOWN_STATE packet : not implemented")

        if self.backend.Request(packet):
            packet.TransformToDhcpAckPacket()
            self.SendPacket(packet)
        elif self.backend.Discover(packet):
            packet.TransformToDhcpNackPacket()
            self.SendPacket(packet)
        else:
            pass # We aren't authoritative, so don't reply if we don't know them.

    # FIXME: These are not yet implemented.
    def HandleDhcpDecline(self, packet):
        s.syslog(s.LOG_INFO, "Get DHCPDECLINE packet")
        self.backend.Decline(packet)
        
    def HandleDhcpRelease(self, packet):
        s.syslog(s.LOG_INFO,"Get DHCPRELEASE packet")
        self.backend.Release(packet)
        
    def HandleDhcpInform(self, packet):
        s.syslog(s.LOG_INFO, "Get DHCPINFORM packet")

        if self.backend.Request(packet) :
            packet.TransformToDhcpAckPacket()
            # FIXME : Remove lease_time from options
            self.SendPacket(packet)

        # FIXME : what if false ?

class ArpspoofWorker(Thread):
    def __init__(self, queue):
        Thread.__init__(self)
        self.queue = queue
        self.iface = config.xen.iface

    def run(self):
        while True:
            (ip, gw) = self.queue.get()
            try:
                p = Popen(['timeout', '-s', 'KILL', '5', 'arpspoof', '-i', self.iface, '-t', gw, ip], stdout=PIPE, stderr=PIPE)
                (out, err) = p.communicate()
                if p.returncode != 124:
                    s.syslog(s.LOG_ERR, "arpspoof returned %s for IP %s gateway %s" % (p.returncode, ip, gw))
                else:
                    s.syslog(s.LOG_INFO, "aprspoof'd for IP %s gateway %s" % (ip, gw))
                sys.stderr.write(err)
                sys.stdout.write(out)
            except Exception as e:
                s.syslog(s.LOG_ERR, "Could not run arpspoof for IP %s gateway %s: %s" % (ip, gw, e))
            self.queue.task_done()

if '__main__' == __name__:
    options = { "server_listen_port":67,
                "client_listen_port":68,
                "listen_address":"0.0.0.0"}

    myip = socket.gethostbyname(socket.gethostname())
    if not myip:
        print "invirt-dhcpserver: cannot determine local IP address by looking up %s" % socket.gethostname()
        sys.exit(1)
    
    dhcp_options['server_identifier'] = ipv4(myip).int()

    queue = Queue()

    backend = DhcpBackend(queue)
    server = DhcpServer(backend, options)

    for x in range(2):
        worker = ArpspoofWorker(queue)
        worker.daemon = True
        worker.start()

    while True : server.GetNextDhcpPacket()