Persisting pids given by parallel ssh connections and its processes

diff --git a/src/nepi/resources/linux/node.py b/src/nepi/resources/linux/node.py
index ab10bae..16b9b00 100644 (file)
--- a/src/nepi/resources/linux/node.py
+++ b/src/nepi/resources/linux/node.py
@@ -440,6 +440,42 @@ class LinuxNode(ResourceManager):
                    "sudo -S killall -u {} || /bin/true ; ".format(self.get("username")))
         else:
             if self.state >= ResourceState.READY:
+                ########################
+                #Collect all process (must change for a more intelligent way)
+                ppid = []
+                pids = []
+                avoid_pids = "ps axjf | awk '{print $1,$2}'"
+                (out, err), proc = self.execute(avoid_pids)
+                if len(out) != 0:
+                    for line in out.strip().split("\n"):
+                        parts = line.strip().split(" ")
+                        ppid.append(parts[0])
+                        pids.append(parts[1])
+
+                #Collect all process below ssh -D
+                tree_owner = 0
+                ssh_pids = []
+                sshs = "ps aux | grep 'sshd' | awk '{print $2,$12}'"
+                (out, err), proc = self.execute(sshs)
+                if len(out) != 0:
+                    for line in out.strip().split("\n"):
+                        parts = line.strip().split(" ")
+                        if parts[1].startswith('root@pts'):
+                            ssh_pids.append(parts[0])
+                        elif parts[1] == "-D":
+                            tree_owner = parts[0]
+
+                avoid_kill = []
+                temp = []
+                #Search for the child process of the pid's collected at the first block.
+                for process in ssh_pids:
+                    temp = self.search_for_child(process, pids, ppid)
+                    avoid_kill = list(set(temp))
+                
+                if len(avoid_kill) > 0:
+                    avoid_kill.append(tree_owner) 
+                ########################
+
                 import pickle
                 pids = pickle.load(open("/tmp/save.proc", "rb"))
                 pids_temp = dict()
@@ -449,9 +485,17 @@ class LinuxNode(ResourceManager):
                     for line in out.strip().split("\n"):
                         parts = line.strip().split(" ")
                         pids_temp[parts[0]] = parts[1]
+                    # creates the difference between the machine pids freezed (pickle) and the actual
+                    # adding the avoided pids filtered above (avoid_kill) to allow users keep process
+                    # alive when using besides ssh connections  
                     kill_pids = set(pids_temp.items()) - set(pids.items())
                     kill_pids = ' '.join(dict(kill_pids).keys())
 
+                    # removing pids from beside connections and its process
+                    kill_pids = kill_pids.split(' ')
+                    kill_pids = list(set(kill_pids) - set(avoid_kill))
+                    kill_pids = ' '.join(kill_pids)
+
                     cmd = ("killall tcpdump || /bin/true ; " +
                            "kill $(ps aux | grep '[.]nepi' | awk '{print $2}') || /bin/true ; " +
                            "kill {} || /bin/true ; ".format(kill_pids))
@@ -464,6 +508,16 @@ class LinuxNode(ResourceManager):
 
         (out, err), proc = self.execute(cmd, retry = 1, with_lock = True)
 
+    def search_for_child(self, pid, pids, ppid, family=[]):
+        """ Recursive function to search for child. List A contains the pids and list B the parents (ppid)
+        """
+        family.append(pid)
+        for key, value in enumerate(ppid):
+            if value == pid:
+                child = pids[key]
+                self.search_for_child(child, pids, ppid)
+        return family
+        
     def clean_home(self):
         """ Cleans all NEPI related folders in the Linux host
         """