Merge pull request #78 from adisbladis/ruleperf

Rule performance fixes

opensnitch/rule.py

@@ -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 = {}
+
+        with db._lock:
+            for r in db._load_rules():
+                self._add_rule(r)
 
-    def get_verdict( self, connection ):
+    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 @@ def _create_table(self):
             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):

opensnitch/snitch.py

@@ -28,7 +28,7 @@
 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 @@ def run(self):
             self.pkt.accept()
 
         else:
-            if verdict == Rule.DROP:
+            if RuleVerdict(verdict) == RuleVerdict.DROP:
                 drop_packet(self.pkt, self.conn)
 
             else:
@@ -141,18 +141,18 @@ def pkt_callback(self, pkt):
 
             # 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")
 

opensnitch/ui/dialog.py

@@ -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 @@
 
 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 @@ def handle_connection(self):
         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 @@ def _action_changed(self):
             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,19 +174,19 @@ def _action(self, verdict, apply_to_all=False):
             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)
 
             # 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)