Merge pull request #78 from adisbladis/ruleperf

Rule performance fixes
This commit is contained in:
Simone Margaritelli 2017-05-21 14:06:14 +02:00 committed by GitHub
commit eb3f96c303
3 changed files with 92 additions and 71 deletions

View file

@ -16,84 +16,104 @@
# program. If not, go to http://www.gnu.org/licenses/gpl.html
# or write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
import logging
from collections import namedtuple
from threading import Lock
from enum import Enum
import logging
import sqlite3
class Rule:
Rule = namedtuple('Rule', ('app_path',
'verdict',
'address',
'port',
'proto'))
class RuleVerdict(Enum):
ACCEPT = 0
DROP = 1
DROP = 1
class RuleSaveOption(Enum):
ONCE = 0
UNTIL_QUIT = 1
FOREVER = 2
def __init__( self, app_path=None, verdict=ACCEPT, address=None, port=None, proto=None ):
self.app_path = app_path
self.verdict = verdict
self.address = address
self.port = port
self.proto = proto
def matches( self, c ):
if self.app_path != c.app_path:
return False
def matches(rule, conn):
if rule.app_path != conn.app_path:
return False
elif self.address is not None and self.address != c.dst_addr:
return False
elif rule.address is not None and rule.address != conn.dst_addr:
return False
elif self.port is not None and self.port != c.dst_port:
return False
elif rule.port is not None and rule.port != conn.dst_port:
return False
elif self.proto is not None and self.proto != c.proto:
return False
elif rule.proto is not None and rule.proto != conn.proto:
return False
else:
return True
else:
return True
class Rules:
def __init__(self, database):
self.mutex = Lock()
self.db = RulesDB(database)
self.rules = self.db.load_rules()
db = self.db = RulesDB(database)
self.rules = {}
def get_verdict( self, connection ):
with db._lock:
for r in db._load_rules():
self._add_rule(r)
def get_verdict(self, connection):
with self.mutex:
for r in self.rules:
if r.matches(connection):
for r in self.rules.get(connection.app_path, []):
if matches(r, connection):
return r.verdict
return None
def _remove_rules_for_path( self, path, remove_from_db=False ):
for rule in self.rules:
if rule.app_path == path:
self.rules.remove(rule)
def _remove_rules_for_path(self, path, remove_from_db=False):
try:
del self.rules[path]
except KeyError:
pass
if remove_from_db is True:
self.db.remove_all_app_rules(path)
def add_rule( self, connection, verdict, apply_to_all=False, save_option=Rule.UNTIL_QUIT ):
def _add_rule(self, rule):
self.rules.setdefault(rule.app_path, set()).add(rule)
def add_rule(self, connection, verdict, apply_to_all=False,
save_option=RuleSaveOption.UNTIL_QUIT.value):
with self.mutex:
logging.debug( "Adding %s rule for '%s' (all=%s)" % (
"ALLOW" if verdict == Rule.ACCEPT else "DENY",
connection,
"true" if apply_to_all == True else "false" ) )
r = Rule()
r.verdict = verdict
r.app_path = connection.app_path
logging.debug("Adding %s rule for '%s' (all=%s)",
"ALLOW" if RuleVerdict(verdict) == RuleVerdict.ACCEPT else "DENY", # noqa
connection,
"true" if apply_to_all is True else "false")
if apply_to_all is True:
self._remove_rules_for_path( r.app_path, (save_option == Rule.FOREVER) )
self._remove_rules_for_path(
connection.app_path,
(RuleSaveOption(save_option) == RuleSaveOption.FOREVER))
elif apply_to_all is False:
r.address = connection.dst_addr
r.port = connection.dst_port
r.proto = connection.proto
r = Rule(
connection.app_path,
verdict,
connection.dst_addr if not apply_to_all else None,
connection.dst_port if not apply_to_all else None,
connection.proto if not apply_to_all else None)
self.rules.append(r)
self._add_rule(r)
if save_option == Rule.FOREVER:
if RuleSaveOption(save_option) == RuleSaveOption.FOREVER:
self.db.save_rule(r)
@ -115,18 +135,18 @@ class RulesDB:
c = conn.cursor()
c.execute("CREATE TABLE IF NOT EXISTS rules (app_path TEXT, verdict INTEGER, address TEXT, port INTEGER, proto TEXT, UNIQUE (app_path, verdict, address, port, proto))") # noqa
def load_rules(self):
with self._lock:
conn = self._get_conn()
c = conn.cursor()
c.execute("SELECT * FROM rules")
return [Rule(*item) for item in c.fetchall()]
def _load_rules(self):
conn = self._get_conn()
c = conn.cursor()
c.execute("SELECT * FROM rules")
for item in c.fetchall():
yield Rule(*item)
def save_rule(self, rule):
with self._lock:
conn = self._get_conn()
c = conn.cursor()
c.execute("INSERT INTO rules VALUES (?, ?, ?, ?, ?)", (rule.app_path, rule.verdict, rule.address, rule.port, rule.proto,)) # noqa
c.execute("INSERT INTO rules VALUES (?, ?, ?, ?, ?)", (rule.app_path, rule.verdict.value, rule.address, rule.port, rule.proto,)) # noqa
conn.commit()
def remove_all_app_rules(self, app_path):

