Added func to find route to victim, reorganised firewall

chaos_monkey/network/firewall.py

@@ -3,6 +3,11 @@ import sys
 import platform
 
 
+def _run_netsh_cmd(command, args):
+    cmd = subprocess.Popen("netsh %s %s" % (command, " ".join(['%s="%s"' % (key, value) for key, value in args.items()
+                                                               if value])), stdout=subprocess.PIPE)
+    return cmd.stdout.read().strip().lower().endswith('ok.')
+
 class FirewallApp(object):
     def is_enabled(self, **kwargs):
         return False
@@ -26,12 +31,6 @@ class FirewallApp(object):
         return
 
 
-def _run_netsh_cmd(command, args):
-    cmd = subprocess.Popen("netsh %s %s" % (command, " ".join(['%s="%s"' % (key, value) for key, value in args.items()
-                                                               if value])), stdout=subprocess.PIPE)
-    return cmd.stdout.read().strip().lower().endswith('ok.')
-
-
 class WinAdvFirewall(FirewallApp):
     def __init__(self):
         self._rules = {}
@@ -93,7 +92,7 @@ class WinAdvFirewall(FirewallApp):
     def close(self):
         try:
             for rule in self._rules.keys():
-                _run_netsh_cmd('advfirewall firewall delete rule', {'name': rule})
+                self.remove_firewall_rule({'name': rule})
         except:
             pass
 

chaos_monkey/network/info.py

@@ -5,6 +5,7 @@ import socket
 import struct
 import psutil
 import ipaddress
+from subprocess import check_output
 from random import randint
 
 if sys.platform == "win32":
@@ -112,3 +113,17 @@ def get_ips_from_interfaces():
                 continue
             res.append(str(addr))
     return res
+
+if sys.platform == "win32":
+    def get_ip_for_connection(target_ip):
+        return None
+else:
+    def get_ip_for_connection(target_ip):
+        try:
+            query_str = 'ip route get %s' % target_ip
+            resp = check_output(query_str.split())
+            substr = resp.split()
+            src = substr[substr.index('src')+1]
+            return src
+        except Exception:
+            return None