Fix traceroute

monkey/infection_monkey/network/tools.py

@@ -5,12 +5,11 @@ import select
 import socket
 import struct
 import time
-
+import re
-from six import text_type
-import ipaddress
 
 DEFAULT_TIMEOUT = 10
 BANNER_READ = 1024
+IP_ADDR_RE = r'[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}'
 
 LOG = logging.getLogger(__name__)
 SLEEP_BETWEEN_POLL = 0.5
@@ -176,13 +175,21 @@ def tcp_port_to_service(port):
 
 def traceroute(target_ip, ttl):
     """
-    Traceroute for a specific IP.
+    Traceroute for a specific IP/name.
-    :param target_ip: Destination
+    :param target_ip: IP/name of target
     :param ttl: Max TTL
     :return: Sequence of IPs in the way
     """
     if sys.platform == "win32":
-        try:
+        return _traceroute_windows(target_ip, ttl)
+    else:  # linux based hopefully
+        return _traceroute_linux(target_ip, ttl)
+
+
+def _traceroute_windows(target_ip, ttl):
+    """
+    Traceroute for a specific IP/name - Windows implementation
+    """
     # we'll just use tracert because that's always there
     cli = ["tracert",
            "-d",
@@ -191,20 +198,33 @@ def traceroute(target_ip, ttl):
            target_ip]
     proc_obj = subprocess.Popen(cli, stdout=subprocess.PIPE)
     stdout, stderr = proc_obj.communicate()
-            ip_lines = stdout.split('\r\n')[3:-3]
+    ip_lines = stdout.split('\r\n')
     trace_list = []
-            for line in ip_lines:
+
-                tokens = line.split()
+    first_line_index = None
-                last_token = tokens[-1]
+    for i in range(len(ip_lines)):
-                try:
+        if re.search(r'^\s*1', ip_lines[i]) is not None:
-                    ip_addr = ipaddress.ip_address(text_type(last_token))
+            first_line_index = i
-                except ValueError:
+            break
-                    ip_addr = ""
+
+    for i in range(first_line_index, first_line_index + ttl):
+        if re.search(r'^\s*' + str(i - first_line_index + 1), ip_lines[i]) is None:  # If trace is finished
+            break
+
+        re_res = re.search(IP_ADDR_RE, ip_lines[i])
+        if re_res is None:
+            ip_addr = None
+        else:
+            ip_addr = re_res.group()
         trace_list.append(ip_addr)
+
     return trace_list
-        except:
+
-            return []
+
-    else:  # linux based hopefully
+def _traceroute_linux(target_ip, ttl):
+    """
+    Traceroute for a specific IP/name - Linux implementation
+    """
     # implementation note: We're currently going to just use ping.
     # reason is, implementing a non root requiring user is complicated (see traceroute(8) code)
     # while this is just ugly
@@ -212,7 +232,6 @@ def traceroute(target_ip, ttl):
     current_ttl = 1
     trace_list = []
     while current_ttl <= ttl:
-            try:
         cli = ["ping",
                "-c", "1",
                "-w", "1",
@@ -220,12 +239,18 @@ def traceroute(target_ip, ttl):
                target_ip]
         proc_obj = subprocess.Popen(cli, stdout=subprocess.PIPE)
         stdout, stderr = proc_obj.communicate()
-                ip_line = stdout.split('\n')
+        ips = re.findall(IP_ADDR_RE, stdout)
-                ip_line = ip_line[1]
+        if len(ips) < 2:
-                ip = ip_line.split()[1]
+            raise Exception("Unexpected output")
-                trace_list.append(ipaddress.ip_address(text_type(ip)))
+        elif ips[-1] in trace_list:  # Failed getting this hop
-            except (IndexError, ValueError):
+            trace_list.append(None)
-                # assume we failed parsing output
+        else:
-                trace_list.append("")
+            trace_list.append(ips[-1])
+            dest_ip = ips[0]  # first ip is dest ip. must be parsed here since it can change between pings
+
+            if dest_ip == ips[-1]:
+                break
+
         current_ttl += 1
+
     return trace_list