analyzer.py

"""
Module is used for analyzing link relationships
"""
import requests
from requests.exceptions import HTTPError

from bs4 import BeautifulSoup
from ete3 import Tree, TreeStyle, TextFace, add_face_to_node
from .getweblinks import get_urls_from_page
from .pagereader import read

class LinkTree:
    """
    This is a class that represents a tree of links within TorBot. This can be used to build a tree,
    examine the number of nodes, if a node exists within a tree, displaying the tree and
    downloading the tree, and will be expanded in the future to meet further needs.

    Attributes:
        root (str): Represents root link
        tld (bool): Decides whether or not to use additional top-level-domains besides .tor
        stop_depth (int): Depth of which to stop searching for links
    """
    def __init__(self, root, tld=False, stop_depth=1):
        self._tree = build_tree(root, tld=tld, stop=stop_depth)

    def __len__(self):
        return len(self._tree)

    def __contains__(self, link):
        return link in self._tree

    def save(self, file_name):
        """
        Saves LinkTree to file with given file_name
        Current file types supported are .png, .pdf, .svg

        Args:
            file_name (str): Name of file being saved to
        """
        style = TreeStyle()
        style.show_leaf_name = False
        def my_layout(node):
            node_style = TextFace(node.name, tight_text=True)
            add_face_to_node(node_style, node, column=0, position='branch-bottom')
        style.layout_fn = my_layout
        self._tree.render(file_name, tree_style=style)

    def show(self):
        """
        Allows user to quickly view LinkTree
        """
        style = TreeStyle()
        style.show_leaf_name = False
        def my_layout(node):
            node_style = TextFace(node.name, tight_text=True)
            add_face_to_node(node_style, node, column=0, position='branch-bottom')
        style.layout_fn = my_layout
        self._tree.show(tree_style=style)

def get_node_children(link, tld):
    """
    Returns children for link node

    Args:
        link (str): link node to get children for
        tld (bool): Additional top-level-domains
    Returns:
        children (list): A list of children from linknode
    """
    try:
        resp = requests.get(link)
        soup = BeautifulSoup(resp.text, 'html.parser')
        children = get_urls_from_page(soup, tld)
    except (HTTPError, ConnectionError):
        children = []
    return children

def initialize_tree(link, tld):
    """
    Creates root of tree
    Args:
        link (str): link node to be used as root
        tld (bool): Additional top-level-domains
    Returns:
        root (ete3.Tree): root node of tree
        to_visit (list): Children of root node
    """
    root = Tree(name=link)
    html_content = read(link)
    soup = BeautifulSoup(html_content, 'html.parser')
    to_visit = get_urls_from_page(soup, extension=tld)
    return root, to_visit

def build_tree(link, tld, stop=1, *, rec=0, to_visit=None, tree=None):
    """
    Builds tree using Breadth First Search. You can specify stop depth.
    Rec & tree arguments are used for recursion.

    *NOTE: This function uses a GET request for each url found, this can
    be very expensive so avoid if possible try to acquire the urls to
    be traversed and use bfs function.

    Args:
        link (str): root node
        tld (boolean): specifies if all top-level-domains will be allowed or not
        stop (int): stops traversing at this depth if specified
        rec (int): used for recursion
        tree (ete3.Tree): a tree node used for recursion

    Returns:
        tree (ete3.Tree): built tree
    """
    if rec == 0:
        tree, to_visit = initialize_tree(link, tld)

    sub_tree = Tree(name=tree.name)

    if rec == stop:
        # If recursion is 0 then sub_tree will be root
        return sub_tree if rec == 0 else tree

    children_to_visit = list()
    for link_name in to_visit:
        link_node = sub_tree.add_child(name=link_name)
        link_children = get_node_children(link_name, tld)
        # No need to find children if we aren't going to visit them
        if stop != rec + 1:
            for child in link_children:
                link_node.add_child(name=child)
                children_to_visit.append(child)
    rec += 1

    # If we've reached stop depth then return tree
    if stop == rec:
        return sub_tree

    new_tree = tree.add_child(sub_tree)
    return build_tree(to_visit, tld, stop, rec=rec, tree=new_tree)