trust_explorer.py

#!/usr/bin/python3

######
# Author: @sixdub
# Description: Shell based wrapper for networkx library that 
# provides functionality to analyze domain trust relationships 
# generated by powerview
#
######

import networkx as nx
import matplotlib.pyplot as plt
from cmd import Cmd
import sys, os, argparse

#Define my graph and the filename read in
G = nx.DiGraph()
filename=""

#Instance of cmd.Cmd object which gives us the shell
class GraphShell(Cmd):
    global G
    global filename
    last_output=''

    #Set our custom prompt
    def __init__(self):
        Cmd.__init__(self)
        self.prompt = "TrustExplore> "

    #Allow user to run local shell commands and print output
    def do_shell(self, line):
        "Run a shell command"
        if line:
            print("[*] Running shell command:", line)
            output = os.popen(line).read()
            print(output)
            self.last_output = output
        else:
            print("[-] Error: Cmd Needed")

    #Handle our commands for exit, Ctrl-D and EOF
    def do_exit(self,line):
        "Exit the application. Can also use Ctrl-D"
        return True

    def do_EOF(self, line):
        "Exit the application with Crtl-D"
        return True

    #Dump the graph in GML format
    def do_gml_dump(self, line):
        "Creates output as GML of access graph"
        ofile = filename+".gml"
        nx.write_gml(G, ofile)
        print("[*] %s written!"%ofile)

    #Dump the graph in GraphML format
    def do_graphml_dump(self, line):
        "Creates output as GraphML of access graph"
        ofile = filename+".graphml"
        nx.write_graphml(G, ofile)
        print("\n[*] %s written!"%ofile)
        print ("\n[*] Red = ParentChild, Green = External, Blue = CrossLink\n")

    def do_list_nodes(self,line):
        "List all nodes for the graph"
        print("[*] All nodes: ")
        for n in nx.nodes_iter(G):
            print(n)

    #Command go show all shortest paths
    def do_path(self, args):
        "Display the shortest path between two nodes"

        arglist = args.split(" ")

        if arglist[0] and arglist[1]:
            #Grab the args
            node1=arglist[0].upper()
            node2=arglist[1].upper()
        else:
            print("[-] Error: Args Needed")

        #ensure they exist
        if G.has_node(node1) and G.has_node(node2):
            if (nx.has_path(G,node1,node2)):
                print("[*] Shortest Paths from %s to %s" %(node1,node2))
                #Get the shortest paths
                paths = nx.all_shortest_paths(G, node1, node2)

                #Print all paths in pretty format
                for p in paths:
                    outputpath = "[*] "
                    for n in p:
                        outputpath+=n+" -> "
                    print(outputpath[:-4])
            else:
                print("[-] No path exist :(")
        else:
            print("[-] Node %s or %s does not exist :(" % (node1, node2))

    #Show all paths
    def do_all_paths(self,args):
        "Display all paths between two nodes"

        arglist = args.split(" ")

        if arglist[0] and arglist[1]:
            #Grab the args
            node1=arglist[0].upper()
            node2=arglist[1].upper()
        else:
            print("[-] Error: Args Needed")

        #ensure they exist
        if G.has_node(node1) and G.has_node(node2):
            if (nx.has_path(G,node1,node2)):
                print("[*] All Paths from %s to %s" %(node1,node2))
                #Get the shortest paths
                paths = nx.all_simple_paths(G, node1, node2)

                #Print all paths in pretty format
                for p in paths:
                    outputpath = "[*] "
                    for n in p:
                        outputpath+=n+" -> "
                    print(outputpath[:-4])
            else:
                print("[-] No path exist :(")
        else:
            print("[-] Node %s or %s does not exist :(" % (node1, node2))

    #Show all domains that can be reached from a source domain 
    def do_connected(self, args):
        "Show all nodes able to be reached from a source node"
        if args:
            node = args.upper()
            if G.has_node(node):
                conn_count = 0
                print("[*] Domains reachable from \"%s\""%(node))
                for dest in G.nodes():
                    if nx.has_path(G, node, dest):
                        print(dest)
                        conn_count+=1
                print("[*] %d domains reachable from source"%conn_count)
            else:
                print("[-] Error: No node in the graph")
        else:
            print("[-] Error: Args Needed")

    #Print all neighbors of a certain node
    def do_neighbors(self,args):
        "Show all the neighbors for a certain node"

        if args:
            node = args.upper()
            if G.has_node(node):
                l = G.neighbors(node)
                
                print("[*] Neighbors:")
                for n in l:
                    print(n)
            else:
                print("[-] Error: No node in the graph")
        else:
            print("[-] Error: Args Needed")

    #print all isolated nodes
    def do_isolated(self, args):
        "Show all nodes that are isolated"
        print("[*] Isolated Nodes:")
        for n in G.nodes():
            if len(G.neighbors(n)) ==1:
                print(n)

    #calculate degree centrality and print top 5
    def do_center(self,args):
        "Show the top 5 most central nodes"
        d = nx.out_degree_centrality(G)
        cent_items=[(b,a) for (a,b) in d.iteritems()]
        cent_items.sort()
        cent_items.reverse()
        print("[*] Most Central Nodes")
        for i in range(0,5):
            if cent_items[i]:
                print(cent_items[i])

    #print some statistics
    def do_summary(self, args):
        "Show statistics on my trust map"
        ncount = len(G)
        ecount = G.number_of_edges()
        print("[*] Summary:")
        print("Filename: %s"%filename)
        print("Node Count: %d"%ncount)
        print("Edge Count: %d"%ecount)

    #notify the user if a node exist
    def do_is_node(self, args):
        "Tell the user if the node is in the graph"
        if args:
            print("[*] "+args.upper()+": "+G.has_node(args.upper()))


