#!/usr/bin/env python

/*
 * Copyright (c) 2014-2015 Jonathan Anderson
 * All rights reserved.
 *
 * This software was developed by SRI International and the University of
 * Cambridge Computer Laboratory under DARPA/AFRL contract FA8750-10-C-0237
 * ("CTSRD"), as part of the DARPA CRASH research programme.
 *
 * This software was developed at the University of Cambridge Computer
 * Laboratory with support from a grant from Google, Inc.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

import argparse
import itertools
import json
import sys

import callgraph
import dot
import soaap


args = argparse.ArgumentParser()
args.add_argument('filename')
args.add_argument('-o', '--output', default = '-')
args.add_argument('-a', '--analyses', default = 'vulnerability_warning',
                  help = ("Which SOAAP analysis outputs to parse. Options are: "
                          + ', '.join(soaap.analyses.keys()) + '.'))
args.add_argument('-c', '--cluster-by', default = '',
                  help = ("Cluster functions in the output graph by"
                          + " 'filename', 'namespace' or 'sandbox',"
                          + " or, using a comma-separated list, by"
                          + " several of these."))
args.add_argument('-s', '--simplify', action = 'store_true')

args = args.parse_args()

f = open(args.filename, 'r')
out = open(args.output, 'w') if args.output != '-' else sys.stdout

data = json.load(f)['soaap']

functions = set()

for analysis in args.analyses.split(','):
    (f,c) = soaap.parse(data, analysis)

    if functions:
        functions = functions.union(f)
        calls.update(c)

    else:
        (functions, calls) = (f,c)

#
# How do we want to group nodes into subgraphs?
#
if args.cluster_by == '':
    # Put everything together in one big cluster.
    cluster_key = lambda node: ''

else:
    cluster_by = args.cluster_by.split(',')
    keys = []

    if 'filename' in cluster_by:
        keys.append(lambda fn: 'filename: %s' % fn.filename())

    if 'namespace' in cluster_by:
        keys.append(lambda fn: 'namespace: %s' % fn.namespace())

    if 'sandbox' in cluster_by:
        keys.append(lambda fn: 'sandbox: %s' % fn.sandbox())

    assert len(keys) == len(cluster_by), '%s -> %s' % (cluster_by, keys)

    def cluster_key(node):
        return ', '.join([ str(k(node)) for k in keys ])


if args.simplify:
    (functions, calls) = callgraph.simplify(functions, calls, cluster_key)


#
# Convert call graph into GraphViz .dot format.
#
style = {
    'default': {
        'label': 'Unsandboxed',
        'fillcolor': '#cccccccc',
    },

    'mitigated': {
        'label': 'Mitigated',
        'fillcolor': '#ffff66cc',
        'shape': 'octagon'
    },

    'prev_vuln': {
        'label': 'Previously vulnerable',
        'fillcolor': '#ff9999cc',
        'shape': 'doubleoctagon'
    },

    'private': {
        'label': 'Private',
        'fillcolor': '#ff99cccc',
        'shape': 'invhouse'
    },

    'sandboxed': {
        'label': 'Sandboxed',
        'fillcolor': '#99ff9999',
        'style': 'dashed,filled',
    },
}


out.write('''
digraph {
    node [ fontname = "Inconsolata" ];
    edge [ fontname = "Avenir" ];

    labeljust = "l";
    labelloc = "b";
    rankdir = "BT";

    subgraph cluster_legend {
        label = "Legend";

        edge [ style = "invis" ];
        ''')


for (index, name) in enumerate(style):
    attrs = {
        'fontname': 'Avenir',
        'fontsize': '48',
        'style': 'filled',
    }
    attrs.update(style[name])

    out.write('      legend_%d %s\n' % (index, dot.attribute_list(attrs)))

for i in xrange(1, len(style)):
    out.write('      legend_%d -> legend_%d;\n' % (i - 1, i))

out.write('}\n\n')


clusters = {}
functions = sorted(functions, key = cluster_key)

for (cluster, group) in itertools.groupby(functions, key = cluster_key):
    name = 'cluster_%d' % (len(clusters) + 1)
    clusters[cluster] = name

    if cluster:
        out.write('    subgraph %s {\n' % name)
        out.write('        bgcolor = "#cccccc33";\n')
        out.write('        label = "%s";\n' % cluster)
        out.write('\n')

    for node in group:
        if node.cve and node.sandbox_name:
            attributes = style['mitigated']

        elif node.cve:
            attributes = style['prev_vuln']

        elif node.owner:
            attributes = style['private']

        elif node.sandbox_name:
            attributes = style['sandboxed']

        else:
            attributes = style['default']

        attributes['label'] = node.fn[:80]
        if 'style' not in attributes:
            attributes['style'] = 'filled'

        attributes = dot.attribute_list(attributes)

        out.write('        "%s" %s;\n' % (node.name(), attributes))

    if cluster:
        out.write('    }\n\n')

for ((source, dest), weight) in calls.items():
    assert source in functions
    assert dest in functions
    assert weight > 0

    label = ('' if args.simplify else
             '%s:%d' % (source.location['file'], source.location['line']))

    attrs = dot.attribute_list({
        'label': label,
        'lhead': clusters[cluster_key(source)],
        'ltail': clusters[cluster_key(dest)],
        'penwidth': weight,
        'weight': weight,
    })

    out.write('    "%s" -> "%s" %s\n' % (source.name(), dest.name(), attrs))

out.write('}\n')
