CSIT-341 Refactor DHCPv4 proxy test
[csit.git] / resources / traffic_scripts / dhcp / send_and_check_proxy_messages.py
@@ -25,27 +25,18 @@ from resources.libraries.python.PacketVerifier import RxQueue, TxQueue
 from resources.libraries.python.TrafficScriptArg import TrafficScriptArg
 
 
-def dhcp_discover(args):
-    """Send DHCP DISCOVER packet."""
-
-    tx_if = args.get_arg('tx_if')
-    rx_if = args.get_arg('rx_if')
+def dhcp_discover(tx_if, rx_if, tx_src_ip, tx_dst_ip, server_ip, proxy_ip,
+                  client_mac):
+    """Send and check DHCP DISCOVER proxy packet."""
 
     rxq = RxQueue(rx_if)
     txq = TxQueue(tx_if)
 
-    tx_src_ip = "0.0.0.0"
-    tx_dst_ip = "255.255.255.255"
-
-    server_ip = args.get_arg('server_ip')
-    proxy_ip = args.get_arg('proxy_ip')
-    client_mac = args.get_arg('client_mac')
-
     sent_packets = []
 
     dhcp_discover = Ether(src=client_mac, dst="ff:ff:ff:ff:ff:ff") / \
                     IP(src=tx_src_ip, dst=tx_dst_ip) / \
-                    UDP(sport=68, dport=67) / \
+                    UDP(sport=UDP_SERVICES.bootpc, dport=UDP_SERVICES.bootps) / \
                     BOOTP(op=1,) / \
                     DHCP(options=[("message-type", "discover"),
                                   "end"])
@@ -81,29 +72,22 @@ def dhcp_discover(args):
     if ether[DHCP].options[0][1] != 1:  # 1 - DISCOVER message
         raise RuntimeError("DHCP DISCOVER message error.")
     print "DHCP DISCOVER message OK."
-    dhcp_offer(args, option_82)
 
+    return option_82
 
-def dhcp_offer(args, option_82):
-    """Send DHCP OFFER packet."""
 
-    rx_if = args.get_arg('tx_if')
-    tx_if = args.get_arg('rx_if')
+def dhcp_offer(rx_if, tx_if, tx_dst_ip, server_ip, proxy_ip, client_ip,
+               server_mac, option_82):
+    """Send and check DHCP OFFER proxy packet."""
 
     rxq = RxQueue(rx_if)
     txq = TxQueue(tx_if)
 
-    tx_dst_ip = "255.255.255.255"
-    server_ip = args.get_arg('server_ip')
-    server_mac = args.get_arg('server_mac')
-    client_ip = args.get_arg('client_ip')
-    proxy_ip = args.get_arg('proxy_ip')
-
     sent_packets = []
 
     dhcp_offer = Ether(src=server_mac, dst="ff:ff:ff:ff:ff:ff") / \
                  IP(src=server_ip, dst=tx_dst_ip) / \
-                 UDP(sport=67, dport=68) / \
+                 UDP(sport=UDP_SERVICES.bootps, dport=UDP_SERVICES.bootpc) / \
                  BOOTP(op=2,
                        yiaddr=client_ip,
                        siaddr=server_ip) / \
@@ -148,29 +132,20 @@ def dhcp_offer(args, option_82):
     if ether[DHCP].options[0][1] != 2:  # 2 - OFFER message
         raise RuntimeError("DHCP OFFER message error.")
     print "DHCP OFFER message OK."
-    dhcp_request(args)
 
 
-def dhcp_request(args):
-    """Send DHCP REQUEST packet."""
-
-    tx_if = args.get_arg('tx_if')
-    rx_if = args.get_arg('rx_if')
+def dhcp_request(tx_if, rx_if, tx_src_ip, tx_dst_ip, server_ip, proxy_ip,
+                 client_ip, client_mac):
+    """Send and check DHCP REQUEST proxy packet."""
 
     rxq = RxQueue(rx_if)
     txq = TxQueue(tx_if)
 
-    tx_dst_ip = "255.255.255.255"
-    server_ip = args.get_arg('server_ip')
-    client_ip = args.get_arg('client_ip')
-    client_mac = args.get_arg('client_mac')
-    proxy_ip = args.get_arg('proxy_ip')
-
     sent_packets = []
 
     dhcp_request = Ether(src=client_mac, dst="ff:ff:ff:ff:ff:ff") / \
-                   IP(src="0.0.0.0", dst=tx_dst_ip) / \
-                   UDP(sport=68, dport=67) / \
+                   IP(src=tx_src_ip, dst=tx_dst_ip) / \
+                   UDP(sport=UDP_SERVICES.bootpc, dport=UDP_SERVICES.bootps) / \
                    BOOTP(op=1,
                          giaddr=proxy_ip,
                          siaddr=server_ip) / \
@@ -213,35 +188,26 @@ def dhcp_request(args):
 
     if ether[DHCP].options[3][0] != 'relay_agent_Information':  # option 82
         raise RuntimeError("Relay agent information error.")
-    option_82 = ether[DHCP].options[3][1]
 
     if ether[DHCP].options[0][1] != 3:  # 2 - REQUEST message
         raise RuntimeError("DHCP REQUEST message error.")
     print "DHCP REQUEST message: OK."
-    dhcp_ack(args, option_82)
-
 
-def dhcp_ack(args, option_82):
-    """Send DHCP ACK packet."""
 
-    rx_if = args.get_arg('tx_if')
-    tx_if = args.get_arg('rx_if')
+def dhcp_ack(rx_if, tx_if, tx_dst_ip, server_ip, proxy_ip, client_ip,
+             server_mac, option_82):
+    """Send and check DHCP ACK proxy packet."""
 
     rxq = RxQueue(rx_if)
     txq = TxQueue(tx_if)
 
-    tx_dst_ip = "255.255.255.255"
-    server_ip = args.get_arg('server_ip')
-    server_mac = args.get_arg('server_mac')
-    client_ip = args.get_arg('client_ip')
-    proxy_ip = args.get_arg('proxy_ip')
     lease_time = 43200  # 12 hours
 
     sent_packets = []
 
     dhcp_ack = Ether(src=server_mac, dst="ff:ff:ff:ff:ff:ff") / \
                IP(src=server_ip, dst=tx_dst_ip) / \
-               UDP(sport=67, dport=68) / \
+               UDP(sport=UDP_SERVICES.bootps, dport=UDP_SERVICES.bootpc) / \
                BOOTP(op=2,
                      yiaddr=client_ip,
                      siaddr=server_ip) / \
@@ -294,12 +260,38 @@ def dhcp_ack(args, option_82):
 
 
 def main():
-    """Send DHCP messages."""
+    """Send DHCP proxy messages."""
 
     args = TrafficScriptArg(['server_ip', 'server_mac', 'client_ip',
                              'client_mac', 'proxy_ip'])
 
-    dhcp_discover(args)
+    tx_if = args.get_arg('tx_if')
+    rx_if = args.get_arg('rx_if')
+
+    tx_src_ip = "0.0.0.0"
+    tx_dst_ip = "255.255.255.255"
+
+    server_ip = args.get_arg('server_ip')
+    client_ip = args.get_arg('client_ip')
+    proxy_ip = args.get_arg('proxy_ip')
+    client_mac = args.get_arg('client_mac')
+    server_mac = args.get_arg('server_mac')
+
+    # DHCP DISCOVER
+    option_82 = dhcp_discover(tx_if, rx_if, tx_src_ip, tx_dst_ip, server_ip,
+                              proxy_ip, client_mac)
+
+    # DHCP OFFER
+    dhcp_offer(tx_if, rx_if, tx_dst_ip, server_ip, proxy_ip, client_ip,
+               server_mac, option_82)
+
+    # DHCP REQUEST
+    dhcp_request(tx_if, rx_if, tx_src_ip, tx_dst_ip, server_ip, proxy_ip,
+                 client_ip, client_mac)
+
+    # DHCP ACK
+    dhcp_ack(tx_if, rx_if, tx_dst_ip, server_ip, proxy_ip, client_ip,
+             server_mac, option_82)
 
     sys.exit(0)