From bdde8dfeedece34ec21495a1de492630e2ec0752 Mon Sep 17 00:00:00 2001
From: daniel goldberg <danielg@guardicore.com>
Date: Mon, 29 Aug 2016 12:08:42 +0300
Subject: [PATCH] Added func to find route to victim, reorganised firewall

---
 chaos_monkey/network/firewall.py | 13 ++++++-------
 chaos_monkey/network/info.py     | 15 +++++++++++++++
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/chaos_monkey/network/firewall.py b/chaos_monkey/network/firewall.py
index 54c8a1870..678b6f67f 100644
--- a/chaos_monkey/network/firewall.py
+++ b/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
 
diff --git a/chaos_monkey/network/info.py b/chaos_monkey/network/info.py
index 9d7f2861e..d5e507846 100644
--- a/chaos_monkey/network/info.py
+++ b/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
\ No newline at end of file