#!/usr/bin/perl
# graphnode - graph a node from QCMDB
#
# Written by Digant C Kasundra <digant@stanford.edu>
# Copyright 2009 Board of Trustees, Leland Stanford Jr. University
#
# See LICENSE for licensing terms.

#############################################################################
# Modules and declarations
#############################################################################

use strict;

use AppConfig;
use GraphViz;
use Log::Log4perl qw(get_logger);
use QCMDB::Config;
use QCMDB::Database;
use Switch;
use XML::Simple;
use Data::Dumper;


my $CFGFILE = '/etc/extractor.conf';

# globals
my $qdb = QCMDB::Database->new(QCMDB::Config->new($CFGFILE))
    or die "Could not connect to the DB.";
my $dbh = $qdb->{dbh};
my $g = GraphViz->new(
        layout =>'dot',
        node => { fontsize => '10' },
        edge => { fontsize => '10' },
    );
my ($log, $seen, $mdrId);

##############################################################################
# Internal helper methods
##############################################################################

# Get a node based on localId 
sub _get_node {
    my ($id) = @_;
    my ($node, $props, $sql);

    $sql = sprintf("select ci.id, ci.localId, tn.name as type " . 
        "from cmdb_items as ci, cmdb_typenames as tn " .
        "where ci.localId = '%s' and ci.mdrId like'%s' and ci.type=tn.id",
        $id, $mdrId);
    eval {
        $node = $dbh->selectrow_hashref ($sql);
        $dbh->commit;
    };
    if ($@) {
        $log->error("$@");
        $dbh->rollback;
        return;
    }
    if ($node eq undef) {
        return;
    }

    # Find the node's properties
    eval {
        $sql = sprintf("select cpn.name, cip.value " .
            "from cmdb_propnames as cpn, cmdb_itemprops as cip " .
            "where cip.iid = %d and cpn.id=cip.pid", $node->{id});
        $props = $dbh->selectall_hashref($sql, 'name');
        $dbh->commit;
    };
    if ($@) {
        $log->error("$@");
        $dbh->rollback;
        return;
    }

    foreach my $key (keys %$props) {
        $node->{props}->{$key} = $props->{$key}->{value};
    }

    # And just to make it easier to grab...
    $node->{name} = $props->{Name}->{value};

    return $node;
}

# Get an existing relationship if found, or return error
sub _get_edge {
    my ($localId) = @_;
    my ($edge, $props);

    # Find the node by name
    eval {
        my $sql = "select cr.id, cr.localId, " .
            "ci.localId as source, cii.localId as target, " .
            "tn.name as type " . 
            "from cmdb_rels as cr, cmdb_typenames as tn, " .
            "cmdb_items as ci, cmdb_items as cii " .
            "where cr.localId = ? and cr.mdrId like ? and cr.type=tn.id " .
            "and ci.id = cr.source and cii.id = cr.target";
        $edge= $dbh->selectrow_hashref ($sql, undef, $localId, $mdrId);
        $dbh->commit;
    };
    if ($@) {
        $log->error("$@");
        $dbh->rollback;
        return;
    }
    if ($edge eq undef) {
        return;
    }

    # Find the edge's properties
    eval {
        my $sql = sprintf("select cpn.name, crp.value " .
            "from cmdb_propnames as cpn, cmdb_relprops as crp " .
            "where crp.rid = %d and cpn.id=crp.pid", $edge->{id});
        $props = $dbh->selectall_hashref($sql, 'name');
        $dbh->commit;
    };
    if ($@) {
        $log->error("$@");
        $dbh->rollback;
        return;
    }

    foreach my $key (keys %$props) {
        $edge->{props}->{$key} = $props->{$key}->{value};
    }

    # And just to make it easier to grab...
    $edge->{name} = $props->{Name}->{value};

    return $edge;
}

