Merge pull request #18 from bias-keenly/feature/dns

Add a flag to specify DNS servers that DHCP clients receive

ptftplib/dhcpserver.py

@@ -58,6 +58,7 @@ DHCP_OP_DHCPACK = 5
 # DHCP options we care about.
 DHCP_OPTION_SUBNET = 1                # Subnet mask
 DHCP_OPTION_ROUTER = 3                # Router
+DHCP_OPTION_DNS = 6                   # Domain Name Servers (DNS)
 DHCP_OPTION_REQUESTED_IP = 50         # Requested IP address
 DHCP_OPTION_LEASE_TIME = 51           # Lease time for the IP address
 DHCP_OPTION_OP = 53                   # The DHCP operation (see above)
@@ -226,9 +227,10 @@ class DhcpPacket(object):
 
 class DHCPServer(object):
     def __init__(self, interface, bootfile, router=None, tftp_server=None,
-                 answer_all_requests=False):
+                 answer_all_requests=False, name_servers=None):
         self.interface = interface
         self.ip, self.netmask, self.mac = get_ip_config_for_iface(interface)
+        self.name_servers = name_servers
         self.bootfile = bootfile
         self.router = router or self.ip
         self.tftp_server = tftp_server or self.ip
@@ -306,13 +308,20 @@ class DHCPServer(object):
             DHCP_OP_DHCPDISCOVER: DHCP_OP_DHCPOFFER,
             DHCP_OP_DHCPREQUEST: DHCP_OP_DHCPACK
             }[request_pkt.op]
-        dhcp_options = (
+        dhcp_options = [
             (DHCP_OPTION_OP, chr(reply_kind)),
             (DHCP_OPTION_LEASE_TIME, struct.pack('!L', DHCP_LEASE_TIMEOUT)),
             (DHCP_OPTION_SUBNET, _pack_ip(self.netmask)),
             (DHCP_OPTION_ROUTER, _pack_ip(self.router)),
             (DHCP_OPTION_SERVER_ID, _pack_ip(self.ip)),
-            )
+            ]
+
+        if self.name_servers:
+            dns_servers = ''
+            for name_server in self.name_servers:
+                dns_servers += _pack_ip(name_server)
+            dhcp_options.append((DHCP_OPTION_DNS, dns_servers))
+
         buf = []
         for code, data in dhcp_options:
             buf.append(struct.pack('!BB', code, len(data)))
@@ -386,6 +395,9 @@ def main():
     parser.add_option("-a", "--answer-all-dhcp-requests", dest="answer_all_requests",
                       help="Enables DHCP response to all clients, "
                            "default is PXE clients only", action="store_true", default=False)
+    parser.add_option("-n", "--name-servers", dest="name_servers", action='append', metavar='NAME_SERVER',
+                      help="Domain Name Servers (DNS) IPs to provide to DHCP client. "
+                           "Use multiple flags to specify up to 3 DNS servers", default=None)
     parser.add_option("-v", "--verbose", dest="loglevel", action="store_const",
                       const=logging.INFO, help="Output information messages",
                       default=logging.WARNING)
@@ -396,6 +408,11 @@ def main():
         parser.print_help()
         return 1
 
+    if options.name_servers and len(options.name_servers) > 3:
+        print('Error: up to 3 DNS servers allowed\n')
+        parser.print_help()
+        return 1
+
     iface, bootfile = args
 
     # Setup notification logging

ptftplib/pxeserver.py

@@ -42,11 +42,12 @@ l = notify.getLogger('pxed')
 
 
 class DHCPThread(threading.Thread):
-    def __init__(self, iface, bootfile, router, answer_all_requests):
+    def __init__(self, iface, bootfile, router, answer_all_requests, name_servers):
         threading.Thread.__init__(self)
         self.setDaemon(True)
         self.server = dhcpserver.DHCPServer(iface, bootfile, router=router,
-                                            answer_all_requests=answer_all_requests)
+                                            answer_all_requests=answer_all_requests,
+                                            name_servers=name_servers)
 
     def run(self):
         self.server.serve_forever()
@@ -67,6 +68,9 @@ def main():
     parser.add_option("-a", "--answer-all-dhcp-requests", dest="answer_all_requests",
                       help="Enables DHCP response to all clients, "
                            "default is PXE clients only", action="store_true", default=False)
+    parser.add_option("-n", "--name-servers", dest="name_servers", action='append', metavar='NAME_SERVER',
+                      help="Domain Name Servers (DNS) IPs to provide to DHCP client. "
+                           "Use multiple flags to specify up to 3 DNS servers", default=None)
     parser.add_option("-v", "--verbose", dest="loglevel", action="store_const",
                       const=logging.INFO, help="Output information messages",
                       default=logging.WARNING)
@@ -78,6 +82,11 @@ def main():
         parser.print_help()
         return 1
 
+    if options.name_servers and len(options.name_servers) > 3:
+        print('Error: up to 3 DNS servers allowed\n')
+        parser.print_help()
+        return 1
+
     iface, root, bootfile = args
 
     # Setup notification logging
@@ -92,7 +101,8 @@ def main():
                                 fmt='%(levelname)s(%(name)s): %(message)s')
 
     try:
-        dhcp = DHCPThread(iface, bootfile, options.router, options.answer_all_requests)
+        dhcp = DHCPThread(iface, bootfile, options.router, options.answer_all_requests,
+                          options.name_servers)
         tftp = tftpserver.TFTPServer(iface, root,
                                      strict_rfc1350=options.strict_rfc1350)
     except tftpserver.TFTPServerConfigurationError as e: