Add support for "recursive" keyword to traffic-shape a matched proces…

…s' descendants

README.md

@@ -122,6 +122,19 @@ processes:
     match:
       - exe: /opt/discord/Discord
 
+  Riot:
+    download-priority: 2
+    upload-priority: 2
+
+    # The process that actually creates network traffic for electron-based applications
+    # is not uniquely identifiable. Instead we match a uniquely identifiable parent
+    # process, in this case "riot-desktop", and set recursive to True. This instructs
+    # TrafficToll to traffic shape the connections of the matched process and all it's
+    # descendants
+    recursive: True
+    match:
+      - name: riot-desktop
+
   JDownloader 2:
     download: 300kbps
     # The download-priority and upload-priority if omitted while another process

example.yaml

@@ -94,6 +94,19 @@ processes:
     match:
       - exe: /opt/discord/Discord
 
+  Riot:
+    download-priority: 2
+    upload-priority: 2
+
+    # The process that actually creates network traffic for electron-based applications
+    # is not uniquely identifiable. Instead we match a uniquely identifiable parent
+    # process, in this case "riot-desktop", and set recursive to True. This instructs
+    # TrafficToll to traffic shape the connections of the matched process and all it's
+    # descendants
+    recursive: True
+    match:
+      - name: riot-desktop
+
   JDownloader 2:
     download: 300kbps
     # The download-priority and upload-priority if omitted while another process

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "TrafficToll"
-version = "1.1.0"
+version = "1.2.0"
 description = "NetLimiter-like bandwidth limiting and QoS for Linux"
 authors = ["cryzed <[email protected]>"]
 

traffictoll/__main__.py

@@ -12,14 +12,15 @@ def main() -> None:
     logger.stop(0)
     logger.add(sys.stderr, level=arguments.logging_level)
 
+    # noinspection PyBroadException
     try:
         cli_main(arguments)
     except KeyboardInterrupt:
         logger.info("Aborted")
     except DependencyError as error:
         logger.error("Missing dependency: {}", error)
-    except Exception as error:
-        logger.error("Unexpected error occurred: {}", error)
+    except Exception:
+        logger.exception("Unexpected error occurred:")
 
 
 if __name__ == "__main__":

traffictoll/cli.py

@@ -128,8 +128,11 @@ def main(arguments: argparse.Namespace) -> None:
             logger.warning(
                 "No conditions for: {!r} specified, it will never be matched", name
             )
+            continue
 
-        predicate = ProcessFilterPredicate(name, conditions)
+        predicate = ProcessFilterPredicate(
+            name, conditions, process.get("recursive", False)
+        )
         process_filter_predicates.append(predicate)
 
         # Set up classes for download/upload limiting

traffictoll/net.py

@@ -1,50 +1,57 @@
 import collections
 import itertools
 import re
-from typing import List, DefaultDict, Iterable
+from typing import DefaultDict, Iterable, Set
 
 import psutil
 from loguru import logger
 
+# noinspection PyProtectedMember
+from psutil._common import pconn
+
 ProcessFilterPredicate = collections.namedtuple(
-    "ProcessFilterPredicate", ["name", "conditions"]
+    "ProcessFilterPredicate", ["name", "conditions", "recursive"]
 )
 
 
-def _match_process(process: psutil.Process, predicate: str) -> bool:
-    name, regex = predicate
-    value = getattr(process, name)()
-    if isinstance(value, int):
-        value = str(value)
-    elif isinstance(value, (list, tuple)):
-        value = " ".join(value)
+def _match_process(process: psutil.Process, predicate: ProcessFilterPredicate) -> bool:
+    for condition in predicate.conditions:
+        name, regex = condition
+        value = getattr(process, name)()
+        if isinstance(value, int):
+            value = str(value)
+        elif isinstance(value, (list, tuple)):
+            value = " ".join(value)
+
+        if not re.match(regex, value):
+            return False
 
-    return bool(re.match(regex, value))
+    return True
 
 
 def filter_net_connections(
     predicates: Iterable[ProcessFilterPredicate],
-) -> DefaultDict[str, List[psutil._common.pconn]]:
-    filtered: DefaultDict[str, List[psutil._common.pconn]] = collections.defaultdict(
-        list
-    )
-    connections = psutil.net_connections()
-    for connection, predicate in itertools.product(connections, predicates):
-        # Stop no specified conditions from matching every process
-        if not (predicate.conditions and connection.pid):
-            continue
-
+) -> DefaultDict[str, Set[pconn]]:
+    filtered: DefaultDict[str, Set[pconn]] = collections.defaultdict(set)
+    for process, predicate in itertools.product(psutil.process_iter(), predicates):
         try:
-            process = psutil.Process(connection.pid)
-            if all(
-                _match_process(process, condition) for condition in predicate.conditions
-            ):
-                filtered[predicate.name].append(connection)
+            if not _match_process(process, predicate):
+                continue
+
+            connections = filtered[predicate.name]
+            connections.update(process.connections())
         except psutil.NoSuchProcess:
-            logger.warning(
+            logger.debug(
                 "Process with PID {} died while filtering network connections",
-                connection.pid,
+                process.pid,
             )
             continue
 
+        if predicate.recursive:
+            for child in process.children(recursive=True):
+                try:
+                    connections.update(child.connections())
+                except psutil.NoSuchProcess:
+                    pass
+
     return filtered