From 0693de8f9009dd70b9b15f5faec474bcd621020d Mon Sep 17 00:00:00 2001
From: Arseniy Kuznetsov <k.arseniy@gmail.com>
Date: Sat, 6 Feb 2021 16:19:03 +0100
Subject: [PATCH] DS refactor, fixes/optimizations

---
 mktxp/cli/output/base_out.py | 29 +++++++++++++++++++----------
 mktxp/cli/output/dhcp_out.py |  3 +++
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/mktxp/cli/output/base_out.py b/mktxp/cli/output/base_out.py
index 7454155..fceba6a 100644
--- a/mktxp/cli/output/base_out.py
+++ b/mktxp/cli/output/base_out.py
@@ -21,21 +21,13 @@ from mktxp.cli.config.config import config_handler
 class BaseOutputProcessor:
     OutputCapsmanEntry = namedtuple('OutputCapsmanEntry', ['dhcp_name', 'dhcp_address', 'mac_address', 'rx_signal', 'interface', 'ssid', 'tx_rate', 'rx_rate', 'uptime'])
     OutputWiFiEntry = namedtuple('OutputWiFiEntry', ['dhcp_name', 'dhcp_address', 'mac_address', 'signal_strength', 'signal_to_noise', 'interface', 'tx_rate', 'rx_rate', 'uptime'])
-    OutputDHCPEntry = namedtuple('OutputDHCPEntry', ['host_name', 'comment', 'address', 'active_address', 'mac_address', 'server', 'expires_after'])
+    OutputDHCPEntry = namedtuple('OutputDHCPEntry', ['host_name', 'server', 'mac_address', 'address', 'active_address', 'expires_after'])
 
     @staticmethod
     def augment_record(router_entry, registration_record, dhcp_lease_records):
         try:
             dhcp_lease_record = next((dhcp_lease_record for dhcp_lease_record in dhcp_lease_records if dhcp_lease_record['mac_address']==registration_record['mac_address']))
-            dhcp_name = dhcp_lease_record.get('host_name')
-            dhcp_comment = dhcp_lease_record.get('comment')
-
-            if dhcp_name and dhcp_comment:
-                dhcp_name = f'{dhcp_name[0:20]} ({dhcp_comment[0:20]})' if not router_entry.config_entry.use_comments_over_names else dhcp_comment
-            elif dhcp_comment:
-                dhcp_name = dhcp_comment
-            else:
-                dhcp_name = dhcp_lease_record.get('mac_address') if not dhcp_name else dhcp_name
+            dhcp_name = BaseOutputProcessor.dhcp_name(router_entry, dhcp_lease_record)
             dhcp_address = dhcp_lease_record.get('address', '')
         except StopIteration:
             dhcp_name = registration_record['mac_address']
@@ -62,6 +54,23 @@ class BaseOutputProcessor:
         if registration_record.get('rx_signal'):
             registration_record['rx_signal'] = BaseOutputProcessor.parse_signal_strength(registration_record['rx_signal'])
 
+    @staticmethod
+    def dhcp_name(router_entry, dhcp_lease_record, drop_comment = False):
+        dhcp_name = dhcp_lease_record.get('host_name')
+        dhcp_comment = dhcp_lease_record.get('comment')
+        
+        if dhcp_name and dhcp_comment:
+            dhcp_name = f'{dhcp_name[0:20]} ({dhcp_comment[0:20]})' if not router_entry.config_entry.use_comments_over_names else dhcp_comment
+        elif dhcp_comment:
+            dhcp_name = dhcp_comment
+        else:
+            dhcp_name = dhcp_lease_record.get('mac_address') if not dhcp_name else dhcp_name        
+
+        if drop_comment:
+            del dhcp_lease_record['comment']
+        
+        return dhcp_name
+
     @staticmethod
     def parse_rates(rate):
         wifi_rates_rgx = config_handler.re_compiled.get('wifi_rates_rgx')
diff --git a/mktxp/cli/output/dhcp_out.py b/mktxp/cli/output/dhcp_out.py
index 8fa0b28..74d2eeb 100644
--- a/mktxp/cli/output/dhcp_out.py
+++ b/mktxp/cli/output/dhcp_out.py
@@ -40,9 +40,12 @@ class DHCPOutput:
         output_table = []
         for key in dhcp_by_server.keys():
             for record in dhcp_by_server[key]:
+                record['host_name'] = BaseOutputProcessor.dhcp_name(router_entry, record, drop_comment = True)
                 output_table.append(BaseOutputProcessor.OutputDHCPEntry(**record))
                 num_records += 1
             output_table.append({})
         print()
         print(tabulate(output_table, headers = "keys",  tablefmt="github"))
         print(tabulate([{0:'DHCP Clients:', 1:num_records}], tablefmt="text"))
+
+