From 993f46be91db80f0b9fdac9ccee1bfc6abfe9d95 Mon Sep 17 00:00:00 2001
From: adisbladis <adis@blad.is>
Date: Tue, 20 Jun 2017 20:27:06 +0800
Subject: [PATCH] Improve getting pid for udp connections

---
 opensnitch/proc.py   | 20 ++++++++++++++------
 opensnitch/snitch.py |  6 +-----
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/opensnitch/proc.py b/opensnitch/proc.py
index 121bc9de..ebda5ca8 100644
--- a/opensnitch/proc.py
+++ b/opensnitch/proc.py
@@ -22,14 +22,22 @@ import os
 
 
 def get_pid_by_connection(src_addr, src_p, dst_addr, dst_p, proto='tcp'):
-    pids = (connection.pid for connection in psutil.net_connections(kind=proto)
-            if connection.laddr == (src_addr, int(src_p)) and
-            connection.raddr == (dst_addr, int(dst_p)))
-
     # We always take the first element as we assume it contains only one
     # It should not be possible to keep two connections which are the same.
-    for p in pids:
-        return p
+    for conn in psutil.net_connections(kind=proto):
+        if proto == 'tcp':
+            if conn.laddr != (src_addr, int(src_p)):
+                continue
+
+            if conn.raddr != (dst_addr, int(dst_p)):
+                continue
+
+        # UDP gives us a very limited dataset to work with
+        elif proto == 'udp':
+            if conn.laddr[1] != int(src_p):
+                continue
+
+        return conn.pid
 
     logging.warning("Could not find process for %s connection %s:%s -> %s:%s",
                     proto,
diff --git a/opensnitch/snitch.py b/opensnitch/snitch.py
index 8ef9ceeb..c3d1fe19 100644
--- a/opensnitch/snitch.py
+++ b/opensnitch/snitch.py
@@ -82,10 +82,6 @@ class NetfilterQueueWrapper(threading.Thread):
                 logging.debug("Could not detect protocol for packet.")
                 return
 
-            elif conn.app.pid is None and conn.proto != 'icmp':
-                logging.debug("Could not detect process for connection.")
-                return
-
             # Get verdict, if verdict cannot be found prompt user in thread
             verd = self.snitch.rules.get_verdict(conn)
             if verd is None:
@@ -94,7 +90,7 @@ class NetfilterQueueWrapper(threading.Thread):
                 self.snitch.dbus_service.prompt(
                     conn.id,
                     conn.hostname,
-                    conn.dst_port,
+                    conn.dst_port or 0,
                     conn.dst_addr,
                     conn.proto,
                     conn.app.pid or 0,