multi_interface_roam: multi_interface_roam: dhcp.py Source File

00001 #! /usr/bin/env python
00002 
00003 import roslib; roslib.load_manifest('multi_interface_roam')
00004 import asmach as smach
00005 import scapy.all as scapy
00006 from twisted.internet.defer import inlineCallbacks, returnValue
00007 from netlink_monitor import monitor, IFSTATE
00008 import async_helpers
00009 import l2socket
00010 import event
00011 import random
00012 import time
00013 import traceback
00014 import ipaddr
00015 import state_publisher
00016 import config
00017 
00018 # FIXME Add support for multiple leases.
00019 # FIXME Add a way to signal x amount of time before we lose the lease.
00020 
00021 class DhcpLease:
00022     def __init__(self):
00023         self.timeout_time = {}
00024         self.public_config = {}
00025 
00026 class DhcpData:
00027     def __init__(self, iface):
00028         self.iface = iface
00029         self.socket = None
00030         self.link_addr_state = monitor.get_state_publisher(self.iface, IFSTATE.LINK_ADDR) 
00031         self.error_event = event.Event()
00032         self.binding_publisher = state_publisher.StatePublisher(None)
00033         self.error_timeout = 5
00034         self.exp_backoff_min = 0.2
00035         self.exp_backoff_max = 0.5
00036         self.exp_backoff_timeout = config.get_parameter('dhcp_timeout', 2)
00037         self.leases = {}
00038 
00039     def start_socket(self):
00040         if not self.socket or self.socket.port.fileno() == -1:
00041             self.socket = async_helpers.ReadDescrEventStream(l2socket.L2Port, iface = self.iface, 
00042                     filter='udp and dst port 68 and src port 67')
00043 
00044     def stop_socket(self):
00045         self.socket = None
00046 
00047 class DhcpState(smach.State):
00048     def __init__(self, *args, **kwargs):
00049         smach.State.__init__(self, input_keys=['dhcp'], output_keys=['dhcp'], *args, **kwargs)
00050 
00051 def find_dhcp_option(key, default, dhcp):
00052     for opt in dhcp.options:
00053         if opt == key:
00054             return None
00055         if opt[0] == key:
00056             return opt[1:]
00057     return default
00058 
00059 
00060 def dhcp_type_str(id):
00061     try:
00062         return scapy.DHCPTypes[id[0]]
00063     except KeyError:
00064         return "<unknown %s>"%repr(id)
00065  
00066 
00067 
00068 class ExchangeRetryExponentialBackoff:
00069     def init_timeouts(self, ud):
00070         self.cur_retry_max = ud.dhcp.exp_backoff_min
00071         return ud.dhcp.exp_backoff_timeout
00072 
00073     def get_next_retry(self, ud):
00074         interval = self.cur_retry_max * random.uniform(0.5, 1)
00075         self.cur_retry_max = min(ud.dhcp.exp_backoff_max, 2 * self.cur_retry_max)
00076         return interval
00077 
00078 
00079 
00080 class ExchangeRetryHalve:
00081     def init_timeouts(self, ud):
00082         return ud.dhcp.lease.timeout_time[self.type] - time.time()
00083 
00084     def get_next_retry(self, ud):
00085         return 60 + (ud.dhcp.lease.timeout_time[self.type] - time.time()) / 2
00086 
00087 
00088         
00089 class ExchangeDiscover:
00090     message_type = "discover"
00091 
00092     def init_xid(self, ud):
00093         ud.dhcp.lease.xid = random.randint(0, 0xFFFF)
00094 
00095     def validate(self, ud, pkt):
00096         try:
00097             ip = pkt.payload
00098             udp = ip.payload 
00099             bootp = udp.payload
00100             dhcp = bootp.payload
00101             
00102             if not self.validate_common(ud, pkt, ip, udp, bootp, dhcp):
00103                 print "validate_common returned False"
00104                 return False
00105 
00106             message_type = dhcp_type_str(find_dhcp_option('message-type', None, dhcp))
00107             if message_type != 'offer':
00108                 print "Ignoring packet based on unexpected message_type %s"%message_type
00109                 return False
00110                 
00111             ip = ud.dhcp.lease.public_config['ip'] = bootp.yiaddr
00112             ud.dhcp.lease.server_ip = find_dhcp_option('server_id', "0.0.0.0", dhcp)[0]
00113             ud.dhcp.lease.server_mac = pkt.src
00114             ud.dhcp.lease.public_config['gateway'] = find_dhcp_option('router', "0.0.0.0", dhcp)[0]
00115             netmask = ud.dhcp.lease.public_config['netmask'] = find_dhcp_option('subnet_mask', "0.0.0.0", dhcp)[0]
00116             net = ipaddr.IPv4Network("%s/%s"%(ip, netmask))
00117             ud.dhcp.lease.public_config['netmask_bits'] = net.prefixlen
00118             ud.dhcp.lease.public_config['network'] = net.network
00119             ud.dhcp.lease.public_config['ip_slashed'] = "%s/%i"%(ip, net.prefixlen)
00120             ud.dhcp.lease.public_config['network_slashed'] = "%s/%i"%(net.network, net.prefixlen)
00121 
00122             # TODO do more stuff here?
00123             return 'success'
00124             
00125 
00126         except:
00127             traceback.print_exc()
00128             print "Excepiton validating packet."
00129             return False
00130 
00131 
00132 class ExchangeRequest:
00133     message_type = "request"
00134     
00135     def init_xid(self, ud):
00136         pass
00137         
00138     def validate(self, ud, pkt):
00139         try:
00140             ip = pkt.payload
00141             udp = ip.payload 
00142             bootp = udp.payload
00143             dhcp = bootp.payload
00144             
00145             if not self.validate_common(ud, pkt, ip, udp, bootp, dhcp):
00146                 print "validate_common returned False"
00147                 return False
00148 
00149             message_type = dhcp_type_str(find_dhcp_option('message-type', None, dhcp))
00150             if message_type == 'nak':
00151                 return 'fail' # TODO Check that from right server?
00152             
00153             if message_type != 'ack':
00154                 print "Ignoring packet based on unexpected message_type %s"%message_type
00155                 return False
00156 
00157             lease_time = find_dhcp_option('lease_time', None, dhcp)
00158             if not lease_time:
00159                 print "Ignoring packet with no lease_time"
00160             lease_time = lease_time[0] 
00161             renewal_time = find_dhcp_option('renewal_time', (random.uniform(0.45, 0.55) * lease_time, ), dhcp)[0]
00162             rebind_time = find_dhcp_option('rebinding_time', (random.uniform(0.825, 0.925) * lease_time, ), dhcp)[0]
00163             #lease_time = 10
00164             #rebind_time = 2
00165             #renewal_time = 1
00166             ud.dhcp.lease.timeout_time['BOUND'] = renewal_time * 0.99 + self.send_time
00167             ud.dhcp.lease.timeout_time['RENEW'] = rebind_time * 0.99 + self.send_time
00168             ud.dhcp.lease.timeout_time['REBIND'] = lease_time * 0.99 + self.send_time
00169 
00170 
00171             # FIXME Read other options here.
00172             return 'success'
00173 
00174         except:
00175             traceback.print_exc()
00176             print "Excepiton validating packet."
00177             return False
00178 
00179 
00180 
00181 class NoLink(DhcpState):
00182     def __init__(self):
00183         DhcpState.__init__(self, outcomes=['bound', 'init', 'init_reboot'])
00184 
00185     @inlineCallbacks
00186     def execute_async(self, ud):
00187         ud.dhcp.binding_publisher.set(None)
00188         ud.dhcp.hwaddr = yield async_helpers.wait_for_state(ud.dhcp.link_addr_state, lambda x: x != False)
00189         network_id = "" # Include MAC, network_id
00190         if network_id not in ud.dhcp.leases:
00191             ud.dhcp.leases[network_id] = DhcpLease()
00192         ud.dhcp.lease = ud.dhcp.leases[network_id]
00193         returnValue('init')
00194 
00195 
00196 
00197 class Init(DhcpState):
00198     def __init__(self):
00199         DhcpState.__init__(self, outcomes=['done', 'nolink'])
00200 
00201     def execute_async(self, ud):
00202         ud.dhcp.binding_publisher.set(None)
00203         ud.dhcp.start_socket()
00204         return 'done'
00205 
00206 ETHER_BCAST='ff:ff:ff:ff:ff:ff'
00207 IP_BCAST='255.255.255.255'
00208 IP_ZERO='0.0.0.0'
00209 
00210 class Exchange(DhcpState):
00211     def __init__(self):
00212         DhcpState.__init__(self, outcomes=['success', 'fail', 'nolink'])
00213 
00214     def send(self, ud):
00215         hwbytes = scapy.mac2str(ud.dhcp.hwaddr)
00216 
00217         # Prepare options
00218         options = [
00219                 ("message-type", self.message_type), 
00220                 ("param_req_list", 
00221                     chr(scapy.DHCPRevOptions["renewal_time"][0]),
00222                     chr(scapy.DHCPRevOptions["rebinding_time"][0]),
00223                     chr(scapy.DHCPRevOptions["lease_time"][0]),
00224                     chr(scapy.DHCPRevOptions["subnet_mask"][0]),
00225                     chr(scapy.DHCPRevOptions["router"][0]),
00226                     )
00227                 ]
00228         if self.type in ["REQUEST", "REBOOT", ]:
00229             options.append(('requested_addr', ud.dhcp.lease.public_config['ip']))
00230         if self.type in ["REQUEST", ]:
00231             options.append(('server_id', ud.dhcp.lease.server_ip))
00232         options.append('end')
00233         pkt = scapy.DHCP(options=options)
00234 
00235         # Prepare BOOTP
00236         pkt = scapy.BOOTP(chaddr=[hwbytes], xid=ud.dhcp.lease.xid)/pkt
00237         if self.type in [ "RENEW", "REBIND", ]:
00238             pkt.ciaddr = ud.dhcp.lease.public_config['ip'] 
00239 
00240         # Prepare UDP/IP
00241         pkt = scapy.IP(src=pkt.ciaddr, dst=IP_BCAST)/scapy.UDP(sport=68, dport=67)/pkt
00242         if self.type in [ "RENEW" ]:
00243             pkt.dst = ud.dhcp.lease.server_ip
00244 
00245         # Prepare Ethernet
00246         pkt = scapy.Ether(src=ud.dhcp.hwaddr, dst=ETHER_BCAST)/pkt
00247         if self.type in [ "RENEW" ]:
00248             pkt.dst = ud.dhcp.lease.server_mac
00249         #print "Out:", repr(scapy.Ether(str(pkt)))
00250         ud.dhcp.socket.port.send(str(pkt))
00251                     
00252     def validate_common(self, ud, pkt, ip, udp, bootp, dhcp): 
00253         # Should we be receiving this packet?
00254         if pkt.dst != ud.dhcp.hwaddr and pkt.dst != ETHER_BCAST:
00255             print "Discarding packet based on destination MAC: %s != %s"%(pkt.dst, ud.dhcp.hwaddr)
00256             return False
00257 
00258         # Does the xid match?
00259         if pkt.xid != ud.dhcp.lease.xid:
00260             print "Discarding packet based on xid: %i != %i"%(pkt.xid, ud.dhcp.lease.xid)
00261             return False
00262        
00263         # TODO Check that from right server?
00264         
00265         return True
00266 
00267     @inlineCallbacks
00268     def execute_async(self, ud):
00269         # Parameters that will depend on the state.
00270         timeout = async_helpers.Timeout(self.init_timeouts(ud))
00271 
00272         # Make sure we aren't discarding incoming dhcp packets.
00273         ud.dhcp.socket.set_discard(False)
00274 
00275         # Generate an xid for the exchange if necessary.
00276         self.init_xid(ud)
00277 
00278         while True:
00279             # Send a request packet
00280             try:
00281                 self.send_time = time.time()
00282                 self.send(ud)
00283             except:
00284                 traceback.print_exc()
00285                 returnValue('fail')
00286 
00287             # How long to wait before retry
00288             interval = self.get_next_retry(ud)
00289 
00290             while True:
00291                 # Wait for an event
00292                 events = yield async_helpers.select(
00293                         async_helpers.StateCondition(ud.dhcp.link_addr_state, lambda x: x == False), 
00294                         ud.dhcp.socket, 
00295                         async_helpers.Timeout(interval),
00296                         timeout)
00297 
00298                 if 0 in events: # Lost link
00299                     returnValue('nolink')
00300     
00301                 if 1 in events: # Got packet
00302                     pkt = scapy.Ether(ud.dhcp.socket.recv())
00303                     #print "In:", repr(pkt)
00304 
00305                     result = self.validate(ud, pkt)
00306                     if result:
00307                         returnValue(result)
00308     
00309                 if 2 in events:
00310                     break
00311     
00312                 if 3 in events:
00313                     returnValue('fail')
00314 
00315 
00316 
00317 class Rebooting  (Exchange, ExchangeRetryExponentialBackoff, ExchangeRequest ): type = "REBOOT"
00318 class Selecting  (Exchange, ExchangeRetryExponentialBackoff, ExchangeDiscover): type = "SELECT"
00319 class Requesting (Exchange, ExchangeRetryExponentialBackoff, ExchangeRequest ): type = "REQUEST"
00320 class Renewing   (Exchange, ExchangeRetryHalve,              ExchangeRequest ): type = "RENEW"
00321 class Rebinding  (Exchange, ExchangeRetryHalve,              ExchangeRequest ): type = "REBIND"
00322 
00323 
00324 
00325 class Error(DhcpState):
00326     def __init__(self):
00327         DhcpState.__init__(self, outcomes=['done', 'nolink'])
00328  
00329     @inlineCallbacks
00330     def execute_async(self, ud):
00331         ud.dhcp.binding_publisher.set(None)
00332         ud.dhcp.error_event.trigger()
00333         ud.dhcp.socket.set_discard(True)
00334         events = yield async_helpers.select(
00335                 async_helpers.StateCondition(ud.dhcp.link_addr_state, lambda x: x == False),
00336                 async_helpers.Timeout(ud.dhcp.error_timeout)
00337                 )
00338 
00339         if 0 in events:
00340             returnValue('nolink')
00341 
00342         returnValue('done')
00343 
00344 
00345 class Bound(DhcpState):
00346     def __init__(self):
00347         DhcpState.__init__(self, outcomes=['timeout', 'nolink'])
00348  
00349     @inlineCallbacks
00350     def execute_async(self, ud):
00351         ud.dhcp.socket.set_discard(True)
00352         ud.dhcp.binding_publisher.set(ud.dhcp.lease.public_config)
00353         events = yield async_helpers.select(
00354             async_helpers.StateCondition(ud.dhcp.link_addr_state, lambda x: x == False),
00355             async_helpers.Timeout(ud.dhcp.lease.timeout_time['BOUND'] - time.time()),
00356             )
00357 
00358         if 0 in events:
00359             returnValue('nolink')
00360 
00361         returnValue('timeout')
00362 
00363 
00364 def dhcp_client(iface):
00365     sm = smach.StateMachine(outcomes=[], input_keys=['dhcp'])
00366     smadd = smach.StateMachine.add
00367     with sm:
00368         smadd('NOLINK',      NoLink(),      transitions = {'bound'  :'BOUND', 'init':'INIT', 'init_reboot':'INIT_REBOOT'})
00369         smadd('INIT_REBOOT', Init(),        transitions = {'done'   :'REBOOT',                   'nolink':'NOLINK'})
00370         smadd('REBOOT',      Rebooting(),   transitions = {'success':'BOUND',   'fail':'INIT',   'nolink':'NOLINK'})
00371         smadd('INIT',        Init(),        transitions = {'done'   :'SELECT',                   'nolink':'NOLINK'})
00372         smadd('SELECT',      Selecting(),   transitions = {'success':'REQUEST', 'fail':'ERROR',  'nolink':'NOLINK'})
00373         smadd('REQUEST',     Requesting(),  transitions = {'success':'BOUND',   'fail':'ERROR',  'nolink':'NOLINK'})
00374         smadd('BOUND',       Bound(),       transitions = {'timeout':'RENEW',                    'nolink':'NOLINK'})
00375         smadd('RENEW',       Renewing(),    transitions = {'success':'BOUND',   'fail':'REBIND', 'nolink':'NOLINK'})
00376         smadd('REBIND',      Rebinding(),   transitions = {'success':'BOUND',   'fail':'INIT',   'nolink':'NOLINK'})
00377         smadd('ERROR',       Error(),       transitions = {'done'   :'INIT',                     'nolink':'NOLINK'})
00378 
00379     ud = smach.UserData()
00380     ud.dhcp = DhcpData(iface)
00381 #    def shutdown(value):
00382 #        reactor.fireSystemEvent('shutdown')
00383 #    def ignore_eintr(error):
00384 #        if error.type == IOError:
00385 #            import errno
00386 #            if error.value.errno == errno.EINTR:
00387 #                return None
00388 #        return error
00389     sm.execute_async(ud)#.addCallback(shutdown).addErrback(ignore_eintr)
00390     return ud.dhcp
00391 
00392 if __name__ == "__main__":
00393     import sys
00394     from twisted.internet import reactor
00395 
00396     if len(sys.argv) != 2:
00397         print "usage: dhcp.py <interface>"
00398         sys.exit(1)
00399 
00400     iface = sys.argv[1]
00401 
00402     dhcp_client(iface)
00403     reactor.run()