BuildSoftwareGraph.py

import os
import sys
import json
import networkx as nx
# import matplotlib.pyplot as plt
import glob
from BuildVulPredicates import findVulnerableBinaries, createPredicatesForVulBinaries, fixName, findOpenSourceBinaries,\
    createPredicatesForOpenSourceBinaries

from networkx.drawing.nx_pydot import graphviz_layout

Internet = "Internet"
DEBUG = False

def printDebug(debugMessage):
    if DEBUG:
        print(debugMessage)

# Read given JSON file and create a directed graph
def readGraph(binaryTopologyFile):
    if not os.path.isfile(binaryTopologyFile):  # File does not exist
        print("File {} does not exist".format(binaryTopologyFile))
        raise Exception()

    file = open(binaryTopologyFile)
    try:
        json_object = json.load(file)
    except ValueError:
        print("Could not load {} json file.".format(file))
        return

    # Handle the nodes graph, which look like:
    # "binary_name": {
    #  "version": [],
    #  "peers": [
    #   {"name": "binary_name",
    #    "type": "exec/environment",
    #    "info": []
    #   }
    #  ]
    # }
    swGraph = nx.DiGraph()
    graph = json_object['graph']
    for nodeName in graph:
        printDebug("Node Name: " + nodeName)
        printDebug("Versions of {}: {}".format(nodeName, str(graph[nodeName]['version'])))
        if not swGraph.has_node(nodeName):  # If this is a new node, add it
            swGraph.add_node(nodeName, Versions=graph[nodeName]['version'])
        else:  # Otherwise, update its versions
            swGraph.nodes[nodeName]['Versions'] = graph[nodeName]['version']

        # Handle node's children, which are under 'peers'
        printDebug("Going to handle children nodes of {}: {}".format(nodeName, str(graph[nodeName]['peers'])))
        for child in graph[nodeName]['peers']:
            childName = child['name'] if child['name'].casefold() != Internet.casefold() else Internet
            printDebug("Node Name: " + childName)
            if not swGraph.has_node(childName):  # If this is a new node, add it
                swGraph.add_node(childName, Versions=[])
            if childName == Internet:
                swGraph.add_edge(Internet, nodeName, Type=None, Info=None)  # Add edge
            elif nodeName != childName:  # For unclear cases in Karonte output that node refers to itself
                swGraph.add_edge(nodeName, childName, Type=child['type'], Info=child['info'])  # Add edge

    return swGraph

# Create MulVAL (software) model (interaction rules) for the inter-binaries dependencies (data flaw)
def createSoftwareModel(swGraph, filePathToCreate):
    fileStream = open(filePathToCreate, "w")
    fileStream.write("attackGoal(vulnerableSoftware(Software)).\n\n")
    # fileStream.write("vulExists('CVE-2022-33099', uhttpd, 'NETWORK', 'availability_loss', 'HIGH').\n")  # For DEBUG only
    # fileStream.write("vulExists('CVE-2022-33099', busybox, 'NETWORK', 'availability_loss', 'CRITICAL').\n\n")  # For DEBUG only
    for edge in swGraph.edges(data=True):
        # For border binaries, create externalInteraction(_source, _software, _interaction_type)
        # For others, create 'dataFlow' interaction rule with flow type
        if edge[0] == Internet:
            fileStream.write("externalInteraction('{}', {}, {}).\n".format(Internet, fixName(edge[1]), Internet))
        else:
            flowType = edge[2]['Type']
            fileStream.write("dataFlow({}, {}, '{}').\n".format(fixName(edge[0]), fixName(edge[1]), flowType))
    fileStream.write("\n")
    fileStream.close()

# Create list of binaries, which will be used for building vulnerability predicates
def createBinariesList(swGraph, filePathToCreate):
    fileStream = open(filePathToCreate, "w")
    binariesWithVersion = []
    for node in swGraph.nodes():
        if len(swGraph.nodes[node]['Versions']) == 0:
            binaryWithVersion = (node, '*')
            if binaryWithVersion not in binariesWithVersion:
                fileStream.write("{}, {}\n".format(node, '*'))
                binariesWithVersion.append(binaryWithVersion)
        else:
            for version in swGraph.nodes[node]['Versions']:
                if version.startswith('v') or version.startswith('V'):
                    version = version[1:]
                binaryWithVersion = (node, version)
                if binaryWithVersion not in binariesWithVersion:
                    fileStream.write("{}, {}\n".format(node, version))
                    binariesWithVersion.append(binaryWithVersion)
    fileStream.close()
    return binariesWithVersion