# Get relationships where the specified itemId is the source. Returns array
# of edges
sub _get_edges_with_source {
    my ($iid) = @_;
    my $edgeIds;
    my @edges;
    eval {
        my $sql = sprintf("select localId from cmdb_rels where source=%d", 
            $iid); 
        $edgeIds = $dbh->selectall_arrayref($sql);
        $dbh->commit;
    };
    if ($@) {
        $log->error("$@");
        $dbh->rollback;
        return;
    }

    # load all of these edges
    foreach my $id (@$edgeIds) {
        my $edge = _get_edge(@$id[0]);
        if ($edge) {
            push (@edges, $edge);
        }
    }

    return @edges;
}

# Get relationships where the specified itemId is the target. Returns array
# of edges
sub _get_edges_with_target {
    my ($iid) = @_;
    my $edgeIds;
    my @edges;
    eval {
        my $sql = sprintf("select localId from cmdb_rels where target=%d", 
            $iid); 
        $edgeIds = $dbh->selectall_arrayref($sql);
        $dbh->commit;
    };
    if ($@) {
        $log->error("$@");
        $dbh->rollback;
        return;
    }

    # load all of these edges
    foreach my $id (@$edgeIds) {
        my $edge = _get_edge(@$id[0]);
        if ($edge) {
            push (@edges, $edge);
        }
    }

    return @edges;
}

# translates type names into style information for 'fontcolor', 'color',
# 'fillcolor', 'style'
sub _get_style_for_type{
    my ($type) = @_;
    switch ($type) {
        case 'application'  
                {return ('white','black', 'black', 'filled')}
        case 'built-on'     
                {return ('red4','red4', undef, 'filled')}
        case 'computerSystem'  
                {return ('white','black', 'SteelBlue', 'filled')}
        case 'ipEndpoint'   
                {return ('black','black', 'gainsboro', 'filled')}
        case 'lanEndpoint' 
                {return ('black','black', 'gainsboro', 'filled')}
        case 'operatingSystem' 
                {return ('black','black', 'LightSteelBlue', 'filled')}
        case 'resolves-to'     
                {return ('gray50','gray50', undef, 'filled')}
        case 'runs-operating-system' 
                {return ('gray50','gray50', undef, 'filled')}
        case 'softwareServer'  
                {return ('black','black', 'Khaki', 'filled')}
        case 'virtualSystem'   
                {return ('white','black', 'SteelBlue', 'filled')}
        case 'with-ip-address' 
                {return ('gray50','gray50', undef, 'filled')}
        case 'with-mac-address' 
                {return ('gray50','gray50', undef, 'filled')}
        else    {return ('black', 'black', 'white', 'filled')}
    }
}

sub _add_node_to_graph {
    my ($node) = @_;
    my $label = $node->{name}.'\n&lt;'.$node->{type}.'&gt;';
    my ($fc, $c, $fill, $st) = _get_style_for_type($node->{type});
    $log->debug("Add node to graph: $label ($fc, $c, $fill, $st)");
    $g->add_node($node->{localId}, 
            label     => $label, 
            shape     => 'record',
            fontcolor =>$fc,
            color     =>$c,
            fillcolor =>$fill,
            style     =>$st,
        );

}

sub _add_edge_to_graph {
    my ($edge) = @_;
    my $label;
    if ($edge->{name}) { 
        $label = $edge->{name};
    } else {
        $label = "UNKNOWN";
        $log->error("No name in " . Dumper($edge));
    }
    my ($fc, $c, undef, $st) = _get_style_for_type($edge->{name});
    $log->debug("Add edge to graph: $label ($fc, $c, $st)");
    $g->add_edge($edge->{source} => $edge->{target}, 
            label     => $label,
            fontcolor =>$fc,
            color     =>$c,
            style     =>$st,
        );
}

