tap into ss2 run as module mode

cherusk · Nov 24, 2023 · 19107c9 · 19107c9
1 parent 4fe33d4
commit 19107c9
Showing 1 changed file with 2 additions and 15 deletions.
diff --git a/prometheus_ss_exporter/stats.py b/prometheus_ss_exporter/stats.py
@@ -29,24 +29,11 @@ def __init__(self):
         self.args.all = False
         self.args.unix = False
 
-    def _reset_io(self):
-        if sys.version_info[0] == 2:
-            import cStringIO
-            self.stream_sink = cStringIO.StringIO()
-        else:
-            self.stream_sink = io.StringIO()
-
     def provide_tcp_stats(self):
-        _stdout = sys.stdout
-        sys.stdout = self.stream_sink
-
-        ss2.run(self.args)
 
-        # catch stdout
-        sys.stdout = _stdout
-        sk_stats_raw = self.stream_sink.getvalue()
+        ss2.RUN_AS_MODULE = True
 
-        self._reset_io()
+        sk_stats_raw = ss2.run(self.args)
 
         sk_stats_parsed = dict(TCP=dict(flows=list()))
         try: