Fix visualization

Author: KingAkeem

src/modules/analyzer.py

@@ -3,14 +3,12 @@
 """
 from requests.exceptions import HTTPError
 
-from ete3 import Tree, TreeStyle, TextFace, add_face_to_node
-from .link import LinkNode
-from .utils import multi_thread
+from ete3 import faces, Tree, TreeStyle, TextFace, add_face_to_node
 
 
 def default_layout(node):
     node_style = TextFace(node.name, tight_text=True)
-    add_face_to_node(node_style, node, column=0, position='branch-bottom')
+    faces.add_face_to_node(node_style, node, column=0, position='branch-bottom')
 
 
 default_style = TreeStyle()
@@ -51,6 +49,7 @@ def save(self, file_name, tree_style=default_style):
             file_name (str): Name of file being saved to
             tree_style (TreeStyle): Styling of downloaded tree
         """
+        self._tree.layout_fn = default_layout
         self._tree.render(file_name, tree_style)
 
     def show(self, tree_style=default_style):
@@ -60,7 +59,8 @@ def show(self, tree_style=default_style):
         Args:
             tree_style (TreeStyle): Styling of downloaded tree
         """
-        self._tree.show(tree_style)
+        self._tree.layout_fn = default_layout
+        self._tree.show(tree_style=tree_style)
 
 
 def build_tree(node, stop=1, rec=0):
@@ -76,16 +76,15 @@ def build_tree(node, stop=1, rec=0):
         tree (ete3.Tree): Built tree.
     """
 
+    print('Adding node for: ', node.get_name())
     tree = Tree(name=node.get_name())
 
     if rec == stop:
         return tree
     else:
         rec += 1
 
-    node.load_data()
     for child in node.get_children():
-        child.load_data()
         if child.get_children():
             tree.add_child(build_tree(child, stop, rec))
         else:

src/modules/collect_data.py

@@ -8,7 +8,6 @@
 
 from bs4 import BeautifulSoup
 from dotenv import load_dotenv
-from .utils import multi_thread
 from .utils import find_file
 from threadsafe.safe_csv import SafeDictWriter
 from progress.bar import Bar

src/modules/link.py

@@ -48,11 +48,7 @@ def get_json_data(node):
     json = []
     for anchor_tag in node._node.find_all('a'):
         link = anchor_tag.get('href')
-        title = "Not Available"
-        if validate_link(link):
-            node = LinkNode(link)
-            title = node.get_name() 
-            json.append({"link":link,"title":title})
+        json.append({"link":link,"tag":anchor_tag})
     return json    
 
 
@@ -91,9 +87,6 @@ def __init__(self, link):
         self._link = link
 
     def load_data(self):
-        if self._loaded:
-            return
-
         response = requests.get(self._link)
         status = str(response.status_code)
         try:
@@ -102,40 +95,43 @@ def load_data(self):
             self._node = BeautifulSoup(response.text, 'html.parser')
             self.status = color(status, 'green')
             self._name = self._node.title.string
+            self._emails = get_emails(self)
+            self._children = get_children(self)
+            self._emails = get_emails(self)
+            self._images = get_images(self)
+            self._json_data = get_json_data(self)
         except Exception:
             self._node = None
             self.status = color(status, 'yellow')
             self._name = 'TITLE NOT FOUND'
+        finally:
+            self._loaded = True
 
-        self._emails = get_emails(self)
-        self._children = get_children(self)
-        self._emails = get_emails(self)
-        self._images = get_images(self)
-        self._json_data = get_json_data(self)
-        self._loaded = True
 
     def get_link(self):
         return self._link
 
     def get_name(self):
+        if not self._loaded:
+            self.load_data()
         return self._name
 
     def get_children(self):
         if not self._loaded:
-            raise Exception("node is not loaded")
+            self.load_data()
         return self._children
 
     def get_emails(self):
         if not self._loaded:
-            raise Exception("node is not loaded")
+            self.load_data()
         return self._emails 
 
     def get_json(self):
         if not self._loaded:
-            raise Exception("node is not loaded")
+            self.load_data()
         return self._json_data
 
     def get_meatadta(self):
         if not self._loaded:
-            raise Exception("node is not loaded")
+            self.load_data()
         return self._metadata

src/modules/link_io.py

@@ -28,7 +28,6 @@ def display_children(node):
     Args:
         node (LinkNode): root of children to be displayed
     """
-    node.load_data()
     children = node.get_children()
     sucess_msg = color(f'Links Found - {len(children)}', 'green')
     print(sucess_msg + '\n' + '---------------------------------')
@@ -43,7 +42,6 @@ def display(node):
         link (str): link to get status of
     """
     try:
-        node.load_data()
         title = node.get_name()
         status = node.status
     except Exception:

src/torBot.py

@@ -150,7 +150,6 @@ def main():
         print_tor_ip_address()
         # -m/--mail
         if args.mail:
-            node.load_data()
             emails = node.get_emails()
             print(emails)
             if args.save: