HighDDG.scala

//Decompile the function at the cursor, then build data-flow graph (AST)
//@category    PCode

// Use with https://github.com/edmcman/ghidra-scala-loader

// XXX: Fix the listener so we can click on the graph. Requires slight change to PCodeDfgDisplayListener

// XXX: Make a button of prompt to write to a dot file

import ghidra.app.decompiler._
import ghidra.app.script.GhidraScript
import ghidra.app.services.GraphDisplayBroker

import ghidra.framework.plugintool.PluginTool

import ghidra.util.task.{Task, TaskMonitor}
import ghidra.util.Msg
import ghidra.service.graph._
import ghidra.program.model.pcode._
import ghidra.app.plugin.core.decompile.actions.{PCodeDfgGraphType, PCodeDfgDisplayOptions}

import ghidra.program.model.pcode.SequenceNumber

import scala.collection.JavaConversions._
import scala.collection.mutable.HashMap

class GraphSettings(val tool: PluginTool) extends PCodeDfgDisplayOptions(tool) {
}

// Ghidra's AttributeVertex::toString is horrible
class NewAttributedVertex(id: String, name: String) extends AttributedVertex(id, name) {
    override def toString(): String = name
}

// This is a copy of PCodeDfgGraphTask that allows us to access the graph.
class DfgGraph(val tool: PluginTool, val graphService: GraphDisplayBroker, val hf: HighFunction) /*extends Task*/ {

    Msg.info(this, "Building Data Flow Graph")

    // What is this?
    def getOpKey(op : PcodeOpAST) = {
        val sq = op.getSeqnum
        val addr = sq.getTarget.toString(false)
        s"${sq.toString}"
    }

    def getVarnodeKey(vn : VarnodeAST) = {
        val defOp: Option[PcodeOp] = Option(vn.getDef)

        defOp match {
            case Some(defOp) => s"def ${defOp.getSeqnum.toString}"
            case None => s"vn ${vn.getUniqueId.toString}"
        }   
    }

    def highvarToName(highvar: HighVariable) = {
        // If this is a symbol, use that name. Otherwise use the highvar name.
        Option(highvar.getSymbol) match {
            case Some(symbol) => Option(symbol.getName).map(s => s"highsym: $s")
            case None => Option(highvar.getName).map(n => s"highvar: $n")
        }
    }

    def createVarnodeVertex(graph: AttributedGraph, vn: VarnodeAST) = {
        val id = getVarnodeKey(vn)
        //var name = vn.getAddress.toString(true)
        var name = vn.toString
        var vertexType = PCodeDfgGraphType.DEFAULT_VERTEX

        if (vn.isConstant) {
            vertexType = PCodeDfgGraphType.CONSTANT
        } else if (vn.isRegister) {
            val reg = Option(hf.getFunction.getProgram.getRegister(vn.getAddress, vn.getSize))
            reg.map(r => name = r.getName)
            vertexType = PCodeDfgGraphType.REGISTER
        } else if (vn.isUnique) {
            vertexType = PCodeDfgGraphType.UNIQUE
        } else if (vn.isPersistent) {
            vertexType = PCodeDfgGraphType.PERSISTENT
        } else if (vn.isAddrTied) {
            vertexType = PCodeDfgGraphType.ADDRESS_TIED
        }

        val highvar: Option[HighVariable] = Option(vn.getHigh)
        val highvarname = highvar.map(highvarToName)

        highvarname match {
            case Some(Some(highvarname)) => name = name + "\n" + highvarname
            case _ => ()
        }

        if (vn.isInput) {
            name = "Input: " + name
        }

        name = Option(vn.getDef) match {
            case None => name
            case Some(defOp) => name + s"\n@${defOp.getSeqnum.getTarget.toString(false)}"
        }

        val vert = new NewAttributedVertex(id, name)
        graph.addVertex(vert)
        //Msg.info(this, s"Varnode id=$id name=$name")
        vert.setVertexType(vertexType)

        // Set inputs to triangle shape
        if (vn.isInput) {
            vert.setAttribute(PCodeDfgDisplayOptions.SHAPE_ATTRIBUTE, VertexShape.TRIANGLE_DOWN.getName)
        }

        vert
    }

    def getVarnodeVertex(graph: AttributedGraph, vertices: HashMap[Int, AttributedVertex], hf: HighFunction, vn: VarnodeAST): AttributedVertex = {
        val id = vn.getUniqueId
        vertices.get(id) match {
            case Some(v) => v
            case None => {
                val v = createVarnodeVertex(graph, vn)
                vertices(id) = v
                v
            }
        }
    }

    def createOpVertex(graph: AttributedGraph, op: PcodeOpAST) = {
        var name = op.getMnemonic // toString is very verbose
        val id = getOpKey(op)
        val opcode = op.getOpcode

        opcode match {
            case PcodeOp.LOAD => {
                val vn = op.getInput(0)
                val addrspace = hf.getFunction.getProgram.getAddressFactory.getAddressSpace(vn.getOffset.toInt)
                name = name + " " + addrspace.getName
            }
            case PcodeOp.INDIRECT => {
                val vn = Option(op.getInput(1))
                vn.map(vn => Option(hf.getOpRef(vn.getOffset.toInt))
                             .map(indop => name = name + s" (${indop.getMnemonic})")
                )

            }
            case _ => {

            }

        }
        name = name + s"\n@${op.getSeqnum.getTarget.toString(false)}"

        val vert = new NewAttributedVertex(id, name)
        graph.addVertex(vert)
        vert.setVertexType(PCodeDfgGraphType.OP)

        vert
    }

    def createEdge(graph: AttributedGraph, in: AttributedVertex, out: AttributedVertex): AttributedEdge = {
        val edge = graph.addEdge(in, out)
        edge.setEdgeType(PCodeDfgGraphType.DEFAULT_EDGE)
        edge
    }

    def getGraph(): AttributedGraph = {
        val graph = new AttributedGraph(s"Data Flow Graph for ${hf.getFunction.getName}", new PCodeDfgGraphType)

        val vertices = new HashMap[Int, AttributedVertex]()

        for (pcodeOp <- hf.getPcodeOps) {
            //val pcodeOpcode = op.getOpcode

            Msg.info(this, s"Pcode op $pcodeOp")

            val opVertex = createOpVertex(graph, pcodeOp)

            for ((input,index) <- pcodeOp.getInputs.zipWithIndex) {
                (pcodeOp.getOpcode, index) match {
                    //case (_, 0) => ()
                    case (PcodeOp.LOAD | PcodeOp.STORE, 0) => () 
                    case (PcodeOp.INDIRECT, 1) => ()
                    case _ => {
                        // Create input node and edge
                        val inv = getVarnodeVertex(graph, vertices, hf, input.asInstanceOf[VarnodeAST])
                        createEdge(graph, inv, opVertex)
                    }
                }
                // Create output node
                val outvarnode: Option[VarnodeAST] = Option(pcodeOp.getOutput.asInstanceOf[VarnodeAST])
                val outvertex = outvarnode.map(outvarnode => getVarnodeVertex(graph, vertices, hf, outvarnode))
                outvertex.map(outvertex => createEdge(graph, opVertex, outvertex))
            }
        }

        graph
    }

    def displayGraph(monitor: TaskMonitor) = {
        val display = graphService.getDefaultGraphDisplay(false, monitor)
        val displayOptions = new GraphSettings(tool)
        val graph = getGraph ()
        display.setGraph(graph, displayOptions, s"Data Flow Graph for ${hf.getFunction.getName}", true, monitor)
    }
}

import org.jgrapht.nio.{Attribute, AttributeType, DefaultAttribute}

// Sigh, the default Dot exporter uses the vertices' ID as the label instead of the name
class DotNameExporter extends ghidra.graph.exporter.DotGraphExporter {
    override def getAttributes(v: Attributed): java.util.Map[String, Attribute] = {
        val attributeMap = new java.util.HashMap[String, Attribute]()

        for (entry <- v.entrySet()) {
            val key = entry.getKey()
            val value = entry.getValue()
            attributeMap.put(key, new DefaultAttribute(value, AttributeType.STRING))
            if (key == "Name") {
                attributeMap.put("label", new DefaultAttribute(value, AttributeType.STRING))
            }
        }
        attributeMap
    }
}

class HighDDG extends GhidraScript {

  override def run() = {
    println("Hello world, I'm written in Scala...!")

    val tool = state.getTool
    assert (tool != null)
    val graphService = tool.getService(classOf[GraphDisplayBroker])
    assert (graphService != null)

    // graphService.getGraphExporters().map(e => e.getName()).foreach(println)

    val ifc = new DecompInterface
    if (!ifc.openProgram(currentProgram)) {
      throw new DecompileException("Decompiler", "Unable to initialize: " + ifc.getLastMessage)
    }

    val res = ifc.decompileFunction(currentProgram.getFunctionManager.getFunctionContaining(currentAddress), 60, null)
    val hf = res.getHighFunction
    val dg = new DfgGraph(tool, graphService, hf)
    dg.displayGraph(monitor)

    // Dump graph to file
    val file = askFile("Export graph to DOT file? Provide output filename or cancel.", "Export graph to file")
    val exporter = new DotNameExporter
    exporter.exportGraph(dg.getGraph (), file)    
  }
}