if __name__ == '__main__':

    parser = argparse.ArgumentParser()

    parser.add_argument("-f", "--inputfile", help='Input .csv file to process.', required=True)
    parser.add_argument("-r", "--relationship", help='PARENTCHILD, EXTERNAL, CROSSLINK')
    parser.add_argument("-g", "--graphml", action='store_true', help='Dump graphml instead of proceeding to a shell.')

    args = parser.parse_args()

    # open our file
    with open(args.inputfile, "r") as f:
        data = f.readlines()

    # save the name for later
    filename = args.inputfile
    c = 0

    # print "\n[*] Red = ParentChild, Green = External, Blue = CrossLink\n"

    # for every line in our CSV
    for line in data[1:]:
        # strip off quotes and newline characters
        stripdata = line.replace('"',"").replace('\n','')

        # split the CSV and store normalized values
        values = stripdata.split(',')
        node1 = values[0].upper()
        node2 = values[1].upper()
        relationship = values[2].upper()
        edgetype=values[4].upper()

        if( (args.relationship and (args.relationship.upper() in relationship)) or (not args.relationship)):

            # assign edge colors based on the kind of relationship
            ecolor ='#000000'

            # blue
            if "CROSSLINK" in relationship:
                ecolor='#0000CC'
            # red
            elif "PARENTCHILD" in relationship:
                ecolor='#FF0000'
            # green
            elif "EXTERNAL" in relationship:
                ecolor='#009900'

            # Add the nodes with labels
            G.add_node(node1, label=node1)
            G.add_node(node2, label=node2)

            # add the edges to the graph
            if "BIDIRECTIONAL" in edgetype:
                G.add_edge(node1, node2, color=ecolor)
                G.add_edge(node2, node1, color=ecolor)
            elif "OUTBOUND" in edgetype:
                G.add_edge(node2, node1, color=ecolor)
            elif "INBOUND" in edgetype:
                G.add_edge(node1, node2, color=ecolor)
            else:
                print("[-] UNRECOGNIZED RELATIONSHIP DIRECTION")
                exit()

            c+=1

    if args.graphml:
        GraphShell().do_graphml_dump(args)

    else:
        print("[*] %d relationships read in... starting shell" % c)
        GraphShell().cmdloop()