# Input: The file of SW graph to process, SW model file to create, and (temp) file to create with list of binaries
# Create MulVAL model (interaction rules) for the software graph and relevant vulnerabilities
def createSoftwareModelAndVul(binaryTopologyFile, swModelFile, binariesListFile):
    if swModelFile is None:
        # dirIndex = binaryTopologyFile.rfind('/') + 1  # Exclude '/'; if no directory, change returned -1 to 0
        dirIndex = 0  # Create SW model in the same directory as the topology
        swModelFile = binaryTopologyFile[dirIndex:binaryTopologyFile.rfind('.')]+".p"
    if binariesListFile is None:
        # binariesListFile = "BinariesList.csv"
        dirIndex = 0  # Create SW model in the same directory as the topology
        binariesListFile = binaryTopologyFile[dirIndex:binaryTopologyFile.rfind('.')]+"-BinariesList.csv"

    print("The file to read: " + binaryTopologyFile)
    swGraph = readGraph(binaryTopologyFile)

    createSoftwareModel(swGraph, swModelFile)
    print("The software model file {} was created.".format(swModelFile))
    binaries = createBinariesList(swGraph, binariesListFile)
    print("The binaries list file {} was created.".format(binariesListFile))
    vulBinaries = findVulnerableBinaries(binaries)
    createPredicatesForVulBinaries(vulBinaries, swModelFile)
    print("The software model file {} was updated with vulnerable binaries.".format(swModelFile))
    openSourceBinaries = findOpenSourceBinaries(binaries)
    createPredicatesForOpenSourceBinaries(openSourceBinaries, swModelFile)
    print("The software model file {} was updated with open-source binaries.".format(swModelFile))

# Draw binary topology graph and save in a file
# mode: 1 - circular_layout, 2 - graphviz_layout(dot), 3 - graphviz_layout(circo)
def drawTopologyGraph(binaryTopologyFile, mode):
    print("The file to read: " + binaryTopologyFile)
    swGraph = readGraph(binaryTopologyFile)
    try:
        longestPath = nx.dag_longest_path(swGraph)
        print(f"The longest path is {longestPath}")
    except:
        print("Failed in dag_longest_path()")
    # Remove nodes with no edges
    nodesToRemove = []
    for node in swGraph.nodes:
        if len(list(swGraph.successors(node))) == 0 and len(list(swGraph.predecessors(node))) == 0:
            nodesToRemove.append(node)
    for node in nodesToRemove:
        swGraph.remove_node(node)

    if mode == 1:
        layout = nx.circular_layout(swGraph)
        # layout = nx.shell_layout(swGraph)
        nx.draw_networkx(swGraph, layout)
    elif mode == 2:
        pos = graphviz_layout(swGraph, prog="dot")
        for k, v in pos.items():
            pos[k] = (-v[1], v[0])

        nx.draw_networkx_nodes(swGraph, pos=pos, node_shape='s', node_size=300, node_color='none', edgecolors='k')
        nx.draw_networkx_edges(swGraph, pos=pos, node_shape='s', width=1, node_size=300)
        nx.draw_networkx_labels(swGraph, pos=pos, font_size=5)
    elif mode == 3:
        pos = graphviz_layout(swGraph, prog="circo")
        nx.draw_networkx_nodes(swGraph, pos=pos, node_shape='s', node_size=100, node_color='none', edgecolors='k')
        nx.draw_networkx_edges(swGraph, pos=pos, node_shape='s', width=1, node_size=100)
        nx.draw_networkx_labels(swGraph, pos=pos, font_size=3)

    swGraphFile = binaryTopologyFile[0:binaryTopologyFile.rfind('.')]+"_Graph.pdf"
    # plt.savefig(swGraphFile, dpi=600)
    # if DEBUG:
        # plt.show()

def main():
    # The file of software graph to process
    binaryTopologyFile = "BinaryTopologies/V4.1/TP*.json"
    # binaryTopologyFile = "BinaryTopologies/V4.1/*.json"
    if len(sys.argv) > 1:
        binaryTopologyFile = sys.argv[1]
    # The file to create with list of predicates for software model (dependencies)
    swModelFile = sys.argv[2] if len(sys.argv) > 2 else None
    # The file to create with list of binaries
    binariesListFile = sys.argv[3] if len(sys.argv) > 3 else None

    if binaryTopologyFile.find('*') == -1:  # A single file to process
        createSoftwareModelAndVul(binaryTopologyFile, swModelFile, binariesListFile)
        # drawTopologyGraph(binaryTopologyFile, 3)
    else:  # Run over all the JSON files in the given directory
        # Build the list of JSON files to process
        topologyFiles = glob.glob(binaryTopologyFile)
        for topologyFile in topologyFiles:
            createSoftwareModelAndVul(topologyFile, swModelFile, binariesListFile)
            # drawTopologyGraph(topologyFile, 3)


if __name__ == "__main__":
    main()