source: branches/wsgi/packages/sipb-xen-dhcp/code/dhcpserver.py @ 1597

#!/usr/bin/python
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

import event_logger
if '__main__' == __name__:
    event_logger.init("stdout", 'DEBUG', {})
from event_logger import Log

import psycopg2
import time
import sipb_xen_database
from sqlalchemy import create_engine

dhcp_options = {'subnet_mask': '255.255.0.0',
                'router': '18.181.0.1',
                'domain_name_server': '18.70.0.160,18.71.0.151,18.72.0.3',
                'ip_address_lease_time': 60*60*24}

class DhcpBackend:
    def __init__(self, database=None):
        if database is not None:
            self.database = database
            sipb_xen_database.connect(create_engine(database))
    def findNIC(self, mac):
        sipb_xen_database.clear_cache()
        for i in range(3):
            try:
                value = sipb_xen_database.NIC.get_by(mac_addr=mac)
            except psycopg2.OperationalError:
                time.sleep(0.5)
                if i == 2:  #Try twice to reconnect.
                    raise
                #Sigh.  SQLAlchemy should do this itself.
                sipb_xen_database.connect(create_engine(self.database))
            else:
                break
        return value
    def find_interface(self, packet):
        chaddr = hwmac(packet.GetHardwareAddress())
        nic = self.findNIC(str(chaddr))
        if nic is None or nic.ip is None:
            return ("18.181.0.60", None)
        ipstr = ''.join(reversed(['%02X' % i for i in ipv4(nic.ip).list()]))
        for line in open('/proc/net/route'):
            parts = line.split()
            if parts[1] == ipstr:
                Log.Output(Log.debug, "find_interface found "+str(nic.ip)+" on "+parts[0])
                return ("18.181.0.60", parts[0])
        return ("18.181.0.60", 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):
        Log.Output(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
        if ip is None:  #Deactivated?
            return False

        options = {}
        if nic.hostname and '.' in nic.hostname:
            options['host_name'], options['domain_name'] = nic.hostname.split('.', 1)
        elif nic.machine.name:
            options['host_name'] = nic.machine.name
            options['domain_name'] = 'xvm.mit.edu'
        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'] = ['mit.edu']
        if ip is not None:
            ip = ipv4(ip)
            Log.Output(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
        return False
        
    def Request(self, packet):
        Log.Output(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:
            Log.Output(Log.info,"Unknown MAC address: "+str(chaddr))
            return False
        
        if yiaddr!="0.0.0.0" and yiaddr == request :
            Log.Output(Log.info,"Ack ip "+str(yiaddr)+" for "+str(chaddr))
            return True
        else:
            Log.Output(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
        Log.Output(Log.debug, "__init__ DhcpServer")

    def SendDhcpPacketTo(self, To, packet):
        (ip, intf) = self.backend.find_interface(packet)
        if intf:
            out_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            out_socket.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST,1)
            out_socket.setsockopt(socket.SOL_SOCKET, IN.SO_BINDTODEVICE, intf)
            #out_socket.bind((ip, self.listen_port))
            ret = out_socket.sendto(packet.EncodePacket(), (To,self.emit_port))
            out_socket.close()
            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)
            Log.Output(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 :
            Log.Output(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()

        Log.Output(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] :
            Log.Output(Log.info, "Get DHCPREQUEST_SELECTING_STATE packet")

        elif sid == [0,0,0,0] and ciaddr == [0,0,0,0] and ip :
            Log.Output(Log.info, "Get DHCPREQUEST_INITREBOOT_STATE packet")

        elif sid == [0,0,0,0] and ciaddr != [0,0,0,0] and not ip :
            Log.Output(Log.info,"Get DHCPREQUEST_INITREBOOT_STATE packet")

        else : Log.Output(Log.info,"Get DHCPREQUEST_UNKNOWN_STATE packet : not implemented")

        if self.backend.Request(packet) : packet.TransformToDhcpAckPacket()
        else : packet.TransformToDhcpNackPacket()

        self.SendPacket(packet)



    # FIXME: These are not yet implemented.
    def HandleDhcpDecline(self, packet):
        Log.Output(Log.info, "Get DHCPDECLINE packet")
        self.backend.Decline(packet)
        
    def HandleDhcpRelease(self, packet):
        Log.Output(Log.info,"Get DHCPRELEASE packet")
        self.backend.Release(packet)
        
    def HandleDhcpInform(self, packet):
        Log.Output(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 ?

if '__main__' == __name__:
    options = { "server_listen_port":67,
                "client_listen_port":68,
                "listen_address":"0.0.0.0"}
    backend = DhcpBackend('postgres://sipb-xen@sipb-xen-dev/sipb_xen')
    server = DhcpServer(backend, options)

    while True : server.GetNextDhcpPacket()