View file

@ -28,7 +28,7 @@ import os
from opensnitch.ui import QtApp
from opensnitch.connection import Connection
from opensnitch.dns import DNSCollector
from opensnitch.rule import Rule, Rules
from opensnitch.rule import RuleVerdict, Rules
from opensnitch.procmon import ProcMon
@ -101,7 +101,7 @@ class PacketHandler(threading.Thread):
self.pkt.accept()
else:
if verdict == Rule.DROP:
if RuleVerdict(verdict) == RuleVerdict.DROP:
drop_packet(self.pkt, self.conn)
else:
@ -141,18 +141,18 @@ class Snitch:
# Get verdict, if verdict cannot be found prompt user in thread
verd = self.rules.get_verdict(conn)
if verd == Rule.DROP:
drop_packet(pkt, conn)
elif verd == Rule.ACCEPT:
pkt.accept()
elif verd is None:
if verd is None:
conn.hostname = self.dns.get_hostname(conn.dst_addr)
handler = PacketHandler(conn, pkt, self.rules)
self.connection_futures[conn.id] = handler.future
self.qt_app.prompt_user(conn)
elif RuleVerdict(verd) == RuleVerdict.DROP:
drop_packet(pkt, conn)
elif RuleVerdict(verd) == RuleVerdict.ACCEPT:
pkt.accept()
else:
raise RuntimeError("Unhandled state")

View file

@ -16,8 +16,8 @@
# program. If not, go to http://www.gnu.org/licenses/gpl.html
# or write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
from opensnitch.rule import RuleVerdict, RuleSaveOption
from PyQt5 import QtCore, QtGui, uic, QtWidgets
from opensnitch.rule import Rule
import threading
import queue
import sys
@ -33,7 +33,7 @@ DIALOG_UI_PATH = "%s/dialog.ui" % RESOURCES_PATH
class Dialog(QtWidgets.QDialog, uic.loadUiType(DIALOG_UI_PATH)[0]):
DEFAULT_RESULT = (Rule.ONCE, Rule.ACCEPT, False)
DEFAULT_RESULT = (RuleSaveOption.ONCE, RuleVerdict.ACCEPT, False)
MESSAGE_TEMPLATE = "<b>%s</b> (pid=%s) wants to connect to <b>%s</b> on <b>%s port %s%s</b>" # noqa
add_connection_signal = QtCore.pyqtSignal()
@ -71,7 +71,8 @@ class Dialog(QtWidgets.QDialog, uic.loadUiType(DIALOG_UI_PATH)[0]):
with self.rule_lock:
verd = self.rules.get_verdict(connection)
if verd is not None:
self.set_conn_result(connection, Rule.ONCE, verd, False)
self.set_conn_result(connection, RuleSaveOption.ONCE,
verd, False)
# Lock needs to be released before callback can be triggered
if verd is not None:
@ -149,16 +150,16 @@ class Dialog(QtWidgets.QDialog, uic.loadUiType(DIALOG_UI_PATH)[0]):
self.block_button.hide()
def _allow_action(self):
self._action(Rule.ACCEPT, False)
self._action(RuleVerdict.ACCEPT, False)
def _deny_action(self):
self._action(Rule.DROP, False)
self._action(RuleVerdict.DROP, False)
def _whitelist_action(self):
self._action(Rule.ACCEPT, True)
self._action(RuleVerdict.ACCEPT, True)
def _block_action(self):
self._action(Rule.DROP, True)
self._action(RuleVerdict.DROP, True)
def set_conn_result(self, connection, option, verdict, apply_to_all):
try:
@ -173,11 +174,11 @@ class Dialog(QtWidgets.QDialog, uic.loadUiType(DIALOG_UI_PATH)[0]):
s_option = self.action_combo_box.currentText()
if s_option == "Once":
option = Rule.ONCE
option = RuleSaveOption.ONCE
elif s_option == "Until Quit":
option = Rule.UNTIL_QUIT
option = RuleSaveOption.UNTIL_QUIT
elif s_option == "Forever":
option = Rule.FOREVER
option = RuleSaveOption.FOREVER
self.set_conn_result(self.connection, option,
verdict, apply_to_all)
@ -185,7 +186,7 @@ class Dialog(QtWidgets.QDialog, uic.loadUiType(DIALOG_UI_PATH)[0]):
# We need to freeze UI thread while storing rule, otherwise another
# connection that would have been affected by the rule will pop up
# TODO: Figure out how to do this nicely when separating UI
if option != Rule.ONCE:
if option != RuleSaveOption.ONCE:
self.rules.add_rule(self.connection, verdict,
apply_to_all, option)