monkey/chaos_monkey/network/firewall.py

178 lines
5.1 KiB
Python

import subprocess
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
def add_firewall_rule(self, **kwargs):
return False
def remove_firewall_rule(self, **kwargs):
return False
def listen_allowed(self, **kwargs):
return True
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self.close()
def close(self):
return
class WinAdvFirewall(FirewallApp):
def __init__(self):
self._rules = {}
def is_enabled(self):
try:
cmd = subprocess.Popen('netsh advfirewall show currentprofile', stdout=subprocess.PIPE)
out = cmd.stdout.readlines()
for l in out:
if l.startswith('State'):
state = l.split()[-1].strip()
return state == "ON"
except:
return None
def add_firewall_rule(self, name="Firewall", dir="in", action="allow", program=sys.executable, **kwargs):
netsh_args = {'name': name,
'dir': dir,
'action': action,
'program': program}
netsh_args.update(kwargs)
try:
if _run_netsh_cmd('advfirewall firewall add rule', netsh_args):
self._rules[name] = netsh_args
return True
else:
return False
except:
return None
def remove_firewall_rule(self, name="Firewall", **kwargs):
netsh_args = {'name': name}
netsh_args.update(kwargs)
try:
if _run_netsh_cmd('advfirewall firewall delete rule', netsh_args):
if name in self._rules:
del self._rules[name]
return True
else:
return False
except:
return None
def listen_allowed(self, **kwargs):
if not self.is_enabled():
return True
for rule in self._rules.values():
if rule.get('program') == sys.executable and \
'in' == rule.get('dir') and \
'allow' == rule.get('action') and \
4 == len(rule.keys()):
return True
return False
def close(self):
try:
for rule in self._rules.keys():
self.remove_firewall_rule({'name': rule})
except:
pass
class WinFirewall(FirewallApp):
def __init__(self):
self._rules = {}
def is_enabled(self):
try:
cmd = subprocess.Popen('netsh firewall show state', stdout=subprocess.PIPE)
out = cmd.stdout.readlines()
for l in out:
if l.startswith('Operational mode'):
state = l.split('=')[-1].strip()
elif l.startswith('The service has not been started.'):
return False
return state == "Enable"
except:
return None
def add_firewall_rule(self, rule='allowedprogram', name="Firewall", mode="ENABLE", program=sys.executable,
**kwargs):
netsh_args = {'name': name,
'mode': mode,
'program': program}
netsh_args.update(kwargs)
try:
if _run_netsh_cmd('firewall add %s' % rule, netsh_args):
netsh_args['rule'] = rule
self._rules[name] = netsh_args
return True
else:
return False
except:
return None
def remove_firewall_rule(self, rule='allowedprogram', name="Firewall", mode="ENABLE", program=sys.executable,
**kwargs):
netsh_args = {'program': program}
netsh_args.update(kwargs)
try:
if _run_netsh_cmd('firewall delete %s' % rule, netsh_args):
if name in self._rules:
del self._rules[name]
return True
else:
return False
except:
return None
def listen_allowed(self, **kwargs):
if not self.is_enabled():
return True
for rule in self._rules.values():
if rule.get('program') == sys.executable and 'ENABLE' == rule.get('mode'):
return True
return False
def close(self):
try:
for rule in self._rules.values():
self.remove_firewall_rule(**rule)
except:
pass
if sys.platform == "win32":
try:
win_ver = int(platform.version().split('.')[0])
except:
win_ver = 0
if win_ver > 5:
app = WinAdvFirewall()
else:
app = WinFirewall()
else:
app = FirewallApp()