MaltegoTransform.py

#
# Updated Maltego Python library
# 2013/03/30
# See TRX documentation
#
# RT

import sys
from xml.dom import minidom

BOOKMARK_COLOR_NONE = "-1"
BOOKMARK_COLOR_BLUE = "0"
BOOKMARK_COLOR_GREEN = "1"
BOOKMARK_COLOR_YELLOW = "2"
BOOKMARK_COLOR_ORANGE = "3"
BOOKMARK_COLOR_RED = "4"
LINK_STYLE_NORMAL = "0"
LINK_STYLE_DASHED = "1"
LINK_STYLE_DOTTED = "2"
LINK_STYLE_DASHDOT = "3"
UIM_FATAL = "FatalError"
UIM_PARTIAL = "PartialError"
UIM_INFORM = "Inform"
UIM_DEBUG = "Debug"


class MaltegoEntity(object):
    value = ""
    weight = 100
    displayInformation = []
    additionalFields = []
    iconURL = ""
    entityType = "Phrase"

    def __init__(self, eT=None, v=None, additionFields=[]):
        if eT is not None:
            self.entityType = eT
        if v is not None:
            self.value = v
        self.additionalFields = additionFields
        self.weight = 100
        self.displayInformation = []
        self.iconURL = ""

    def setType(self, eT=None):
        if eT is not None:
            self.entityType = eT

    def setValue(self, eV=None):
        if eV is not None:
            self.value = eV

    def setWeight(self, w=None):
        if w is not None:
            self.weight = w

    def addDisplayInformation(self, di=None, dl="Info"):
        if di is not None:
            self.displayInformation.append([dl, di])

    def addProperty(
        self, fieldName=None, displayName=None, matchingRule=False, value=None
    ):
        self.additionalFields.append([fieldName, displayName, matchingRule, value])

    def setIconURL(self, iU=None):
        if iU is not None:
            self.iconURL = iU

    def setLinkColor(self, color):
        self.addProperty("link#maltego.link.color", "LinkColor", "", color)

    def setLinkStyle(self, style):
        self.addProperty("link#maltego.link.style", "LinkStyle", "", style)

    def setLinkThickness(self, thick):
        self.addProperty("link#maltego.link.thickness", "Thickness", "", str(thick))

    def setLinkLabel(self, label):
        self.addProperty("link#maltego.link.label", "Label", "", label)

    def setBookmark(self, bookmark):
        self.addProperty("bookmark#", "Bookmark", "", bookmark)

    def setNote(self, note):
        self.addProperty("notes#", "Notes", "", note)

    def returnEntity(self):
        r = ""
        r += '<Entity Type="' + str(self.entityType) + '">'
        r += "<Value>" + str(self.value) + "</Value>"
        r += "<Weight>" + str(self.weight) + "</Weight>"
        if len(self.displayInformation) > 0:
            r += "<DisplayInformation>"
            for i in range(len(self.displayInformation)):
                r += (
                    '<Label Name="'
                    + self.displayInformation[i][0]
                    + '" Type="text/html"><![CDATA['
                    + str(self.displayInformation[i][1])
                    + "]]></Label>"
                )
            r += "</DisplayInformation>"
        if len(self.additionalFields) > 0:
            r += "<AdditionalFields>"
            for i in range(len(self.additionalFields)):
                if str(self.additionalFields[i][2]) != "strict":
                    r += (
                        '<Field Name="'
                        + str(self.additionalFields[i][0])
                        + '" DisplayName="'
                        + str(self.additionalFields[i][1])
                        + '">'
                        + str(self.additionalFields[i][3])
                        + "</Field>"
                    )
                else:
                    r += (
                        '<Field MatchingRule="'
                        + str(self.additionalFields[i][2])
                        + '" Name="'
                        + str(self.additionalFields[i][0])
                        + '" DisplayName="'
                        + str(self.additionalFields[i][1])
                        + '">'
                        + str(self.additionalFields[i][3])
                        + "</Field>"
                    )
            r += "</AdditionalFields>"
        if len(self.iconURL) > 0:
            r += "<IconURL>" + self.iconURL + "</IconURL>"
        r += "</Entity>"
        return r