# Recursive function for finding nodes
sub _find_nodes {
    my ($depth, $localId, $dir) = @_;
    my (@upEdges, @downEdges);
    # Load node
    my $node = _get_node($localId);

    $log->debug("_find_nodes: depth: $depth; dir: $dir; id: $localId");
    if (!$node) {
        $log->error("Failed to find $localId\n");
        return;
    }
    $log->debug("Found $localId");

    # Add to graph
    _add_node_to_graph($node);

    # decrease the depth count and if we are still positive, find the
    # children and recurse
    $depth--;

    # MID FUNCTION EXIT!
    if ($depth == 0) {
        # returning when not an error in the middle of a function is bad but
        # will make it cleaner to not have half the function in an else
        # block
        return $node;
    }

    if ($dir eq 'down') {
        $log->debug("Getting edges down from $localId");
        @downEdges = _get_edges_with_source($node->{id});
    } elsif ($dir eq 'up') {
        $log->debug("Getting edges up from $localId");
        @upEdges = _get_edges_with_target($node->{id});
    } else {
        $log->debug("Getting edges up and down from $localId");
        @downEdges = _get_edges_with_source($node->{id});
        @upEdges = _get_edges_with_target($node->{id}); 
    }

    # iterate over children and if the target children are found, add
    # that edge
    if (($dir eq 'down') or ($dir eq 'both')) {
        foreach my $cEdge (@downEdges) {
            # Skip if not a valid edge
            if (!$cEdge->{localId}) { 
                $log->debug("Skipping invalid edge " . Dumper($cEdge));
                next;
            }
            if ($seen->{$cEdge->{localId}}==1) {
                $log->debug("Already saw edge $cEdge->{localId}");
                next;
            } else {
                # mark edge as visited
                $log->debug("Traversing edge $cEdge->{localId}");
                $seen->{$cEdge->{localId}} = 1;
            }
            if (_find_nodes($depth, $cEdge->{target}, $dir)) {
                _add_edge_to_graph($cEdge);
            } else {
                $log->error("Couldn't find node $dir from $cEdge->{localId}");
            }
        }
    }
    if (($dir eq 'up') or ($dir eq 'both')) {
        foreach my $cEdge (@upEdges) {
            # Skip if not a valid edge
            if (!$cEdge->{localId}) { 
                $log->debug("Skipping invalid edge " . Dumper($cEdge));
                next;
            }
            if ($seen->{$cEdge->{localId}}==1) {
                $log->debug("Already saw edge $cEdge->{localId}");
                next;
            } else {
                # mark edge as visited
                $log->debug("Traversing edge $cEdge->{localId}");
                $seen->{$cEdge->{localId}} = 1;
            }
            if (_find_nodes($depth, $cEdge->{source}, $dir)) {
                _add_edge_to_graph($cEdge);
            } else {
                $log->error("Couldn't find node $dir from $cEdge->{localId}");
            }
        }
    }

    return $node;
}

sub _print_usage {
    print "Usage: graphnode <mdrId> <localId> <depth> <up|down|both>\n";
}

###########################################################################
# MAIN
###########################################################################

my $CONF = AppConfig->new({CREATE => 1}, 'logconfig=s');
if (!$CONF->file($CFGFILE)) {
    die "Could not get logconfig file location."
}

# Setup logging
Log::Log4perl->init($CONF->logconfig());
$log = Log::Log4perl->get_logger("GN");

unless (defined($ARGV[3])) {
    _print_usage();
    exit;
}

$mdrId = $ARGV[0];
my $localId = $ARGV[1];
my $depth = $ARGV[2] + 1;
my $dir = $ARGV[3];

if ($dir eq 'down') {
    _find_nodes($depth, $localId, 'down');
} elsif ($dir eq 'up') {
    _find_nodes($depth, $localId, 'up');
} elsif ($dir eq 'both') {
    _find_nodes($depth, $localId, 'both');
} else {
    _print_usage();
    exit;
}

$log->debug("Generating PNG...");
print $g->as_png;
