Something that is hopefully a DHCP server package
[invirt/packages/invirt-dhcp.git] / code / dhcpserver.py
diff --git a/code/dhcpserver.py b/code/dhcpserver.py
new file mode 100644 (file)
index 0000000..3133643
--- /dev/null
@@ -0,0 +1,265 @@
+#!/usr/bin/python
+import sys
+sys.path.append('pydhcplib/')
+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',
+                'domain_name': 'mit.edu',
+                '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):
+        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()
+                
+            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
+        if nic.hostname and '.' in nic.hostname:
+            hostname = nic.hostname
+        elif nic.machine.name:
+            hostname = nic.machine.name + '.servers.csail.mit.edu'
+        else:
+            hostname = None
+        if ip is not None:
+            ip = ipv4(ip)
+            Log.Output(Log.debug,"dhcp_backend : Discover result = "+str(ip))
+            packet_parameters = self.getParameters(host_name=hostname)
+
+            # 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()