Generating Tree with Python for Rumal

demo.py

import pymongo
import simplejson
from bson import ObjectId
from urlparse import urlparse

db = pymongo.MongoClient().thug

def graph_populate_node(analysis_id, new_node):
    '''Adds information from different 
       collections to nodes.'''
    url_id = new_node["url_id"]
    if not isinstance(analysis_id, ObjectId):
        analysis_id = ObjectId(analysis_id)
    url         = db.urls.find_one({"_id": url_id})
    location    = db.locations.find_one({"analysis_id": analysis_id, "url_id": url_id}) or {}
    exploits    = [x for x in db.exploits.find({"analysis_id": analysis_id, "url_id": url_id})]

    for key, value in location.iteritems():
        if isinstance(value, ObjectId):
            location[key] = str(value)

    for exploit in exploits:
        for key, value in exploit.iteritems():
            if isinstance(value, ObjectId):
                exploit[key] = str(value)
    # new_node['url_id']  = str(url_id),
    new_node['url']  = url and url['url'] or '-'
    new_node['domain']  = url and urlparse(url['url']).hostname or '-'
    new_node['location']  = location and location or {}
    new_node['exploits']  = exploits
    new_node['children']  = []

    return new_node

def get_immediate_children(analysis_id,source_id):
    ''' Gets a list of all immediate children of 
        node with given url_id as source_id'''
    return db.connections.find({"analysis_id": ObjectId(analysis_id),"source_id": ObjectId(source_id)})

def generate_tree(analysis_id):
    root_url_id = db.connections.find({"analysis_id": ObjectId(analysis_id)}).sort("chain_id")[0]['source_id']
    # print root_url_id
    # flat_tree_nodes holds the tree in the form of a list
    # of nodes with each node having a node id(nid or index in list)
    # and a parent parameter which refers to parent node 
    flat_tree_nodes = [{"url_id":root_url_id,"parent":None}] #root node as initial element

    # Traverses through flat_tree_nodes 
    # initially consisting only of root
    # For each node extra details are 
    # added and its children with reference 
    # to parent are added.
    for nid,new_node in enumerate(flat_tree_nodes):
        # new_node["nid"] = nid
        node = graph_populate_node(analysis_id, new_node)
        node['nid'] = nid
        flat_tree_nodes[nid] = node
        # print node["url_id"]
        # print get_immediate_children(analysis_id,new_node["url_id"])
        for x in db.connections.find({"analysis_id": ObjectId(analysis_id),"source_id": ObjectId(node["url_id"])}):
            # print x
            temp = {
            "url_id" : x["destination_id"],
            "parent" : nid
            }
            flat_tree_nodes.append(temp)

    # Now traverse(in reverse)list with nodes that have
    # all information except children. Appends children
    # to parent. The initial element has all nodes appended
    # and gives the required tree.
    for node in reversed(flat_tree_nodes):
        if not len(node["children"]):
            node.pop("children")
        if node["parent"] != None:
            flat_tree_nodes[node["parent"]]["children"].append(node)

    final_tree = flat_tree_nodes[0]
    return final_tree

print generate_tree("54d92dfc3295931ca96354f4")