Java tutorial
/** * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * * Note: Some libraries linking to this code might not be following the same * licencing policy. */ package com.caspida.plugins.graph; import com.caspida.plugins.graph.datastore.UniqueEntityFactory; import com.google.common.hash.HashCode; import com.google.common.hash.Hashing; import org.json.simple.JSONArray; import org.json.simple.JSONObject; import org.json.simple.parser.JSONParser; import org.neo4j.cypher.ExecutionEngine; import org.neo4j.graphdb.*; import org.neo4j.graphdb.index.Index; import org.neo4j.graphdb.index.UniqueFactory; import org.neo4j.server.plugins.Description; import org.neo4j.server.plugins.Name; import org.neo4j.server.plugins.PluginTarget; import org.neo4j.server.plugins.ServerPlugin; import org.neo4j.server.plugins.Source; import org.neo4j.server.plugins.Parameter; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.util.*; /** * A plugin to create/merge/query graphs with entities (Users, Machines, Applications etc) * This plugin supports creation and/or aggregation of small graphs with * transactional support * * It also supports very specific graph navigation queries for display/navigation * of entity graphs using REST API * * Created by ravib on 5/3/14. */ @Description("Entity graph plug-in for creation, aggregations, rank computation" + " and entity clustering visualizations") public class EntityGraphNeo4jPlugin extends ServerPlugin { private static final Logger logger = LoggerFactory.getLogger(EntityGraphNeo4jPlugin.class); private static final Label eventClassLabel = DynamicLabel.label(EntityGraphConstants.EVENT_CLASS); private static final String UNIQUE_FACTORY_INDEX_NAME = "nodes"; private static final String UNIQUE_FACTORY_PROPERTIES_DATA = "data"; private static UniqueFactory.UniqueNodeFactory uniqueNodeFactory = null; private static UniqueFactory.UniqueRelationshipFactory uniqueRelationshipsFactory = null; private static final String USER_UNIQUE_NODE_MERGE_GQL = String.format("MERGE (n:%s {%s: {data}}) RETURN n", EntityGraphConstants.EVENT_CLASS, EntityGraphConstants.NODE_ID, EntityGraphConstants.NODE_ID); private static synchronized ExecutionEngine getQueryExecutionEngine(GraphDatabaseService graphDb) throws Exception { try (Transaction tx = graphDb.beginTx()) { graphDb.schema().constraintFor(DynamicLabel.label(UNIQUE_FACTORY_INDEX_NAME)) .assertPropertyIsUnique(EntityGraphConstants.NODE_ID).create(); tx.success(); } return new ExecutionEngine(graphDb, null); } /** * Create the generic error * @param jsonObject * @param t */ private static void createErrorResponse(JSONObject jsonObject, String api, String message, Throwable t) { JSONObject error = new JSONObject(); error.put(EntityGraphConstants.ERROR_JSON_TAG_API, api); error.put(EntityGraphConstants.ERROR_JSON_TAG_MESSAGE, message); error.put(EntityGraphConstants.ERROR_JSON_TAG_EXCEPTION, t.getMessage()); jsonObject.put(EntityGraphConstants.ERROR_JSON_TAG, error); } /** * The request JSON structure is * { * "containers" : [ * { * "labels" : ["Red", "Blue"], * "header" : ["id", "label", "<property name 1>", * "<property name2>", ...] * "nodes" : * [ * ["id", "<key>", "<hashcode>", "label", <entropy>, <probability>, <score>, <anomalyType>], * ["id", "<key>", "<hashcode>", "label", <entropy>, <probability>, <score>, <anomalyType>], * ... * ] * } * ] * } * * response JSON * { * "nodeIds" : { * "<hashcode>", <nodeid> * } * } * @param graphDb * @param json * @return */ @Name(EntityGraphConstants.API_MERGE_ENTITY_NODES) @Description(EntityGraphConstants.API_MERGE_ENTITY_NODES) @PluginTarget(GraphDatabaseService.class) public String mergeBehaviorGraphNodes(@Source GraphDatabaseService graphDb, @Description("") @Parameter(name = "data", optional = false) String json) { // Perform input validations if ((json == null) || (json.trim().isEmpty())) { throw new IllegalArgumentException("Invalid JSON provided for method : " + json); } JSONObject response = new JSONObject(); JSONObject nodeIds = new JSONObject(); try { JSONParser parser = new JSONParser(); JSONObject request = (JSONObject) parser.parse(json); JSONArray labels = (JSONArray) request.get(EntityGraphConstants.REQ_LABELS_TAG); JSONArray header = (JSONArray) request.get(EntityGraphConstants.REQ_PROPERTY_HEADER_TAG); JSONArray nodes = (JSONArray) request.get(EntityGraphConstants.REQ_NODE_MERGE_NODES_TAG); if (logger.isDebugEnabled()) { logger.debug("Labels ({}), header ({}), node size ({})", labels, header, nodes.size()); } List<UniqueEntityFactory.NodeRequestInfo> nodeRequestData = new ArrayList<>(); // Navigate with the nodes, because there will always be less nodes // and more edges. For every node go over the relationships and // create the node edges for (int idx = 0; idx < nodes.size(); idx++) { JSONArray nodeData = (JSONArray) nodes.get(idx); Map<String, Object> properties = getPropertiesFromHeaderAndArray(header, nodeData); Long nodeId = null; Object tmp; tmp = properties.get(EntityGraphConstants.NODE_NAME); String nodeKey = (tmp != null) ? tmp.toString() : null; tmp = properties.get(EntityGraphConstants.NODE_ID); String nodeHash = (tmp != null) ? tmp.toString() : null; tmp = properties.remove(EntityGraphConstants.EVENT_DB_OBJECT_ID); String nodeIdObj = (tmp != null) ? tmp.toString() : null; tmp = properties.remove(EntityGraphConstants.REQ_LABEL_TAG); String nodeLabel = (tmp != null) ? tmp.toString() : null; if (logger.isDebugEnabled()) { logger.debug("nodeKey ({}), nodeHash ({}), DB object id ({}), label ({})", nodeKey, nodeHash, nodeIdObj, nodeLabel); } if (nodeIdObj != null) { try { nodeId = Long.parseLong(nodeIdObj); } catch (NumberFormatException nfe) { throw new IllegalArgumentException("Invalid nodeId in request : " + nodeIdObj); } } Node eventNode; if ((nodeId != null) && (nodeId >= 0)) { try (Transaction txn = graphDb.beginTx()) { eventNode = graphDb.getNodeById(nodeId); nodeIds.put(nodeHash, Long.toString(eventNode.getId())); } } else { List<String> nodeLabels = new ArrayList<>(); if ((nodeLabel != null) && (!nodeLabel.toString().trim().isEmpty())) { nodeLabels.add(nodeLabel.toString()); } for (Object label : labels) { nodeLabels.add(label.toString()); } nodeRequestData.add(new UniqueEntityFactory.NodeRequestInfo(nodeHash, properties, nodeLabels)); } } Map<String, Long> resultIds = UniqueEntityFactory.createBulkNodes(graphDb, nodeRequestData); if (resultIds != null) { for (Map.Entry<String, Long> entry : resultIds.entrySet()) { nodeIds.put(entry.getKey(), entry.getValue()); } } //numNodes += nodes.size(); response.put(EntityGraphConstants.RES_IDS_TAG, nodeIds); } catch (Exception e) { logger.error("Failed to create entity graph", e); createErrorResponse(response, EntityGraphConstants.API_MERGE_ENTITY_NODES, "Failed to create entity graph", e); } return response.toJSONString(); } private Map<String, Object> getPropertiesFromHeaderAndArray(JSONArray header, JSONArray data) throws Exception { if (header.size() != data.size()) { throw new Exception("Header and data not of same size : Header {" + header + "}, data {" + data + "}"); } Map<String, Object> props = new HashMap<>(); for (int idx = 0; idx < header.size(); idx++) { props.put(header.get(idx).toString(), data.get(idx).toString()); } return props; } /** * Nodes and edges are merged in two different API calls because node creation * returns the nodes hashes to nodeId map * * Adding nodeId to the edge request speeds up the process of edge creation * because the lookUp of nodes by nodeId is much faster then going with the index * * { * containers : [ * { * "labels" : ["a", "b"], * "header" : [<Array of properties names. All edge rels array will * follow the same order>] * "relmap" : { * "hash1" : { * "id" : <e1>, * "rels" : [ * [<hash2>, <e2>, <relationship>, <conditional probability>, * <joint probability>, <mutual information>] * ], * } * ... * } * } * ] * } * @param graphDb * @param json * @return */ @Name(EntityGraphConstants.API_MERGE_ENTITY_EDGES) @Description("Merge entity edges") @PluginTarget(GraphDatabaseService.class) public String mergeBehaviorGraphEdges(@Source GraphDatabaseService graphDb, @Description("") @Parameter(name = "data", optional = true) String json) { // Perform request validations if ((json == null) || (json.trim().isEmpty())) { throw new IllegalArgumentException("Invalid JSON provided for method : " + json); } logger.debug("Merge edges request data : {}", json); JSONObject relIds = new JSONObject(); long numEdges = 0; JSONObject response = new JSONObject(); try { JSONParser parser = new JSONParser(); JSONObject request = (JSONObject) parser.parse(json); JSONArray header = (JSONArray) request.get(EntityGraphConstants.REQ_PROPERTY_HEADER_TAG); JSONObject relationsMap = (JSONObject) request .get(EntityGraphConstants.REQ_EDGE_MERGE_RELATIONSHIP_MAP_TAG); Map<Long, Map<String, Relationship>> eventRelationships = new HashMap<>(); for (Object key : relationsMap.keySet()) { eventRelationships.clear(); String eventHash = key.toString(); JSONObject data = (JSONObject) relationsMap.get(key); Long eventId = null; Node e1 = null; if ((eventId == null) || (eventId < 0)) { e1 = UniqueEntityFactory.getOrCreateNode(graphDb, eventHash, null); } else { try (Transaction txn = graphDb.beginTx()) { e1 = graphDb.getNodeById(eventId); } } List<UniqueEntityFactory.RelationshipRequestInfo> requestInfoList = new ArrayList<>(); JSONArray relationships = (JSONArray) data .get(EntityGraphConstants.REQ_EDGE_MERGE_RELATIONSHIPS_TAG); for (Object objectRel : relationships) { JSONArray rel = (JSONArray) objectRel; Map<String, Object> edgeProperties = getPropertiesFromHeaderAndArray(header, rel); Object tmp; tmp = edgeProperties.remove(EntityGraphConstants.REQ_EDGE_MERGE_END_NODE_HASH_TAG); String endNodeEventHash = (tmp != null) ? tmp.toString() : null; Long endNodeEventId = null; tmp = edgeProperties.remove(EntityGraphConstants.REQ_EDGE_MERGE_END_DB_ID_TAG); String relEventIdObj = (tmp != null) ? tmp.toString() : null; try { endNodeEventId = (relEventIdObj != null) ? Long.parseLong(relEventIdObj.toString()) : null; } catch (NumberFormatException e) { // ignore exception and do a lookup } tmp = edgeProperties.remove(EntityGraphConstants.REQ_EDGE_PROPS_RELNAME_KEY); String relationshipName = (tmp != null) ? tmp.toString() : null; requestInfoList.add(new UniqueEntityFactory.RelationshipRequestInfo(e1, endNodeEventId, endNodeEventHash, relationshipName, edgeProperties)); } UniqueEntityFactory.createBulkRelationships(graphDb, e1, requestInfoList); for (UniqueEntityFactory.RelationshipRequestInfo rri : requestInfoList) { String relKey = rri.getDescr(); Relationship relationship = rri.getRelationship(); if ((relKey == null) || (relationship == null)) { continue; } relIds.put(relKey, relationship.getId()); } numEdges += requestInfoList.size(); } response.put(EntityGraphConstants.RES_IDS_TAG, relIds); } catch (Throwable e) { logger.error("Failed to create entity graph : " + json, e); createErrorResponse(response, EntityGraphConstants.API_MERGE_ENTITY_EDGES, "Failed to merge entity graph", e); } return response.toJSONString(); } /** * request = {"start" : {"type":"User", "value":"userId"}, * "end" : {"type":"Device", * "values": ["deviceId1", "deviceId2", "deviceId3"]}, * "relationship" : "ConnectsTo", * "propertyName" : "<some property name>"} * * response = {"edges": {"deviceId1": ["<property value>", "<nodeId>"], * "deviceId2": ["<property value>", "<nodeId>"], * "deviceId3": ["<property value>", "<nodeId>"] * } * } * * @param graphDb * @param json * @return */ @Name(EntityGraphConstants.API_GET_NEIGHBORHOOD_PROPERTY) @Description("Get properties of the connected node and the edge if it exists") @PluginTarget(GraphDatabaseService.class) public String getNeighborNodeAndEdgeProperty(@Source GraphDatabaseService graphDb, @Description("Request JSON") @Parameter(name = "data", optional = true) String json) { JSONObject response = new JSONObject(); try { JSONParser parser = new JSONParser(); JSONObject request = (JSONObject) parser.parse(json); Object tmp = request.get(EntityGraphConstants.REQ_EDGE_PROPS_START_KEY); if (tmp == null) { return response.toJSONString(); } JSONObject start = (JSONObject) tmp; tmp = request.get(EntityGraphConstants.REQ_EDGE_PROPS_END_KEY); if (tmp == null) { return response.toJSONString(); } JSONObject end = (JSONObject) tmp; tmp = start.get(EntityGraphConstants.REQ_EDGE_PROPS_START_VALUE_KEY); if (tmp == null) { return response.toJSONString(); } String startNodeId = tmp.toString(); tmp = end.get(EntityGraphConstants.REQ_EDGE_PROPS_END_VALUES_KEY); if (tmp == null) { return response.toJSONString(); } JSONArray endValues = (JSONArray) tmp; if (endValues.size() == 0) { return response.toJSONString(); } Set<String> endNodeIds = new HashSet<>(); for (Object endValue : endValues) { endNodeIds.add(endValue.toString()); } tmp = request.get(EntityGraphConstants.REQ_EDGE_PROPS_RELNAME_KEY); if (tmp == null) { return response.toJSONString(); } DynamicRelationshipType relName = DynamicRelationshipType.withName(tmp.toString()); tmp = request.get(EntityGraphConstants.REQ_EDGE_PROPS_PROPNAME_KEY); if (tmp == null) { return response.toJSONString(); } String propertyName = tmp.toString(); response.clear(); JSONObject responseEdges = new JSONObject(); // get all thr relationships try (Transaction txn = graphDb.beginTx()) { Index<Node> nodeIndex = graphDb.index().forNodes(EntityGraphConstants.EVENT_CLASS); Node node = findNode(graphDb, startNodeId, nodeIndex); if (logger.isDebugEnabled()) { logger.debug("Found node with id {} => node id", startNodeId, (node != null) ? node.getId() : "-1"); } if (node != null) { Iterable<Relationship> rels = node.getRelationships(relName); for (Relationship rel : rels) { Node otherNode = (rel.getStartNode().getId() == node.getId()) ? rel.getEndNode() : rel.getStartNode(); Object value = null; Object otherNodeId = null; try { otherNodeId = otherNode.getProperty(EntityGraphConstants.NODE_ID); } catch (Exception e) { } if (otherNodeId == null) { logger.error("Node without a ID property {}", otherNode.getId()); continue; } // if the JSONArray contains this node if (!endNodeIds.remove(otherNodeId.toString())) { // not interested all edges continue; } value = null; try { value = rel.getProperty(propertyName); } catch (Exception e) { } JSONArray dat = new JSONArray(); dat.add(0, (value != null) ? value.toString() : "1.0"); dat.add(1, otherNode.getId()); responseEdges.put(otherNodeId, dat); } } // for the remaining end nodes that were not part of the // relationships find the node id and return for (String endNodeId : endNodeIds) { Node endNode = findNode(graphDb, endNodeId, nodeIndex); JSONArray dat = new JSONArray(); dat.add(0, ""); dat.add(1, (endNode != null) ? endNode.getId() : "-1"); responseEdges.put(endNodeId, dat); } response.put(EntityGraphConstants.RES_EDGE_PROPS_EDGES_KEY, responseEdges); } } catch (Throwable e) { logger.error("Failed to create entity graph : " + json, e); createErrorResponse(response, EntityGraphConstants.API_GET_NEIGHBORHOOD_PROPERTY, "Failed to create entity graph", e); } if (logger.isDebugEnabled()) { logger.debug("Response : {}", response.toJSONString()); } return response.toJSONString(); } public static final String GRAPH_NAMESPACE_SEPARATOR = "|"; public static final String GRAPH_NAMESPACE_SPLIT_FORMAT = "\\|"; public static String getFormattedString(String... tokens) { StringBuilder sb = new StringBuilder(); boolean isFirst = true; for (String eventAttribute : tokens) { if (!isFirst) { sb.append(GRAPH_NAMESPACE_SEPARATOR); } sb.append(eventAttribute); isFirst = false; } return sb.toString(); } public static String[] getTokensFromFormattedString(String formattedString) { return formattedString.split(GRAPH_NAMESPACE_SPLIT_FORMAT); } public static Long hash(String text) { HashCode hc = Hashing.murmur3_128().hashBytes(text.getBytes()); return hc.asLong(); } /** * request = [ * user1, user2, user3 * ] * * response = { * "user1" : [device1, device2, device3] * "user2" : [device1, device5, device6] * } * * @param graphDb * @param json * @return */ @Name(EntityGraphConstants.API_GET_ALL_DEVICES_FOR_USERS) @Description("Get top machines for user") @PluginTarget(GraphDatabaseService.class) public String getAllDevicesForUser(@Source GraphDatabaseService graphDb, @Description("Array of users as a JSON string ") @Parameter(name = "data", optional = true) String json) { JSONObject response = new JSONObject(); try { JSONParser parser = new JSONParser(); JSONArray request = (JSONArray) parser.parse(json); String namespace = getFormattedString("Network", EntityGraphConstants.UBER_ENTITY_GRAPH_NAMESPACE); for (Object userObject : request) { List<String> entities = this.getUberGraphConnectedEntities(graphDb, namespace, userObject.toString(), "ConnectsTo", EntityGraphConstants.NODE_NAME); response.put(userObject, entities); } } catch (Exception e) { logger.error("Failed to get top machines", e); createErrorResponse(response, EntityGraphConstants.API_GET_ALL_DEVICES_FOR_USERS, "Failed to get top machines", e); } return response.toJSONString(); } /** * request = [ * user1, user2, user3 * ] * * response = { * "user1" : [app1, app2, app3] * "user2" : [app9, app4, app6] * } * * @param graphDb * @param json * @return */ @Name(EntityGraphConstants.API_GET_ALL_APPS_FOR_USERS) @Description("Get top applications for user") @PluginTarget(GraphDatabaseService.class) public String getAllApplicationsForUser(@Source GraphDatabaseService graphDb, @Description("Array of users as a JSON string ") @Parameter(name = "data", optional = true) String json) { JSONObject response = new JSONObject(); try { JSONParser parser = new JSONParser(); JSONArray request = (JSONArray) parser.parse(json); String namespace = getFormattedString("Network", EntityGraphConstants.UBER_ENTITY_GRAPH_NAMESPACE); for (Object userObject : request) { List<String> entities = this.getUberGraphConnectedEntities(graphDb, namespace, userObject.toString(), "Uses", EntityGraphConstants.NODE_NAME); response.put(userObject, entities); } } catch (Exception e) { logger.error("Failed to get top apps for user", e); createErrorResponse(response, EntityGraphConstants.API_GET_ALL_APPS_FOR_USERS, "Failed to get top apps for user", e); } return response.toJSONString(); } /** * request = [ * user1, user2, user3 * ] * * response = { * "user1" : [app1, app2, app3] * "user2" : [app9, app4, app6] * } * * @param graphDb * @param json * @return */ @Name(EntityGraphConstants.API_GET_ALL_THREATS_FOR_USERS) @Description("Get top threats for user") @PluginTarget(GraphDatabaseService.class) public String getAllThreatsForUser(@Source GraphDatabaseService graphDb, @Description("Array of users as a JSON string ") @Parameter(name = "data", optional = true) String json) { JSONObject response = new JSONObject(); try { JSONParser parser = new JSONParser(); JSONArray request = (JSONArray) parser.parse(json); String namespace = getFormattedString("Network", EntityGraphConstants.UBER_ENTITY_GRAPH_NAMESPACE); for (Object userObject : request) { List<String> entities = this.getUberGraphConnectedEntities(graphDb, namespace, userObject.toString(), "Anomaly", EntityGraphConstants.NODE_NAME); response.put(userObject, entities); } } catch (Exception e) { logger.error("Failed to get top threats for user", e); createErrorResponse(response, EntityGraphConstants.API_GET_ALL_THREATS_FOR_USERS, "Failed to get top threats for user", e); } return response.toJSONString(); } /* if a node is an anomaly, get its score, if not return -1 */ private Integer getNodeScore(Node node) { /* figure out the type of the node */ Object nodeTypeObject = node.getProperty(EntityGraphConstants.EVENT_TYPE); if (nodeTypeObject == null) { return -1; } String nodeType = nodeTypeObject.toString(); /* check if we are below the minScore so we stop here */ if (nodeType.equalsIgnoreCase("Threat")) { Integer score = Integer.parseInt(node.getProperty(EntityGraphConstants.ANOMALY_SCORE).toString()); return score; } return -1; } private String getNodeType(Node node) { Object tmp = node.getProperty(EntityGraphConstants.EVENT_TYPE); return (tmp != null) ? tmp.toString() : null; } private Boolean shouldInclude(Node node, List<String> entityList) { /* figure out the type of the node */ String otherObjectType = getNodeType(node); if (otherObjectType == null) { return false; } logger.error("Other object type is {}", otherObjectType); // skip the node if is not in the list of types we are interested in // can we fix this so the entities have the right value in their tag switch (otherObjectType) { case "User": if (!entityList.contains(EntityGraphConstants.UBER_GRAPH_USER_LABEL)) { logger.error("Dropping other node, not one of the requested types"); return false; } break; case "Device": if (!entityList.contains(EntityGraphConstants.UBER_GRAPH_DEVICE_LABEL)) { logger.error("Dropping other node, not one of the requested types"); return false; } break; case "Application": if (!entityList.contains(EntityGraphConstants.UBER_GRAPH_APPLICATION_LABEL)) { logger.error("Dropping other node, not one of the requested types"); return false; } break; case "Threat": if (!entityList.contains(EntityGraphConstants.UBER_GRAPH_THREAT_LABEL)) { logger.error("Dropping other node, not one of the requested types"); return false; } break; case "Url": if (!entityList.contains(EntityGraphConstants.UBER_GRAPH_URL_LABEL)) { logger.error("Dropping other node, not one of the requested types"); return false; } break; default: logger.error("Invalid entity type {}", otherObjectType); return false; } return true; } @Name(EntityGraphConstants.API_GET_BIPARTITE_GRAPH) @Description("Get bipartite information for various entities") @PluginTarget(GraphDatabaseService.class) public String getBipartite(@Source GraphDatabaseService graphDb, @Description("Input parameters as a JSON string ") @Parameter(name = "entityTypes") List<String> entityList, @Parameter(name = "sourceEntityId", optional = true) String entityId, @Parameter(name = "sourceEntityType") String entityType, @Parameter(name = "minScore", optional = true) Integer minScoreIn) { JSONArray nodes = new JSONArray(); JSONArray edges = new JSONArray(); JSONObject response = new JSONObject(); JSONParser parser = new JSONParser(); Integer minScore = (minScoreIn != null) ? minScoreIn.intValue() : 0; Integer uniqueLinkId = 1; /* keep track if the caller requested to see anomalies, or not * if it did not, we have two options on what we will show, facts or only anomaly related links * if the caller asked for anomalies, then we will only show anomaly related links */ Boolean askedForAnomalies = ((entityType == EntityGraphConstants.UBER_GRAPH_THREAT_LABEL) || entityList.contains(EntityGraphConstants.UBER_GRAPH_THREAT_LABEL)) ? true : false; logger.error("Got input params entity list {} and entityId {} asked for anomalies {}", entityList, entityId, askedForAnomalies); try (Transaction txn = graphDb.beginTx()) { /* build a list of nodes that we will later expand * if entityId is present then it will be a single node * else we will need to pick a set of nodes from the entity types * specified. Should pick some nodes with low selectivity * (i.e. anomalies) if possible */ List<Node> nodeList = new ArrayList<>(); if (entityId != null && entityId.length() > 2) { // XXX better test for "" string logger.error("Looking up single node with id {}", entityId); Index<Node> nodeIndex = graphDb.index().forNodes(EntityGraphConstants.EVENT_CLASS); Node node = nodeIndex.get(EntityGraphConstants.NODE_ID, Long.parseLong(entityId)).getSingle(); if (node == null) { logger.error("Can not locate entity with provided id {}", entityId); return response.toJSONString(); } else { logger.debug("Found entity node with provided id {}", entityId); nodeList.add(node); } } else { /* map the entityType to the right label */ Label label = DynamicLabel.label(entityType); // XXX String eventType = null; /* map the entity id to the value I need for the event type */ switch (entityType) { case "caspida.device": eventType = "Device"; break; case "caspida.anomaly": eventType = "Threat"; break; case "caspida.user": eventType = "User"; break; case "caspida.url": eventType = "Url"; break; default: break; } logger.error("Got event type {}", eventType); /** XXX use the exact label to reduce the scope of the walk */ ResourceIterable<Node> iterableNodes = null; if (eventType != null) { iterableNodes = graphDb.findNodesByLabelAndProperty(eventClassLabel, EntityGraphConstants.EVENT_TYPE, eventType); } if (iterableNodes != null) { ResourceIterator<Node> it = iterableNodes.iterator(); while (it.hasNext()) { Node node = it.next(); logger.error("Adding node {} to the list", node.getId()); nodeList.add(node); } } else { logger.error("Could not find any nodes for request with entity list {}", entityList); return response.toJSONString(); } } /* at this point I have a list of nodes that I need to expand * need to make sure I do not repeat the same nodes, so use a * hash map to keep track of the ones I have visited so far */ HashMap<Node, Integer> nodeHash = new HashMap<>(); for (Node node : nodeList) { String objectType = null; String otherObjectType = null; Iterable<Relationship> rels = node.getRelationships(); /* if the current node is anomaly, skip depending on the score */ Integer score = getNodeScore(node); if (score != -1 && score < minScore) { logger.error("Dropping node due to low score"); continue; } /* now iterate all the neighbhors and put the ones to be expanded in a list * need to filter anomalies for the minScore * if we hit anomaly nodes and askedForAnomalies is false we need to "go over" them * to get to the entity on the other side * */ nodeHash.put(node, 1); List<Node> otherNodeList = new ArrayList<Node>(); for (Relationship rel : rels) { Node otherNode = (rel.getStartNode().getId() == node.getId()) ? rel.getEndNode() : rel.getStartNode(); /* if the other node is anomaly, skip depending on the score */ score = getNodeScore(otherNode); if (score != -1 && score < minScore) { logger.error("Dropping other node due to low score"); continue; } if (score != -1) { /* it is an anomaly */ if (askedForAnomalies == false) { /* need to expand the anomalies and get to its children */ Iterable<Relationship> anomalyRels = otherNode.getRelationships(); for (Relationship anomalyRel : anomalyRels) { Node otherAnomalyNode = (anomalyRel.getStartNode().getId() == otherNode.getId()) ? anomalyRel.getEndNode() : anomalyRel.getStartNode(); if (otherAnomalyNode != node && shouldInclude(otherAnomalyNode, entityList) == true && !nodeHash.containsKey(otherAnomalyNode)) { otherNodeList.add(otherAnomalyNode); nodeHash.put(otherAnomalyNode, 1); } } } else { if (shouldInclude(otherNode, entityList) == true && !nodeHash.containsKey(otherNode)) { otherNodeList.add(otherNode); nodeHash.put(otherNode, 1); } } } else { if (shouldInclude(otherNode, entityList) == true && !nodeHash.containsKey(otherNode)) { otherNodeList.add(otherNode); nodeHash.put(otherNode, 1); } } } /* I have a list of other nodes to expand now, if the list is not empty * output the local node first and then the links and other nodes * if i am going to skip remove the node from the nodeHash */ if (otherNodeList.size() == 0) { nodeHash.remove(node); continue; } logger.error("Other node list {}", otherNodeList); /* add information about the current node */ nodeHash.put(node, 1); JSONObject nodeInfo = new JSONObject(); String[] myTokens = getTokensFromFormattedString( node.getProperty(EntityGraphConstants.NODE_NAME).toString()); nodeInfo.put("nodeId", node.getId()); nodeInfo.put("entityId", node.getProperty(EntityGraphConstants.NODE_ID)); String typeStr = null; objectType = getNodeType(node); switch (objectType) { case "User": typeStr = EntityGraphConstants.UBER_GRAPH_USER_LABEL; nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, myTokens[myTokens.length - 1]); break; case "Device": Object deviceTypeObj = node.getProperty(EntityGraphConstants.DEVICE_TYPE); String deviceType = (deviceTypeObj != null) ? deviceTypeObj.toString() : "Desktop"; typeStr = EntityGraphConstants.UBER_GRAPH_DEVICE_LABEL; nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, myTokens[myTokens.length - 1]); nodeInfo.put("entitySubtype", deviceType); break; case "Application": typeStr = EntityGraphConstants.UBER_GRAPH_APPLICATION_LABEL; nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, myTokens[myTokens.length - 1]); break; case "Threat": typeStr = EntityGraphConstants.UBER_GRAPH_THREAT_LABEL; /* anomalies do not have a meaningful label */ nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, "Anomaly"); nodeInfo.put("entitySubtype", node.getProperty(EntityGraphConstants.ANOMALY_TYPE)); break; case "Url": typeStr = EntityGraphConstants.UBER_GRAPH_URL_LABEL; nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, myTokens[myTokens.length - 1]); default: break; } nodeInfo.put("entityType", typeStr); nodes.add(nodeInfo); for (Node otherNode : otherNodeList) { /* add information about the relationship */ JSONObject edgeInfo = new JSONObject(); edgeInfo.put(EntityGraphConstants.RES_SOURCE_ID_TAG, node.getId()); edgeInfo.put(EntityGraphConstants.RES_TARGET_ID_TAG, otherNode.getId()); edgeInfo.put(EntityGraphConstants.RES_LINK_ID_TAG, uniqueLinkId); uniqueLinkId += 1; edges.add(edgeInfo); /* add information about the other node */ nodeInfo = new JSONObject(); myTokens = getTokensFromFormattedString( otherNode.getProperty(EntityGraphConstants.NODE_NAME).toString()); nodeInfo.put("nodeId", otherNode.getId()); nodeInfo.put("entityId", otherNode.getProperty(EntityGraphConstants.NODE_ID)); typeStr = null; otherObjectType = getNodeType(otherNode); switch (otherObjectType) { case "User": typeStr = EntityGraphConstants.UBER_GRAPH_USER_LABEL; nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, myTokens[myTokens.length - 1]); break; case "Device": Object deviceTypeObj = otherNode.getProperty(EntityGraphConstants.DEVICE_TYPE); String deviceType = (deviceTypeObj != null) ? deviceTypeObj.toString() : "Desktop"; typeStr = EntityGraphConstants.UBER_GRAPH_DEVICE_LABEL; nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, myTokens[myTokens.length - 1]); nodeInfo.put("entitySubtype", deviceType); break; case "Application": typeStr = EntityGraphConstants.UBER_GRAPH_APPLICATION_LABEL; nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, myTokens[myTokens.length - 1]); break; case "Threat": typeStr = EntityGraphConstants.UBER_GRAPH_THREAT_LABEL; /* anomalies do not have a meaningful label */ nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, "Anomaly"); nodeInfo.put("entitySubtype", node.getProperty(EntityGraphConstants.ANOMALY_TYPE)); break; case "Url": typeStr = EntityGraphConstants.UBER_GRAPH_URL_LABEL; nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, myTokens[myTokens.length - 1]); default: break; } nodeInfo.put("entityType", typeStr); nodes.add(nodeInfo); } } } response.put("nodes", nodes); response.put("links", edges); if (logger.isDebugEnabled()) { logger.debug("Response : {}", response.toJSONString()); } return response.toString(); } @Name(EntityGraphConstants.API_GET_ENTITY_PST) @Description("Get an entity's PST model") @PluginTarget(GraphDatabaseService.class) public String getEntityPst(@Source GraphDatabaseService graphDb, @Description("entity id and entity type") @Parameter(name = "entityId") String entityId, @Parameter(name = "entityType") String entityType) { JSONArray nodes = new JSONArray(); Node node = null; Label pstLabel = DynamicLabel.label("PST___Root"); logger.error("Got entity id {}", entityId); String property = EntityGraphConstants.REQ_USER_ID_TAG; logger.error("Got property {}", property); try (Transaction txn = graphDb.beginTx()) { // find the node // use the label first ResourceIterable<Node> iterableNodes = graphDb.findNodesByLabelAndProperty(pstLabel, property, entityId); if (iterableNodes != null) { ResourceIterator<Node> it = iterableNodes.iterator(); if (it.hasNext()) { node = it.next(); } } logger.debug("Found node {}", node); if (node == null) { logger.error("Can not find node with id {}", entityId); return nodes.toString(); } Iterable<Relationship> rels = node.getRelationships(); for (Relationship rel : rels) { Node childNode = (rel.getStartNode().getId() == node.getId()) ? rel.getEndNode() : rel.getStartNode(); Object key = childNode.getProperty(EntityGraphConstants.NODE_NAME); if (key == null) { continue; } JSONObject childInfo = walkPSTGraph(childNode, node, rel); if (childInfo != null) { nodes.add(childInfo); } } } return nodes.toJSONString(); } @Name(EntityGraphConstants.API_GET_USER_PST) @Description("Get a users's PST model") @PluginTarget(GraphDatabaseService.class) public String getUserPst(@Source GraphDatabaseService graphDb, @Description("user id ") @Parameter(name = "userId") String userId) { JSONArray nodes = new JSONArray(); Node node = null; Label pstLabel = DynamicLabel.label("PST___Root"); logger.error("Got user id {}", userId); try (Transaction txn = graphDb.beginTx()) { // find the node // use the label first ResourceIterable<Node> iterableNodes = graphDb.findNodesByLabelAndProperty(pstLabel, EntityGraphConstants.REQ_USER_ID_TAG, userId); if (iterableNodes != null) { ResourceIterator<Node> it = iterableNodes.iterator(); if (it.hasNext()) { node = it.next(); } } logger.debug("Found node {}", node); if (node == null) { logger.error("Can not find node with id {}", userId); return nodes.toString(); } Iterable<Relationship> rels = node.getRelationships(); for (Relationship rel : rels) { Node childNode = (rel.getStartNode().getId() == node.getId()) ? rel.getEndNode() : rel.getStartNode(); Object key = childNode.getProperty(EntityGraphConstants.NODE_NAME); if (key == null) { continue; } JSONObject childInfo = walkPSTGraph(childNode, node, rel); if (childInfo != null) { nodes.add(childInfo); } } } return nodes.toJSONString(); } /* iterate all the nodes up to depth and * output nodes and edges in the provided JSON arrays */ private Integer walkNbrGraph(Node node, Node parentNode, Integer depth, JSONArray nodes, JSONArray edges, Integer minScore) { Integer childrenCount = 0; Integer nodesAdded = 0, added; /* figure out the type of the node */ String nodeType = node.getProperty(EntityGraphConstants.EVENT_TYPE).toString(); if (nodeType == null) { return nodesAdded; } /* check if we are below the minScore so we stop here */ if (nodeType.equalsIgnoreCase("Threat")) { Integer score = Integer.parseInt(node.getProperty(EntityGraphConstants.ANOMALY_SCORE).toString()); if (score < minScore) { return nodesAdded; } } /* iterate children, do that even if depth is 0 to * set properly the hasChildren */ if (depth >= 0) { Iterable<Relationship> rels = node.getRelationships(); for (Relationship rel : rels) { Node childNode = (rel.getStartNode().getId() == node.getId()) ? rel.getEndNode() : rel.getStartNode(); Object key = childNode.getProperty(EntityGraphConstants.NODE_NAME); /* make sure we skip relationships that point back to our parent */ if (parentNode != null && childNode.getId() == parentNode.getId()) { continue; } if (key == null) { continue; } childrenCount += 1; if (depth > 0) { added = walkNbrGraph(childNode, node, depth - 1, nodes, edges, minScore); if (added > 0) { nodesAdded += added; /* add information about the relationship */ JSONObject edgeInfo = new JSONObject(); edgeInfo.put(EntityGraphConstants.RES_SOURCE_ID_TAG, node.getId()); edgeInfo.put(EntityGraphConstants.RES_TARGET_ID_TAG, childNode.getId()); edgeInfo.put(EntityGraphConstants.RES_WEIGHT_TAG, "2"); // XXX missing edgeInfo.put(EntityGraphConstants.RES_LINK_ID_TAG, rel.getId()); edges.add(edgeInfo); } } } } /* add information about the node */ JSONObject nodeInfo = new JSONObject(); String[] myTokens = getTokensFromFormattedString( node.getProperty(EntityGraphConstants.NODE_NAME).toString()); nodeInfo.put(EntityGraphConstants.RES_NODE_ID_TAG, node.getId()); nodeInfo.put(EntityGraphConstants.RES_ANNOTATION_TAG, "looks good"); nodeInfo.put(EntityGraphConstants.RES_CHILDREN_COUNT_TAG, childrenCount); nodeInfo.put("name", myTokens[myTokens.length - 1]); nodeInfo.put(EntityGraphConstants.RES_OBJECT_ID_TAG, node.getProperty(EntityGraphConstants.NODE_ID)); String typeStr = null; switch (nodeType) { case "User": typeStr = EntityGraphConstants.UBER_GRAPH_USER_LABEL; String entityId = node.getProperty(EntityGraphConstants.NODE_ID).toString(); nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, getUserName(entityId, myTokens[myTokens.length - 1])); break; case "Url": typeStr = EntityGraphConstants.UBER_GRAPH_URL_LABEL; nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, myTokens[myTokens.length - 1]); break; case "Device": Object deviceTypeObj = node.getProperty(EntityGraphConstants.DEVICE_TYPE); String deviceType = (deviceTypeObj != null) ? deviceTypeObj.toString() : "Desktop"; typeStr = EntityGraphConstants.UBER_GRAPH_DEVICE_LABEL; nodeInfo.put(EntityGraphConstants.DEVICE_TYPE, deviceType); nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, myTokens[myTokens.length - 1]); break; case "Application": typeStr = EntityGraphConstants.UBER_GRAPH_APPLICATION_LABEL; nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, myTokens[myTokens.length - 1]); break; case "Threat": typeStr = EntityGraphConstants.UBER_GRAPH_THREAT_LABEL; nodeInfo.put(EntityGraphConstants.ANOMALY_TYPE, node.getProperty(EntityGraphConstants.ANOMALY_TYPE)); /* anomalies do not have a meaningful label */ nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, "Anomaly"); nodeInfo.put(EntityGraphConstants.RES_SCORE_TAG, node.getProperty(EntityGraphConstants.ANOMALY_SCORE).toString()); break; default: break; } nodeInfo.put(EntityGraphConstants.RES_TYPE_TAG, typeStr); nodes.add(nodeInfo); nodesAdded += 1; return nodesAdded; } /** * Method resolves userName from the id * @param id * @param token * @return */ private String getUserName(String id, String token) { try { // Ram: We can't have IR lookup here. // Returning the ID for now. //String userName = EntityLookup.getUserName(id); //logger.debug("ID="+id+" Token="+token+" NAME="+userName); //return userName; return token; } catch (Exception e) { logger.error("failed to lookup userName: ", e); //in case of error, return token return token; } } /* iterate all the nodes and output all a json object that contains * all the nodes below the current root */ private JSONObject walkPSTGraph(Node node, Node parentNode, Relationship parentRel) { Boolean hasChildren = false; JSONArray childrenList = new JSONArray(); /* add information about the node */ JSONObject nodeInfo = new JSONObject(); logger.debug("Visiting node {}", node.getId()); if (parentRel != null) { nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, parentRel.getType().name()); nodeInfo.put(EntityGraphConstants.RES_PROBABILITY_TAG, Float.parseFloat(parentRel.getProperty(EntityGraphConstants.RES_PROBABILITY_TAG).toString())); nodeInfo.put(EntityGraphConstants.RES_DELAY_TAG, parentRel.getProperty(EntityGraphConstants.RES_DELAY_TAG)); nodeInfo.put(EntityGraphConstants.RES_COUNT_TAG, Integer.parseInt(parentRel.getProperty(EntityGraphConstants.RES_COUNT_TAG).toString())); } else { nodeInfo.put(EntityGraphConstants.RES_LABEL_TAG, "root"); nodeInfo.put(EntityGraphConstants.RES_PROBABILITY_TAG, 0.0); nodeInfo.put(EntityGraphConstants.RES_DELAY_TAG, 0); nodeInfo.put(EntityGraphConstants.RES_COUNT_TAG, 0); } Iterable<Relationship> rels = node.getRelationships(); for (Relationship rel : rels) { Node childNode = (rel.getStartNode().getId() == node.getId()) ? rel.getEndNode() : rel.getStartNode(); Object key = childNode.getProperty(EntityGraphConstants.NODE_NAME); if (key == null) { continue; } /* make sure we skip relationships that point back to our parent */ if (parentNode != null && childNode.getId() == parentNode.getId()) { continue; } hasChildren = true; JSONObject childInfo = walkPSTGraph(childNode, node, rel); if (childInfo != null) { childrenList.add(childInfo); } } if (hasChildren) { nodeInfo.put(EntityGraphConstants.RES_CHILDREN_TAG, childrenList); } return nodeInfo; } /** * request = * sourceType: User * sourceId: 123 * targetTypes: ["Anomaly", "Device", "Application"] * depth: 1, 2 etc * * response = * nodes [] * nodeId: unique node id * type: Anomaly , User, etc * objectId: Anomaly, User, Device or Application id (the Mysql one) * label: user name, etc * score * properties: depending on the type. * Device type should have a "deviceType" property (Laptop, Handset etc); * Anomaly type should have "anomalyType" property. * annotation * hasChildren: boolean, indicates whether the node is expandable * * links [] * linkId: unique link id * sourceId: source node id * targetId: target node id * weight: a value from 1 to 10, indicates how strong the relation is (most used machines vs occasional etc) * annotation * * @param graphDb * @param sourceType * @param depthIn * @param sourceId * @return */ @Name("getneighborhood") @Description("Get a node's neighborhood") @PluginTarget(GraphDatabaseService.class) public String getUserNeighborhood(@Source GraphDatabaseService graphDb, @Description("Array of users as a JSON string ") @Parameter(name = "sourceType") String sourceType, @Parameter(name = "depth") Integer depthIn, @Parameter(name = "minScore", optional = true) Integer minScoreIn, @Parameter(name = "sourceId") String sourceId ) { JSONArray nodes = new JSONArray(); JSONArray edges = new JSONArray(); JSONObject response = new JSONObject(); Integer minScore = (minScoreIn != null) ? minScoreIn.intValue() : 0; Integer depth = (depthIn != null) ? depthIn.intValue() : 0; String namespace = EntityGraphConstants.UBER_ENTITY_GRAPH_NAMESPACE; logger.info("Got arguments: sourceType {} depth {} sourceId {} namespace {}", sourceType, depth, sourceId, namespace); try (Transaction txn = graphDb.beginTx()) { Index<Node> nodeIndex = graphDb.index().forNodes(EntityGraphConstants.EVENT_CLASS); Node node = nodeIndex.get(EntityGraphConstants.NODE_ID, Long.parseLong(sourceId)).getSingle(); if (node != null && depth != 0) { walkNbrGraph(node, null, depth, nodes, edges, minScore); } } response.put("nodes", nodes); response.put("links", edges); logger.debug("Reponse : {}", response.toJSONString()); return response.toString(); } private List<String> getUberGraphConnectedEntities(GraphDatabaseService graphDb, String namespace, String nodeName, String relationshipType, String propertyName) throws Exception { List<String> entities = new ArrayList<>(); DynamicRelationshipType relationship = DynamicRelationshipType.withName(relationshipType); try (Transaction txn = graphDb.beginTx()) { String fullNodeName = getFormattedString(namespace, nodeName); String nodeHash = hash(nodeName).toString(); logger.debug("Node hash for {} => {}", fullNodeName, nodeHash); Index<Node> nodeIndex = graphDb.index().forNodes(EntityGraphConstants.EVENT_CLASS); Node node = findNode(graphDb, nodeHash, nodeIndex); if (node != null) { Iterable<Relationship> rels = node.getRelationships(relationship); for (Relationship rel : rels) { Node otherNode = (rel.getStartNode().getId() == node.getId()) ? rel.getEndNode() : rel.getStartNode(); String value = null; Object key = otherNode.getProperty(propertyName); if (key == null) { continue; } switch (propertyName) { case "name": String[] tokens = getTokensFromFormattedString(key.toString()); value = tokens[tokens.length - 1]; break; default: value = key.toString(); } entities.add((value != null) ? value : ""); } } } return entities; } private Node findNode(GraphDatabaseService graphDb, String nodeHash, Index<Node> nodeIndex) throws Exception { ResourceIterable<Node> iterableNodes = graphDb.findNodesByLabelAndProperty(eventClassLabel, EntityGraphConstants.NODE_ID, nodeHash); if (iterableNodes != null) { ResourceIterator<Node> it = iterableNodes.iterator(); if (it.hasNext()) { return it.next(); } } // use the traditional index to find the node return nodeIndex.get(EntityGraphConstants.NODE_ID, nodeHash).getSingle(); } }