class MaltegoTransform(object):
    # We were lazy to use a proper XML library to generate
    # our XML. Thus - encode data before you insert!
    # ..Sorry - RT
    entities = []
    exceptions = []
    UIMessages = []

    def __init__(self):
        self.entities = []
        self.exceptions = []
        self.UIMessages = []
        self = None

    def addEntity(self, enType=None, enValue=None):
        me = MaltegoEntity(enType, enValue)
        self.entities.append(me)
        return me

    def addUIMessage(self, message, messageType="Inform"):
        self.UIMessages.append([messageType, message])

    def addException(self, exceptionString):
        self.exceptions.append(exceptionString)

    def throwExceptions(self):
        r = ""
        r += "<MaltegoMessage>"
        r += "<MaltegoTransformExceptionMessage>"
        r += "<Exceptions>"

        for i in range(len(self.exceptions)):
            r += "<Exception>" + self.exceptions[i] + "</Exceptions>"
        r += "</Exceptions>"
        r += "</MaltegoTransformExceptionMessage>"
        r += "</MaltegoMessage>"
        return r

    def returnOutput(self):
        r = ""
        r += "<MaltegoMessage>"
        r += "<MaltegoTransformResponseMessage>"
        r += "<Entities>"
        for i in range(len(self.entities)):
            r += self.entities[i].returnEntity()
        r += "</Entities>"
        r += "<UIMessages>"
        for i in range(len(self.UIMessages)):
            r += (
                '<UIMessage MessageType="'
                + self.UIMessages[i][0]
                + '">'
                + self.UIMessages[i][1]
                + "</UIMessage>"
            )
        r += "</UIMessages>"
        r += "</MaltegoTransformResponseMessage>"
        r += "</MaltegoMessage>"
        return r


class MaltegoMsg:
    def __init__(self, MaltegoXML=""):
        xmldoc = minidom.parseString(MaltegoXML)
        # read the easy stuff like value, limits etc
        self.Value = self.i_getNodeValue(xmldoc, "Value")
        self.Weight = self.i_getNodeValue(xmldoc, "Weight")
        self.Slider = int(self.i_getNodeAttributeValue(xmldoc, "Limits", "SoftLimit"))
        self.Type = self.i_getNodeAttributeValue(xmldoc, "Entity", "Type")

        # read additional fields
        Properties = {}
        try:
            AFNodes = xmldoc.getElementsByTagName("AdditionalFields")[0]
            Settings = AFNodes.getElementsByTagName("Field")
            for node in Settings:
                AFName = node.attributes["Name"].value
                AFValue = self.i_getText(node.childNodes)
                Properties[AFName] = AFValue
        except:
            # sure this is not the right way...;)
            dontcare = 1

        # parse transform settings
        TransformSettings = {}
        try:
            TSNodes = xmldoc.getElementsByTagName("TransformFields")[0]
            Settings = TSNodes.getElementsByTagName("Field")
            for node in Settings:
                TSName = node.attributes["Name"].value
                TSValue = self.i_getText(node.childNodes)
                TransformSettings[TSName] = TSValue
        except:
            dontcare = 1

        # load back into object
        self.Properties = Properties
        self.TransformSettings = TransformSettings

    def i_getText(self, nodelist):
        rc = []
        for node in nodelist:
            if node.nodeType == node.TEXT_NODE:
                rc.append(node.data)
        return "".join(rc)

    def i_getNodeValue(self, node, Tag):
        return self.i_getText(node.getElementsByTagName(Tag)[0].childNodes)

    def i_getNodeAttributeValue(self, node, Tag, Attribute):
        return node.getElementsByTagName(Tag)[0].attributes[Attribute].value

    def getProperty(self, skey):
        if skey in self.Properties.keys():
            return self.Properties[skey]
        else:
            return None

    def getTransformSetting(self, skey):
        if skey in self.TransformSettings.keys():
            return self.TransformSettings[skey]
        else:
            return None