From 7daed1b0d880a82e42f596866486aff2bef9f0ef Mon Sep 17 00:00:00 2001
From: Joshua Boniface <joshua@boniface.me>
Date: Tue, 25 Sep 2018 02:20:32 -0400
Subject: [PATCH] Improve limit handling

---
 client-common/client_lib/network.py | 18 ++++++++++--
 client-common/client_lib/node.py    | 13 +++++++--
 client-common/client_lib/router.py  | 14 +++++++---
 client-common/client_lib/vm.py      | 43 +++++++++++++++++++++--------
 4 files changed, 66 insertions(+), 22 deletions(-)

diff --git a/client-common/client_lib/network.py b/client-common/client_lib/network.py
index c665d9c9..01162999 100644
--- a/client-common/client_lib/network.py
+++ b/client-common/client_lib/network.py
@@ -217,15 +217,27 @@ def get_info(zk_conn, network, long_output):
 def get_list(zk_conn,  limit):
     net_list = []
     full_net_list = zk_conn.get_children('/networks')
+
     for net in full_net_list:
         description = zkhandler.readdata(zk_conn, '/networks/{}'.format(net))
         if limit != None:
             try:
-                if not re.match(limit, net) or not re.match(limit, description):
-                    continue
+                # Implcitly assume fuzzy limits
+                if re.match('\^.*', limit) == None:
+                    limit = '.*' + limit
+                if re.match('.*\$', limit) == None:
+                    limit = limit + '.*'
+
+                if re.match(limit, net) != None:
+                    net_list.append(net)
+                if re.match(limit, description) != None:
+                    net_list.append(net)
             except Exception as e:
                 return False, 'Regex Error: {}'.format(e)
-        net_list.append(net)
+        else:
+            net_list.append(net)
+
+    print(net_list)
 
     net_list_output = []
     description = {}
diff --git a/client-common/client_lib/node.py b/client-common/client_lib/node.py
index 35d792f8..729efe49 100644
--- a/client-common/client_lib/node.py
+++ b/client-common/client_lib/node.py
@@ -169,11 +169,18 @@ def get_list(zk_conn, limit):
     for node in full_node_list:
         if limit != None:
             try:
-                if re.match(limit, node) == None:
-                    continue
+                # Implcitly assume fuzzy limits
+                if re.match('\^.*', limit) == None:
+                    limit = '.*' + limit
+                if re.match('.*\$', limit) == None:
+                    limit = limit + '.*'
+
+                if re.match(limit, node) != None:
+                    node_list.append(node)
             except Exception as e:
                 return False, 'Regex Error: {}'.format(e)
-        node_list.append(node)
+        else:
+            node_list.append(node)
 
     node_list_output = []
     node_daemon_state = {}
diff --git a/client-common/client_lib/router.py b/client-common/client_lib/router.py
index 34425fb9..2e0c90df 100644
--- a/client-common/client_lib/router.py
+++ b/client-common/client_lib/router.py
@@ -121,12 +121,18 @@ def get_list(zk_conn, limit):
     for router in full_router_list:
         if limit != None:
             try:
-                if re.match(limit, router) == None:
-                    continue
+                # Implcitly assume fuzzy limits
+                if re.match('\^.*', limit) == None:
+                    limit = '.*' + limit
+                if re.match('.*\$', limit) == None:
+                    limit = limit + '.*'
+
+                if re.match(limit, router) != None:
+                    router_list.append(router)
             except Exception as e:
-                common.stopZKConnection(zk_conn)
                 return False, 'Regex Error: {}'.format(e)
-        router_list.append(router)
+        else:
+            router_list.append(router)
 
     router_list_output = []
     router_daemon_state = {}
diff --git a/client-common/client_lib/vm.py b/client-common/client_lib/vm.py
index 4f2d52bc..6c695fc3 100644
--- a/client-common/client_lib/vm.py
+++ b/client-common/client_lib/vm.py
@@ -35,6 +35,7 @@ import configparser
 import kazoo.client
 
 import client_lib.ansiiprint as ansiiprint
+import client_lib.zkhandler as zkhandler
 import client_lib.common as common
 
 #
@@ -466,7 +467,7 @@ def get_list(zk_conn, hypervisor, limit):
         # Verify node is valid
         common.verifyNode(zk_conn, hypervisor)
 
-    vm_list_raw = zk_conn.get_children('/domains')
+    full_vm_list = zk_conn.get_children('/domains')
     vm_list = []
     vm_list_output = []
 
@@ -480,23 +481,41 @@ def get_list(zk_conn, hypervisor, limit):
     vm_vcpu = {}
 
     # If we're limited, remove other nodes' VMs
-    for vm in vm_list_raw:
+    for vm in full_vm_list:
+
         # Check we don't match the limit
-        name = zk_conn.get('/domains/{}'.format(vm))[0].decode('ascii')
+        name = zkhandler.readdata(zk_conn, '/domains/{}'.format(vm))
+        vm_hypervisor[vm] = zkhandler.readdata(zk_conn, '/domains/{}/hypervisor'.format(vm))
         if limit != None:
             try:
-                if re.match(limit, name) == None:
-                    continue
+                # Implcitly assume fuzzy limits
+                if re.match('\^.*', limit) == None:
+                    limit = '.*' + limit
+                if re.match('.*\$', limit) == None:
+                    limit = limit + '.*'
+
+                if re.match(limit, vm) != None:
+                    if hypervisor == None:
+                        vm_list.append(vm)
+                    else:
+                        if vm_hypervisor[vm] == hypervisor:
+                            vm_list.append(vm)
+
+                if re.match(limit, name) != None:
+                    if hypervisor == None:
+                        vm_list.append(vm)
+                    else:
+                        if vm_hypervisor[vm] == hypervisor:
+                            vm_list.append(vm)
             except Exception as e:
-                click.echo('Regex Error: {}'.format(e))
-                exit(1)
-        # Check hypervisor to avoid unneeded ZK calls
-        vm_hypervisor[vm] = zk_conn.get('/domains/{}/hypervisor'.format(vm))[0].decode('ascii')
-        if hypervisor == None:
-            vm_list.append(vm)
+                return False, 'Regex Error: {}'.format(e)
         else:
-            if vm_hypervisor[vm] == hypervisor:
+            # Check hypervisor to avoid unneeded ZK calls
+            if hypervisor == None:
                 vm_list.append(vm)
+            else:
+                if vm_hypervisor[vm] == hypervisor:
+                    vm_list.append(vm)
 
     # Gather information for printing
     for vm in vm_list: