refactored traffic module

Signed-off-by: Christoph Massmann <[email protected]>
ma4nn · Nov 17, 2020 · 3059efc · 3059efc
1 parent 4b51389
commit 3059efc
Showing 1 changed file with 48 additions and 53 deletions.
diff --git a/fritzbox_traffic.py b/fritzbox_traffic.py
@@ -24,69 +24,64 @@
 from fritzconnection.lib.fritzstatus import FritzStatus
 from FritzboxConfig import FritzboxConfig
 
-def print_traffic():
-  config = FritzboxConfig()
+class FritzboxTraffic:
+  def __init__(self):
+    config = FritzboxConfig()
+    try:
+      self.__connection = FritzStatus(address=config.server, password=config.password, use_tls=config.useTls)
+    except Exception as e:
+      sys.exit("Couldn't get WAN traffic: " + str(e))
 
-  try:
-    conn = FritzStatus(address=config.server, password=config.password, use_tls=config.useTls)
-  except Exception as e:
-    sys.exit("Couldn't get WAN traffic: " + str(e))
+  def printTraffic(self):
+    traffic = self.__connection.transmission_rate
+    print('down.value %d' % traffic[1])
+    print('up.value %d' % traffic[0])
 
-  traffic = conn.transmission_rate
-  print('down.value %d' % traffic[1])
-  print('up.value %d' % traffic[0])
+    if not os.environ.get('traffic_remove_max') or "false" in os.environ.get('traffic_remove_max'):
+      max_traffic = self.__connection.max_bit_rate
+      print('maxdown.value %d' % max_traffic[1])
+      print('maxup.value %d' % max_traffic[0])
 
-  if not os.environ.get('traffic_remove_max') or "false" in os.environ.get('traffic_remove_max'):
-    max_traffic = conn.max_bit_rate
-    print('maxdown.value %d' % max_traffic[1])
-    print('maxup.value %d' % max_traffic[0])
+  def printConfig(self):
+    max_traffic = self.__connection.max_bit_rate
 
-def print_config():
-  config = FritzboxConfig()
-  try:
-    conn = FritzStatus(address=config.server, password=config.password, use_tls=config.useTls)
-  except Exception as e:
-    print(e)
-    sys.exit("Couldn't get WAN traffic: " + str(e))
-
-  max_traffic = conn.max_bit_rate
-
-  print("graph_title WAN traffic")
-  print("graph_args --base 1000")
-  print("graph_vlabel bit in (-) / out (+) per ${graph_period}")
-  print("graph_category network")
-  print("graph_order down up maxdown maxup")
-  print("down.label received")
-  print("down.type DERIVE")
-  print("down.graph no")
-  print("down.cdef down,8,*")
-  print("down.min 0")
-  print("down.max %d" % max_traffic[1])
-  print("up.label bps")
-  print("up.type DERIVE")
-  print("up.draw LINE")
-  print("up.cdef up,8,*")
-  print("up.min 0")
-  print("up.max %d" % max_traffic[0])
-  print("up.negative down")
-  print("up.info Traffic of the WAN interface.")
-  if not os.environ.get('traffic_remove_max') or "false" in os.environ.get('traffic_remove_max'):
-    print("maxdown.label received")
-    print("maxdown.type GAUGE")
-    print("maxdown.graph no")
-    print("maxup.label MAX")
-    print("maxup.type GAUGE")
-    print("maxup.negative maxdown")
-    print("maxup.draw LINE1")
-    print("maxup.info Maximum speed of the WAN interface.")
+    print("graph_title WAN traffic")
+    print("graph_args --base 1000")
+    print("graph_vlabel bit in (-) / out (+) per ${graph_period}")
+    print("graph_category network")
+    print("graph_order down up maxdown maxup")
+    print("down.label received")
+    print("down.type DERIVE")
+    print("down.graph no")
+    print("down.cdef down,8,*")
+    print("down.min 0")
+    print("down.max %d" % max_traffic[1])
+    print("up.label bps")
+    print("up.type DERIVE")
+    print("up.draw LINE")
+    print("up.cdef up,8,*")
+    print("up.min 0")
+    print("up.max %d" % max_traffic[0])
+    print("up.negative down")
+    print("up.info Traffic of the WAN interface.")
+    if not os.environ.get('traffic_remove_max') or "false" in os.environ.get('traffic_remove_max'):
+      print("maxdown.label received")
+      print("maxdown.type GAUGE")
+      print("maxdown.graph no")
+      print("maxup.label MAX")
+      print("maxup.type GAUGE")
+      print("maxup.negative maxdown")
+      print("maxup.draw LINE1")
+      print("maxup.info Maximum speed of the WAN interface.")
 
 if __name__ == "__main__":
+    traffic = FritzboxTraffic()
     if len(sys.argv) == 2 and sys.argv[1] == 'config':
-        print_config()
+        traffic.printConfig()
     elif len(sys.argv) == 2 and sys.argv[1] == 'autoconf':
         print("yes")  # Some docs say it'll be called with fetch, some say no arg at all
     elif len(sys.argv) == 1 or (len(sys.argv) == 2 and sys.argv[1] == 'fetch'):
         try:
-            print_traffic()
+            traffic.printTraffic()
         except Exception as e:
             sys.exit("Couldn't retrieve fritzbox traffic: " + str(e))