From a4b3b1dc13b9fdc8800d7953dfadc355c5825fa0 Mon Sep 17 00:00:00 2001 From: Saleem Ansari Date: Thu, 31 Jan 2019 01:05:05 +0530 Subject: [PATCH] Add support for creating custom Graph Schemas for JanusGraph --- .gitignore | 1 + README-custom-graph.md | 89 ++++++++ pom.xml | 78 ++++++- src/main/assembly/assembly.xml | 42 ++++ src/main/config/graph-hbase-config.properties | 9 + src/main/config/graphtopology_config.conf | 55 +++++ .../java/metron/graph/SiaGraphTopology.java | 178 +++++++++++++++ src/main/java/metron/graph/SiaJanusBolt.java | 88 ++++++++ src/main/java/metron/graph/SiaJanusDAO.java | 55 +++++ src/main/java/metron/graph/SiaMapperBolt.java | 90 ++++++++ .../resources/azurensg_sample_record.json | 27 +++ .../resources/graph-hbase-config.properties | 9 + .../metron/graph/AzureNSGDataImporter.scala | 208 ++++++++++++++++++ .../metron/graph/ConstantsAndUtils.scala | 24 ++ .../scala/metron/graph/GlobalSchema.scala | 136 ++++++++++++ src/main/scala/metron/graph/GraphDAO.scala | 85 +++++++ .../scala/metron/graph/GraphImporter.scala | 23 ++ .../scala/metron/graph/GraphManager.scala | 130 +++++++++++ src/main/scala/metron/graph/TestMain.scala | 36 +++ 19 files changed, 1358 insertions(+), 5 deletions(-) create mode 100644 README-custom-graph.md create mode 100644 src/main/assembly/assembly.xml create mode 100644 src/main/config/graph-hbase-config.properties create mode 100644 src/main/config/graphtopology_config.conf create mode 100644 src/main/java/metron/graph/SiaGraphTopology.java create mode 100644 src/main/java/metron/graph/SiaJanusBolt.java create mode 100644 src/main/java/metron/graph/SiaJanusDAO.java create mode 100644 src/main/java/metron/graph/SiaMapperBolt.java create mode 100644 src/main/resources/azurensg_sample_record.json create mode 100644 src/main/resources/graph-hbase-config.properties create mode 100644 src/main/scala/metron/graph/AzureNSGDataImporter.scala create mode 100644 src/main/scala/metron/graph/ConstantsAndUtils.scala create mode 100644 src/main/scala/metron/graph/GlobalSchema.scala create mode 100644 src/main/scala/metron/graph/GraphDAO.scala create mode 100644 src/main/scala/metron/graph/GraphImporter.scala create mode 100644 src/main/scala/metron/graph/GraphManager.scala create mode 100644 src/main/scala/metron/graph/TestMain.scala diff --git a/.gitignore b/.gitignore index a8cca1f..21950db 100644 --- a/.gitignore +++ b/.gitignore @@ -66,3 +66,4 @@ pom.xml.tag # Maven log/ target/ +dependency-reduced-pom.xml diff --git a/README-custom-graph.md b/README-custom-graph.md new file mode 100644 index 0000000..de2108d --- /dev/null +++ b/README-custom-graph.md @@ -0,0 +1,89 @@ +# How to build a custom graph? + +Here is an example of how to build a custom graph for your dataset. Lets say we have following data: + +``` +Event -> FlowTuple +FlowTuple -> SourceIP +FlowTuple -> DestinationIP +``` + +In the above exapmle we have some JSON data that has properties for an Event, FlowTuple, and IP Addresses. + +Checkout following file: [azurensg_sample_record.json](./src/main/resources/azurensg_sample_record.json) + + +So now we can store the relationship among these entities as follows: + +``` + Event -> hasFlowTuple -> FlowTuple (via jsonId) + FlowTuple -> hasSourceIp -> IP (via ipSrcAddr) + FlowTuple -> hasDestIp -> IP (via ipDstAddr) +``` + +Checkout following file for more details: [GlobalSchema.scala](./src/main/scala/metron/graph/GlobalSchema.scala) + + +# How to import sample data into graph? + +## Setup a Janus Graph instance locally + +Steps defined here: https://gist.github.com/tuxdna/166dd41902c59ca0470252b5bf4f3dcd + + +## Build and run the Graph importer + +Build the project + +``` +export CP=`mvn dependency:build-classpath | grep -A1 "Dependencies classpath:" | tail -1` +mvn clean compile +java -cp $CP:target/classes:target/test-classes metron.graph.TestMain ./src/main/config/graph-hbase-config.properties +``` + + +[metron.graph.TestMain](./src/main/scala/metron/graph/TestMain.scala) will perform following steps: + + * Connect to graph as specified in [graph-hbase-config.properties](./src/main/config/graph-hbase-config.properties) + * Initialize Graph Schema with properties and indices as defined in [GraphManager.scala](./src/main/scala/metron/graph/GraphManager.scala) + * Read a single JSON record from [azurensg_sample_record.json](./src/main/resources/azurensg_sample_record.json), and add it to graph as defined in [AzureNSGDataImporter.scala](./src/main/scala/metron/graph/AzureNSGDataImporter.scala) + + +## Query some data from Graph + + +[QueryURL](http://localhost:8182/?gremlin=g.V().has('data_source_type','azurensg').hasLabel('flow_tuple').has('ipDstAddr','100.10.2.6').valueMap('ipSrcAddr','ipDstAddr','protocol')) + + +``` +$ wget -q -O - "http://localhost:8182/?gremlin=g.V().has('data_source_type','azurensg').hasLabel('flow_tuple').has('ipDstAddr','100.10.2.6').valueMap('ipSrcAddr','ipDstAddr','protocol')" | python -mjson.tool +{ + "requestId": "db04e13c-15a2-4ccb-9f3b-57a822490712", + "result": { + "data": [ + { + "ipDstAddr": [ + "100.10.2.6" + ], + "ipSrcAddr": [ + "100.10.1.2" + ], + "protocol": [ + "T" + ] + } + ], + "meta": {} + }, + "status": { + "attributes": {}, + "code": 200, + "message": "" + } +} +``` + + +# TODOs + +Integrating with Apache Storm Bolt diff --git a/pom.xml b/pom.xml index 8def69f..2514346 100644 --- a/pom.xml +++ b/pom.xml @@ -13,6 +13,9 @@ UTF-8 + 2.12.4 + 0.3.1 + 2.7.2 @@ -57,7 +60,7 @@ org.janusgraph janusgraph-cassandra - 0.3.1 + ${janusgraph.version} @@ -91,7 +94,7 @@ org.janusgraph janusgraph-hbase - 0.3.1 + ${janusgraph.version} @@ -124,7 +127,7 @@ org.janusgraph janusgraph-solr - 0.3.1 + ${janusgraph.version} servlet-api @@ -156,7 +159,7 @@ org.janusgraph janusgraph-es - 0.3.1 + ${janusgraph.version} @@ -190,7 +193,7 @@ org.janusgraph janusgraph-core - 0.3.1 + ${janusgraph.version} @@ -380,6 +383,20 @@ + + org.scala-lang + scala-reflect + ${scala.version} + + + + org.apache.hbase + hbase-client + 1.1.1 + jar + + + @@ -388,6 +405,9 @@ src/test/resources + + src/main/resources + @@ -451,6 +471,27 @@ + + net.alchim31.maven + scala-maven-plugin + + + scala-compile-first + process-resources + + add-source + compile + + + + scala-test-compile + process-test-resources + + testCompile + + + + @@ -463,6 +504,33 @@ + + + maven-assembly-plugin + 2.5.3 + + + jar-with-dependencies + + + + metron.graph.GraphTopology + + + + src/main/assembly/assembly.xml + + + + + package + + single + + + + + diff --git a/src/main/assembly/assembly.xml b/src/main/assembly/assembly.xml new file mode 100644 index 0000000..9353103 --- /dev/null +++ b/src/main/assembly/assembly.xml @@ -0,0 +1,42 @@ + + + tarball + + tar.gz + + false + ${project.artifactId} + + + ${project.basedir}/src/main/config/graphtopology_config.conf + stormgraph + graphtopology_config.conf + + + + + ${project.basedir}/target + stormgraph + + ${project.artifactId}-${project.version}.jar + + + + diff --git a/src/main/config/graph-hbase-config.properties b/src/main/config/graph-hbase-config.properties new file mode 100644 index 0000000..ad5d990 --- /dev/null +++ b/src/main/config/graph-hbase-config.properties @@ -0,0 +1,9 @@ +gremlin.graph=org.janusgraph.core.JanusGraphFactory +storage.backend=hbase +storage.hostname=localhost +storage.port=2181 +storage.hbase.table=janus_graph1 +cache.db-cache=true +cache.db-cache-clean-wait=20 +cache.db-cache-time=180000 +cache.db-cache-size=0.5 diff --git a/src/main/config/graphtopology_config.conf b/src/main/config/graphtopology_config.conf new file mode 100644 index 0000000..d691ffc --- /dev/null +++ b/src/main/config/graphtopology_config.conf @@ -0,0 +1,55 @@ +#generic topology settings + +top.debug = true +top.generatorSpoutEnabled = false +top.localDeploy = false +top.numWorkers = 1 +top.name = SiaStormGraph +top.spout.name = GraphSpout +top.spout.parallelism = 1 +top.mapperbolt.name = JanusMapper +top.mapperbolt.parallelism = 1 +top.graphbolt.name = JanusBolt +top.graphbolt.parallelism = 1 + +#settings for generator spout + +top.spout.generator.sleep = 1000 +top.spout.generator.randSize = 120 +top.spout.generator.outTupleName = raw +top.spout.generator.sourceFieldName = ip_src +top.spout.generator.destFieldName = ip_dst +top.spout.generator.userField = username + +#settings for kafka spout + +top.spout.kafka.bootStrapServers = localhost:6667 +top.spout.kafka.topic = indexing +top.spout.kafka.consumerGroupId = graphSpout +top.spout.kafka.offsetCommitPeriodMs = 10000 +top.spout.kafka.retry.initialDelay = 500 +top.spout.kafka.retry.delayPeriod = 2 +top.spout.kafka.retry.maxDelay = 10 +top.spout.kafka.maxUncommittedOffsets = 1000000 +#name of tuples that kafka spout outputs are set below +top.spout.kafka.tupleFieldTopic = topic +top.spout.kafka.tupleFieldPartition = partition +top.spout.kafka.tupleFieldOffset = offset +top.spout.kafka.tupleFieldKey = key +top.spout.kafka.tupleFieldValue = value + +#settings for mapper bolt + +top.mapperbolt.tupleToLookFor = value +top.mapperbolt.allowedEdges = connectsTo, uses, usedBy +top.mapperbolt.allowedVertexTypes = valueKey +top.mapperbolt.mappings = ip_src, ip_dst, connectsTo, host, host;username, ip_src, uses, user, host + +#settings for JanusBolt + +top.graphbolt.backEndConfigLocation = /path/to/graph.properties +top.graphbolt.ttlDays = 5 + +#settings for HDFS Access +top.hdfs.uri = "hdfs://localhost:8020" + diff --git a/src/main/java/metron/graph/SiaGraphTopology.java b/src/main/java/metron/graph/SiaGraphTopology.java new file mode 100644 index 0000000..10fe888 --- /dev/null +++ b/src/main/java/metron/graph/SiaGraphTopology.java @@ -0,0 +1,178 @@ +/* + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License + along with this program. If not, see . + */ + +package metron.graph; + +import org.apache.kafka.clients.consumer.ConsumerConfig; +import org.apache.storm.Config; +import org.apache.storm.LocalCluster; +import org.apache.storm.StormSubmitter; +import org.apache.storm.kafka.spout.KafkaSpout; +import org.apache.storm.kafka.spout.KafkaSpoutConfig; +import org.apache.storm.kafka.spout.KafkaSpoutRetryExponentialBackoff; +import org.apache.storm.kafka.spout.KafkaSpoutRetryService; +import org.apache.storm.topology.TopologyBuilder; +import org.apache.storm.tuple.Fields; +import org.apache.storm.tuple.Values; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.BufferedReader; +import java.io.FileInputStream; +import java.io.IOException; +import java.io.InputStreamReader; + + +public class SiaGraphTopology { + + private static final Logger logger = LoggerFactory.getLogger(SiaGraphTopology.class); + + public static void main(String[] args) throws Exception { + + if (args[0] == null) + System.out.println("Please specify the location of the file graphtopology_config.conf"); + + + Config conf = readConfigFromFile(args[0], logger); + + TopologyBuilder builder = new TopologyBuilder(); + + String spoutName = ConfigHandler.checkForNullConfigAndLoad("top.spout.name", conf); + int spoutParallelism = Integer.parseInt(ConfigHandler.checkForNullConfigAndLoad("top.spout.parallelism", conf)); + boolean generateData = Boolean + .parseBoolean(ConfigHandler.checkForNullConfigAndLoad("top.generatorSpoutEnabled", conf)); + + if (generateData) { + + logger.trace("Started initializing generator spout"); + logger.debug( + "Setting up generator spout with name " + spoutName + " and parallelism of " + spoutParallelism); + builder.setSpout(spoutName, new TelemetryLoaderSpout(), spoutParallelism); + logger.trace("Finished initializing the generator spout"); + + } else { + + logger.info("Setting kafka spout..."); + + String bootStrapServers = ConfigHandler.checkForNullConfigAndLoad("top.spout.kafka.bootStrapServers", conf); + String topic = ConfigHandler.checkForNullConfigAndLoad("top.spout.kafka.topic", conf); + String consumerGroupId = ConfigHandler.checkForNullConfigAndLoad("top.spout.kafka.consumerGroupId", conf); + Long offsetCommitPeriodMs = Long + .parseLong(ConfigHandler.checkForNullConfigAndLoad("top.spout.kafka.offsetCommitPeriodMs", conf)); + int initialDelay = Integer + .parseInt(ConfigHandler.checkForNullConfigAndLoad("top.spout.kafka.retry.initialDelay", conf)); + int delayPeriod = Integer + .parseInt(ConfigHandler.checkForNullConfigAndLoad("top.spout.kafka.retry.delayPeriod", conf)); + int maxDelay = Integer + .parseInt(ConfigHandler.checkForNullConfigAndLoad("top.spout.kafka.retry.maxDelay", conf)); + int maxUncommittedOffsets = Integer + .parseInt(ConfigHandler.checkForNullConfigAndLoad("top.spout.kafka.maxUncommittedOffsets", conf)); + + logger.trace("Started initializing kafkaSpoutRetryService"); + logger.debug("Initializing kafkaSpoutRetryService " + " with initial delay " + initialDelay + + " delay period " + delayPeriod + " max delay " + maxDelay); + + KafkaSpoutRetryService kafkaSpoutRetryService = new KafkaSpoutRetryExponentialBackoff( + KafkaSpoutRetryExponentialBackoff.TimeInterval.microSeconds(initialDelay), + KafkaSpoutRetryExponentialBackoff.TimeInterval.milliSeconds(delayPeriod), Integer.MAX_VALUE, + KafkaSpoutRetryExponentialBackoff.TimeInterval.seconds(maxDelay)); + + logger.trace("Finished initializing kafkaSpoutRetryService"); + + String tupleFieldTopic = ConfigHandler.checkForNullConfigAndLoad("top.spout.kafka.tupleFieldTopic", conf); + String tupleFieldPartition = ConfigHandler.checkForNullConfigAndLoad("top.spout.kafka.tupleFieldPartition", + conf); + String tupleFieldOffset = ConfigHandler.checkForNullConfigAndLoad("top.spout.kafka.tupleFieldOffset", conf); + String tupleFieldKey = ConfigHandler.checkForNullConfigAndLoad("top.spout.kafka.tupleFieldKey", conf); + String tupleFieldValue = ConfigHandler.checkForNullConfigAndLoad("top.spout.kafka.tupleFieldValue", conf); + + logger.trace("Started initializing spoutConf"); + KafkaSpoutConfig spoutConf = KafkaSpoutConfig.builder(bootStrapServers, topic) + .setProp(ConsumerConfig.GROUP_ID_CONFIG, consumerGroupId) + .setOffsetCommitPeriodMs(offsetCommitPeriodMs) + .setFirstPollOffsetStrategy(KafkaSpoutConfig.FirstPollOffsetStrategy.UNCOMMITTED_LATEST) + .setMaxUncommittedOffsets(maxUncommittedOffsets).setRetry(kafkaSpoutRetryService) + .setRecordTranslator((r) -> new Values(r.topic(), r.partition(), r.offset(), r.key(), r.value()), + new Fields(tupleFieldTopic, tupleFieldPartition, tupleFieldOffset, tupleFieldKey, + tupleFieldValue)) + .build(); + logger.trace("Finished initializing spoutConf"); + + builder.setSpout(spoutName, new KafkaSpout(spoutConf), spoutParallelism); + + logger.info("Finished setting kafka spout..."); + + } + + String mapperBoltName = ConfigHandler.checkForNullConfigAndLoad("top.mapperbolt.name", conf); + int mapperboltParallelism = Integer + .parseInt(ConfigHandler.checkForNullConfigAndLoad("top.mapperbolt.parallelism", conf)); + + logger.debug("Initializing " + mapperBoltName + " with parallelism " + mapperboltParallelism); + builder.setBolt(mapperBoltName, new SiaMapperBolt(), mapperboltParallelism).shuffleGrouping(spoutName); + + String graphBoltName = ConfigHandler.checkForNullConfigAndLoad("top.graphbolt.name", conf); + int graphBoltParallelism = Integer + .parseInt(ConfigHandler.checkForNullConfigAndLoad("top.graphbolt.parallelism", conf)); + + logger.debug("Initializing " + graphBoltName + " with parallelism " + graphBoltParallelism); + builder.setBolt(graphBoltName, new SiaJanusBolt(), graphBoltParallelism).shuffleGrouping(mapperBoltName); + + boolean debugMode = Boolean.getBoolean(ConfigHandler.checkForNullConfigAndLoad("top.debug", conf)); + conf.setDebug(debugMode); + + int numWorkers = Integer.parseInt(ConfigHandler.checkForNullConfigAndLoad("top.numWorkers", conf)); + conf.setNumWorkers(numWorkers); + + boolean localDeploy = Boolean.parseBoolean(ConfigHandler.checkForNullConfigAndLoad("top.localDeploy", conf)); + + String topologyName = ConfigHandler.checkForNullConfigAndLoad("top.name", conf); + + if (localDeploy) { + LocalCluster cluster = new LocalCluster(); + cluster.submitTopology(topologyName, conf, builder.createTopology()); + } + else + { + System.out.println("[METRON]Submitting topology to remote cluster..."); + StormSubmitter.submitTopology(topologyName, conf, builder.createTopology()); + } + + } + + public static Config readConfigFromFile(String filename, Logger log) throws IOException { + Config conf = new Config(); + System.out.println("[METRON] Reading config file: " + filename); + + FileInputStream fstream = new FileInputStream(filename); + BufferedReader br = new BufferedReader(new InputStreamReader(fstream)); + + String strLine; + + while ((strLine = br.readLine()) != null) { + System.out.println(strLine); + + if (strLine.length() != 0 && !(strLine.charAt(0) == '#')) { + String[] parts = strLine.split("="); + conf.put(parts[0], parts[1]); + log.debug("Setting property " + parts[0] + " to " + parts[1]); + } + } + + br.close(); + + return conf; + } +} \ No newline at end of file diff --git a/src/main/java/metron/graph/SiaJanusBolt.java b/src/main/java/metron/graph/SiaJanusBolt.java new file mode 100644 index 0000000..d9c1af5 --- /dev/null +++ b/src/main/java/metron/graph/SiaJanusBolt.java @@ -0,0 +1,88 @@ +/* + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License + along with this program. If not, see . + */ + +package metron.graph; + +import org.apache.storm.task.OutputCollector; +import org.apache.storm.task.TopologyContext; +import org.apache.storm.topology.OutputFieldsDeclarer; +import org.apache.storm.topology.base.BaseRichBolt; +import org.apache.storm.tuple.Tuple; +import org.json.simple.JSONObject; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; + +public class SiaJanusBolt extends BaseRichBolt { + + /** + * + */ + private static final long serialVersionUID = 3984660977031068498L; + + private String JANUS_CONFIG; + private int TTL_VALUE; + private Logger logger; + private SiaJanusDAO jd; + private String FIELD_TO_LOOK_FOR = "jsonObject"; + private OutputCollector collector; + + @SuppressWarnings("rawtypes") + public void prepare(Map conf, TopologyContext context, OutputCollector collector) { + this.collector = collector; + logger = LoggerFactory.getLogger(SiaJanusBolt.class); + + logger.trace("Initializing janus bolt..."); + + JANUS_CONFIG = ConfigHandler.checkForNullConfigAndLoad("top.graphbolt.backEndConfigLocation", conf); + TTL_VALUE = Integer.parseInt(ConfigHandler.checkForNullConfigAndLoad("top.graphbolt.ttlDays", conf)); + + logger.trace("Initializing Janus DAO..."); + + //java.net.URL fileURL = JanusDAO.class.getResource(JANUS_CONFIG); + //File file = new File(fileURL); + + logger.info("Loading config from: " + JANUS_CONFIG); + jd = new SiaJanusDAO(conf, JANUS_CONFIG, TTL_VALUE); + + logger.debug("Janus bolt initialized..."); + } + + public void execute(Tuple tuple) { + + try { + + if (!tuple.contains(FIELD_TO_LOOK_FOR)) + throw new IllegalArgumentException( + "JsonObject is not present, invalid input in field: " + FIELD_TO_LOOK_FOR); + + JSONObject jsonObject = (JSONObject) tuple.getValueByField(FIELD_TO_LOOK_FOR); + jd.saveJson(jsonObject); + + collector.ack(tuple); + } catch (Exception e) { + collector.fail(tuple); + e.printStackTrace(); + } + + } + + public void declareOutputFields(OutputFieldsDeclarer arg0) { + // TODO Auto-generated method stub + + } + +} diff --git a/src/main/java/metron/graph/SiaJanusDAO.java b/src/main/java/metron/graph/SiaJanusDAO.java new file mode 100644 index 0000000..140233d --- /dev/null +++ b/src/main/java/metron/graph/SiaJanusDAO.java @@ -0,0 +1,55 @@ +/* + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License + along with this program. If not, see . + */ + +package metron.graph; + +import org.apache.commons.configuration.Configuration; +import org.apache.commons.configuration.ConfigurationException; +import org.apache.commons.configuration.PropertiesConfiguration; +import org.janusgraph.core.JanusGraph; +import org.janusgraph.core.JanusGraphFactory; +import org.json.simple.JSONObject; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.net.URISyntaxException; +import java.util.Map; + +public class SiaJanusDAO { + + private JanusGraph g; + + // private int DEFAULT_TTL_DAYS; + + private String CONFIG_FILE; + private Logger logger = LoggerFactory.getLogger(ConfigHandler.class); + + public SiaJanusDAO(Map conf, String configFIle, int ttlDays) { + CONFIG_FILE = configFIle; + g = JanusGraphFactory.open(configFIle); + + // DEFAULT_TTL_DAYS = ttlDays; TODO: implement this later + + // Initialize graph schema and properties + GraphManager.initializeSchema(g); + + } + + public synchronized void saveJson(JSONObject jsonObject){ + GraphImporter.importJsonObject(g, jsonObject); + } + +} diff --git a/src/main/java/metron/graph/SiaMapperBolt.java b/src/main/java/metron/graph/SiaMapperBolt.java new file mode 100644 index 0000000..c4cf501 --- /dev/null +++ b/src/main/java/metron/graph/SiaMapperBolt.java @@ -0,0 +1,90 @@ +/* + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License + along with this program. If not, see . + */ + +package metron.graph; + +import org.apache.storm.task.OutputCollector; +import org.apache.storm.task.TopologyContext; +import org.apache.storm.topology.OutputFieldsDeclarer; +import org.apache.storm.topology.base.BaseRichBolt; +import org.apache.storm.tuple.Fields; +import org.apache.storm.tuple.Tuple; +import org.apache.storm.tuple.Values; +import org.json.simple.JSONObject; +import org.json.simple.parser.JSONParser; +import org.json.simple.parser.ParseException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; + +public class SiaMapperBolt extends BaseRichBolt { + + /** + * + */ + private static final long serialVersionUID = 3035757397365170506L; + private OutputCollector collector; + private String tupleToLookFor; + private JSONParser parser; + private Logger logger; + + @SuppressWarnings("rawtypes") + public void prepare(Map conf, TopologyContext context, OutputCollector collector) { + this.collector = collector; + logger = LoggerFactory.getLogger(SiaMapperBolt.class); + + logger.trace("Initializing parser..."); + parser = new JSONParser(); + tupleToLookFor = ConfigHandler.checkForNullConfigAndLoad("top.mapperbolt.tupleToLookFor", conf); + } + + public void declareOutputFields(OutputFieldsDeclarer declarer) { + declarer.declare(new Fields("jsonObject")); + + } + + public void execute(Tuple tuple) { + + try { + + if (!tuple.contains(tupleToLookFor)) + throw new IllegalArgumentException(tupleToLookFor + " tuple is not present"); + + JSONObject jsonObject = (JSONObject) parser.parse(tuple.getStringByField(tupleToLookFor)); + + logger.debug("Parsed json ojbect: " + jsonObject); + + if (jsonObject.keySet().size() == 0) + throw new IllegalArgumentException(jsonObject + " is not a valid message"); + + collector.emit(new Values(jsonObject)); + + collector.ack(tuple); + } + catch (ParseException e) { + collector.fail(tuple); + logger.error("Failed to pasre object" + tuple.getStringByField(tupleToLookFor)); + e.printStackTrace(); + } + catch(IllegalArgumentException ex){ + collector.fail(tuple); + logger.error("Failed to pasre object" + ex.getMessage()); + ex.printStackTrace(); + } + + } + +} diff --git a/src/main/resources/azurensg_sample_record.json b/src/main/resources/azurensg_sample_record.json new file mode 100644 index 0000000..cd318dd --- /dev/null +++ b/src/main/resources/azurensg_sample_record.json @@ -0,0 +1,27 @@ +{ + "ip_dst_port": 3128, + "rule": "rule1", + "mac": "MAC12", + "protocol": "T", + "original_string": "original_string_was_here", + "ip_dst_addr": "100.10.2.6", + "action": "A", + "ip_src_addr": "100.10.1.2", + "timestamp": 1542644926000, + "direction": "O", + "d": "19", + "system_id": "sid1", + "m": "11", + "version": "1", + "source.type": "demo_azurensg", + "operation_name": "NSGFlowEventExample1", + "json_id": "jid1", + "tuple_id": "tid1", + "ip_src_port": 32948, + "y": "2018", + "resource_id": "/some/resource/id", + "guid": "some_guid", + "time": "2018-11-19T16:28:53.2664109Z", + "category": "NetworkSecurityGroupFlowEvent", + "_version_": 1234 +} diff --git a/src/main/resources/graph-hbase-config.properties b/src/main/resources/graph-hbase-config.properties new file mode 100644 index 0000000..ad5d990 --- /dev/null +++ b/src/main/resources/graph-hbase-config.properties @@ -0,0 +1,9 @@ +gremlin.graph=org.janusgraph.core.JanusGraphFactory +storage.backend=hbase +storage.hostname=localhost +storage.port=2181 +storage.hbase.table=janus_graph1 +cache.db-cache=true +cache.db-cache-clean-wait=20 +cache.db-cache-time=180000 +cache.db-cache-size=0.5 diff --git a/src/main/scala/metron/graph/AzureNSGDataImporter.scala b/src/main/scala/metron/graph/AzureNSGDataImporter.scala new file mode 100644 index 0000000..ac3ff82 --- /dev/null +++ b/src/main/scala/metron/graph/AzureNSGDataImporter.scala @@ -0,0 +1,208 @@ +package metron.graph + +import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource +import org.apache.tinkerpop.gremlin.structure.Vertex +import org.janusgraph.core.{JanusGraph, JanusGraphTransaction} +import org.json.simple.JSONObject + +import metron.graph.ConstantsAndUtils._ + +case class AzureNSGRecord(ipDstPort: Integer, + rule: String, + mac: String, + protocol: String, + originalString: String, + ipDstAddr: String, + action: String, + ipSrcAddr: String, + timestamp: java.util.Date, + direction: String, + d: String, + systemId: String, + m: String, + version: String, + sourceType: String, + operationName: String, + jsonId: String, + tupleId: String, + ipSrcPort: Integer, + y: String, + resourceId: String, + guid: String, + time: java.util.Date, + category: String, + internalVersion: String + ) + +object AzureNSGRecord { + + def loadFromJson(o: JSONObject): AzureNSGRecord = { + val ip_dst_port: java.lang.Integer = if (o.containsKey("ip_dst_port")) Integer.parseInt(o.get("ip_dst_port").toString) else 0 + val ip_src_port: java.lang.Integer = if (o.containsKey("ip_src_port")) Integer.parseInt(o.get("ip_src_port").toString) else 0 + val rule = if (o.containsKey("rule")) o.get("rule").toString else "" + val mac = if (o.containsKey("mac")) o.get("mac").toString else "" + val protocol = if (o.containsKey("protocol")) o.get("protocol").toString else "" + val original_string = if (o.containsKey("original_string")) o.get("original_string").toString else "" + val ip_dst_addr = if (o.containsKey("ip_dst_addr")) o.get("ip_dst_addr").toString else "" + val action = if (o.containsKey("action")) o.get("action").toString else "" + val ip_src_addr = if (o.containsKey("ip_src_addr")) o.get("ip_src_addr").toString else "" + val timestamp = if (o.containsKey("timestamp")) o.get("timestamp").toString else "" + val direction = if (o.containsKey("direction")) o.get("direction").toString else "" + val d = if (o.containsKey("d")) o.get("d").toString else "" + val system_id = if (o.containsKey("system_id")) o.get("system_id").toString else "" + val m = if (o.containsKey("m")) o.get("m").toString else "" + val version = if (o.containsKey("version")) o.get("version").toString else "" + val source_type = if (o.containsKey("source_type")) o.get("source_type").toString else "" + val operation_name = if (o.containsKey("operation_name")) o.get("operation_name").toString else "" + val json_id = if (o.containsKey("json_id")) o.get("json_id").toString else "" + val tuple_id = if (o.containsKey("tuple_id")) o.get("tuple_id").toString else "" + val y = if (o.containsKey("y")) o.get("y").toString else "" + val resource_id = if (o.containsKey("resource_id")) o.get("resource_id").toString else "" + val guid = if (o.containsKey("guid")) o.get("guid").toString else "" + val time = if (o.containsKey("time")) o.get("time").toString else "" + val category = if (o.containsKey("category")) o.get("category").toString else "" + val internal_version = if (o.containsKey("`_version_`")) o.get("`_version_`").toString else "" + + val r = AzureNSGRecord( + ipDstPort = ip_dst_port, + rule = rule, + mac = mac, + protocol = protocol, + originalString = original_string, + ipDstAddr = ip_dst_addr, + action = action, + ipSrcAddr = ip_src_addr, + timestamp = parseDate(timestamp), + direction = direction, + d = d, + systemId = system_id, + m = m, + version = version, + sourceType = source_type, + operationName = operation_name, + jsonId = json_id, + tupleId = tuple_id, + ipSrcPort = ip_src_port, + y = y, + resourceId = resource_id, + guid = guid, + time = parseDate(time), + category = category, + internalVersion = internal_version + ) + + r + + } +} + +object AzureNSGDataImporter { + + def extractEntities(r: AzureNSGRecord, tenantId: String = "demo") = { + + val event = Event( + tenantId = tenantId, + jsonId = r.jsonId, + systemId = r.systemId, + eventId = r.jsonId, + resourceId = r.resourceId, + eventTime = r.time, + category = r.category, + operationName = r.operationName, + // other fields not available + eventName = "", + eventSource = "", + eventType = "", + eventVersion = "", + region = "", + accountId = "", + threatsName = "", + data_source_type = AZURE_NSG) + + val srcIpAddr = IPAddress( + ip = r.ipSrcAddr, + port = r.ipSrcPort, + ipType = "ipv4", + data_source_type = AZURE_NSG) + + val dstIpAddr = IPAddress( + ip = r.ipDstAddr, + port = r.ipDstPort, + ipType = "ipv4", + data_source_type = AZURE_NSG) + + val flowTuple = FlowTuple( + jsonId = r.jsonId, + tupleId = r.tupleId, + mac = r.mac, + timestamp = r.timestamp, + protocol = r.protocol, + direction = r.direction, + action = r.action, + rule = r.rule, + version = r.version, + ipSrcAddr = r.ipSrcAddr, + sourcePort = r.ipSrcPort, + destPort = r.ipDstPort, + ipDstAddr = r.ipDstAddr, + data_source_type = AZURE_NSG + ) + + (event, srcIpAddr, dstIpAddr, flowTuple) + + } + + + def importAzureNSG(graph: JanusGraph, tenantId: String, dataSourceType: String, o: JSONObject) = { + println(tenantId) + println(dataSourceType) + + val r = AzureNSGRecord.loadFromJson(o) + println(r) + val (event, srcIpAddr, dstIpAddr, flowTuple) = AzureNSGDataImporter.extractEntities(r) + println(event) + println(flowTuple) + println(srcIpAddr) + println(dstIpAddr) + + + val tx: JanusGraphTransaction = graph.newTransaction() + + val g: GraphTraversalSource = graph.traversal() + println(g.V().count().next()) + + AzureNSGDataImporter.importData(graph, tx, r, tenantId) + + g.tx.commit() + println(g.V().count().next()) + } + + + def importData(graph: JanusGraph, tx: JanusGraphTransaction, r: AzureNSGRecord, tenantId: String) = { + + /* + Event -> hasFlowTuple -> FlowTuple (via jsonId) + FlowTuple -> hasSourceIp -> IP (via ipSrcAddr) + FlowTuple -> hasDestIp -> IP (via ipDstAddr) + */ + + val (event, srcIpAddr, dstIpAddr, flowTuple) = extractEntities(r, tenantId) + val eventVertex = GraphDAO.addEventToGraph(graph, tx, event, AZURE_NSG) + val srcIpAddrVertex = GraphDAO.addIpAddrToGraph(graph, tx, srcIpAddr, AZURE_NSG) + val dstIpAddrVertex = GraphDAO.addIpAddrToGraph(graph, tx, dstIpAddr, AZURE_NSG) + val flowTupleVertex = GraphDAO.addFlowTupleToGraph(graph, tx, flowTuple, AZURE_NSG) + + addEdge(tx, eventVertex, flowTupleVertex, "hasFlowTuple") + addEdge(tx, flowTupleVertex, srcIpAddrVertex, "hasSourceIp") + addEdge(tx, flowTupleVertex, dstIpAddrVertex, "hasDestIp") + tx.commit() + } + + def currentTime = String.valueOf(System.currentTimeMillis) + + def addEdge(tx: JanusGraphTransaction, v1: Vertex, v2: Vertex, edgeLabel: String) = { + println(s"Add edge: $edgeLabel") + tx.getVertex(v1.id().toString.toLong).addEdge(edgeLabel, tx.getVertex(v2.id().toString.toLong), "createdTime", currentTime) + } + +} diff --git a/src/main/scala/metron/graph/ConstantsAndUtils.scala b/src/main/scala/metron/graph/ConstantsAndUtils.scala new file mode 100644 index 0000000..cbfaafb --- /dev/null +++ b/src/main/scala/metron/graph/ConstantsAndUtils.scala @@ -0,0 +1,24 @@ +package metron.graph + +import java.text.ParseException +import java.util.Date + +object VertexLabels { + val EVENT = "event" + val IP_ADDRESS = "ip_address" + val FLOW_TUPLE = "flow_tuple" +} + +object ConstantsAndUtils { + val AZURE_NSG = "azurensg" + + val EVENT_DATE_FORMAT = new java.text.SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss") + + def parseDate(dateString: String): java.util.Date = { + try { + EVENT_DATE_FORMAT.parse(dateString) + } catch { + case _: ParseException => new Date() + } + } +} diff --git a/src/main/scala/metron/graph/GlobalSchema.scala b/src/main/scala/metron/graph/GlobalSchema.scala new file mode 100644 index 0000000..0a478f5 --- /dev/null +++ b/src/main/scala/metron/graph/GlobalSchema.scala @@ -0,0 +1,136 @@ +package metron.graph + +import java.util.Date + +import VertexLabels._ + +trait VertexEntity { + def vertexLabel: String +} + +case class Event( + tenantId: String, + eventId: String, + eventName: String, + eventSource: String, + eventTime: Date, + eventType: String, + eventVersion: String, + region: String, + accountId: String, + threatsName: String, + resourceId: String, + severity: Double = 0, + count: Int = 0, + accountType: String = "", + connectionDirection: String = "", + api: String = "", + serviceName: String = "", + protocol: String = "", + location: String = "", + sourceData: String = "", + ip: String = "", + domain: String = "", + instanceId: String = "", + orgName: String = "", + userName: String = "", + userId: String = "", + accessKey: String = "", + url: String = "", + timestamp: Date = new Date(), + messageId: String = "", + topicArn: String = "", + subject: String = "", + message: String = "", + signatureVersion: String = "", + signature: String = "", + signingCertUrl: String = "", + unsubscribeUrl: String = "", + messageAttributes: String = "", + awsAccountId: String = "", + azureCloudAccountid: String = "", + entityAccountType: String = "", + entityRegion: String = "", + extraData: String = "", + friendlyType: String = "", + groupName: String = "", + groupTags: String = "", + isProtected: String = "", + msg: String = "", + resourceGroup: String = "", + ruleComplianceTag: String = "", + ruleDescription: String = "", + ruleIsDefault: String = "", + ruleLogic: String = "", + ruleName: String = "", + ruleRemediation: String = "", + ruleSeverity: String = "", + subscriptionId: String = "", + tags: String = "", + targetUserId: String = "", + actionType: String = "", + userAgent: String = "", + sourceIpAddrr: String = "", + errorCode: String = "", + errorMsg: String = "", + responseElement: String = "", + additionalEventData: String = "", + requestID: String = "", + apiVersion: String = "", + mgmtEvent: String = "", + readOnly: String = "", + resources: String = "", + recepientAcctId: String = "", + serviceEventDetails: String = "", + sharedEventID: String = "", + vpcEndPoint: String = "", + isAlert: Boolean = false, + jsonId: String = "", + systemId: String = "", + category: String = "", + operationName: String = "", + data_source_type: String + ) extends VertexEntity { + def vertexLabel = EVENT +} + +object Event extends VertexEntity { + def vertexLabel = EVENT +} + +case class IPAddress( + ip: String, + port: Int, + ipType: String, + data_source_type: String + ) extends VertexEntity { + def vertexLabel = IP_ADDRESS +} + +object IPAddress extends VertexEntity { + def vertexLabel = IP_ADDRESS +} + + +case class FlowTuple( + jsonId: String, + tupleId: String, + mac: String, + timestamp: java.util.Date, + protocol: String, + direction: String, + action: String, + rule: String, + version: String, + sourcePort: Integer, + destPort: Integer, + ipSrcAddr: String, + ipDstAddr: String, + data_source_type: String + ) extends VertexEntity { + def vertexLabel = FLOW_TUPLE +} + +object FlowTuple extends VertexEntity { + def vertexLabel = FLOW_TUPLE +} diff --git a/src/main/scala/metron/graph/GraphDAO.scala b/src/main/scala/metron/graph/GraphDAO.scala new file mode 100644 index 0000000..ad1db88 --- /dev/null +++ b/src/main/scala/metron/graph/GraphDAO.scala @@ -0,0 +1,85 @@ +package metron.graph + +import org.apache.tinkerpop.gremlin.structure +import org.janusgraph.core.JanusGraphTransaction + + +object GraphDAO { + + import org.apache.tinkerpop.gremlin.structure.T + + def addFlowTupleToGraph(graph: structure.Graph, tx: JanusGraphTransaction, o: FlowTuple, source_type: String) = { + println("Add FlowTuple") + val g = graph.traversal() + + val v = g.V().hasLabel(o.vertexLabel) + .has("data_source_type", source_type) + .has("tupleId", o.tupleId) + .tryNext().orElseGet { () => + tx.addVertex(T.label, o.vertexLabel, + "jsonId", o.jsonId, + "tupleId", o.tupleId, + "mac", o.mac, + "timestamp", o.timestamp, + "protocol", o.protocol, + "direction", o.direction, + "action", o.action, + "rule", o.rule, + "version", o.version, + "sourcePort", o.sourcePort, + "destPort", o.destPort, + "ipSrcAddr", o.ipSrcAddr, + "ipDstAddr", o.ipDstAddr, + "data_source_type", source_type) + } + v + } + + def addIpAddrToGraph(graph: structure.Graph, tx: JanusGraphTransaction, o: IPAddress, source_type: String) = { + println("Add IPAddress") + val g = graph.traversal() + + val v = g.V().hasLabel(o.vertexLabel) + .has("data_source_type", source_type) + .has("ip", o.ip) + .tryNext().orElseGet { () => + tx.addVertex(T.label, o.vertexLabel, + "ip", o.ip, + "port", o.port.asInstanceOf[java.lang.Integer], + "ipType", o.ipType, + "data_source_type", source_type) + } + v + } + + + def addEventToGraph(graph: structure.Graph, tx: JanusGraphTransaction, o: Event, source_type: String) = { + println("Add Event") + val g = graph.traversal() + + val v = g.V().hasLabel(o.vertexLabel) + .has("data_source_type", source_type) + .has("event_id", o.eventId) + .tryNext().orElseGet { () => + tx.addVertex(T.label, o.vertexLabel, + "type", o.eventType, + "event_id", o.eventId, + "event_name", o.eventName, + "event_source", o.eventSource, + "event_time", o.eventTime, + "event_type", o.eventType, + "event_version", o.eventVersion, + "timestamp", o.timestamp, + "region", o.region, + "accountId", o.accountId, + "threatsName", o.threatsName, + "resourceId", o.resourceId, + "severity", o.severity.asInstanceOf[java.lang.Double], + "count", o.count.asInstanceOf[java.lang.Integer], + "isAlert", o.isAlert.asInstanceOf[java.lang.Boolean], + "data_source_type", source_type) + } + v + } + +} diff --git a/src/main/scala/metron/graph/GraphImporter.scala b/src/main/scala/metron/graph/GraphImporter.scala new file mode 100644 index 0000000..b93bac8 --- /dev/null +++ b/src/main/scala/metron/graph/GraphImporter.scala @@ -0,0 +1,23 @@ +package metron.graph + +import org.janusgraph.core.JanusGraph +import org.json.simple.JSONObject + + +object GraphImporter { + + def importJsonObject(graph: JanusGraph, o: JSONObject) = { + val sourceType = o.get("source.type").toString + val parts = sourceType.split("_") + val tenantId = parts(0) + val dataSourceType = parts(1) + + dataSourceType match { + case "azurensg" => + AzureNSGDataImporter.importAzureNSG(graph, tenantId, dataSourceType, o) + case _ => + println("Unknown record type") + } + } + +} diff --git a/src/main/scala/metron/graph/GraphManager.scala b/src/main/scala/metron/graph/GraphManager.scala new file mode 100644 index 0000000..6151ac3 --- /dev/null +++ b/src/main/scala/metron/graph/GraphManager.scala @@ -0,0 +1,130 @@ +package metron.graph + +import java.lang +import java.time.Instant +import java.util.{Date, UUID} + +import org.apache.tinkerpop.gremlin.structure.Vertex +import org.janusgraph.core.attribute.Geoshape +import org.janusgraph.core.schema.JanusGraphManagement +import org.janusgraph.core.{JanusGraph, JanusGraphFactory} + +import scala.reflect.runtime.universe._ + +object GraphManager { + + def getMemberFields[T: TypeTag] = { + typeOf[T].members.collect { + case m: MethodSymbol if m.isCaseAccessor => m + }.map { x => + val g = x.getter + val vtype = g.typeSignature.typeSymbol.getClass + (g.name.toString, g.typeSignature.typeSymbol.name.toString) + }.toList + } + + + def initializeSchema(graph: JanusGraph) = { + //Never create new indexes while a transaction is active + graph.tx().rollback() + + val mgmt = graph.openManagement() + + for {(fieldName, fieldType) <- getMemberFields[Event]} { + GraphManager.addPropertyKeyAndIndex(graph, mgmt, fieldName, fieldType, Event.vertexLabel) + } + + for {(fieldName, fieldType) <- getMemberFields[IPAddress]} { + GraphManager.addPropertyKeyAndIndex(graph, mgmt, fieldName, fieldType, IPAddress.vertexLabel) + } + + for {(fieldName, fieldType) <- getMemberFields[FlowTuple]} { + GraphManager.addPropertyKeyAndIndex(graph, mgmt, fieldName, fieldType, FlowTuple.vertexLabel) + } + + mgmt.commit() + } + + /** + * + * @param mgmt JanusGraphManagement instance + * @param fieldName Name of the field + * @param fieldType Type of the field which could be one of: + * "Byte","Short","Int","Integer","Long","Float","Double","Date","UUID","Geoshape","Instant","String" + */ + def addPropertyKeyAndIndex(graph: JanusGraph, mgmt: JanusGraphManagement, fieldName: String, fieldType: String, vertexLabel: String) = { + + // Only following fields can be indexed. Keep everything else as String index as a fallback. + // java.lang.Byte + // java.lang.Short + // java.lang.Integer + // java.lang.Long + // java.lang.Float + // java.lang.Double + // java.util.Date + // java.util.UUID + // org.janusgraph.core.attribute.Geoshape + // java.time.Instant + // java.lang.String + + val clazz = fieldType match { + case "Byte" => classOf[lang.Byte] + case "Short" => classOf[lang.Short] + case "Int" => classOf[Integer] + case "Integer" => classOf[Integer] + case "Long" => classOf[lang.Long] + case "Float" => classOf[lang.Float] + case "Double" => classOf[lang.Double] + case "Date" => classOf[Date] + case "UUID" => classOf[UUID] + case "Geoshape" => classOf[Geoshape] + case "Instant" => classOf[Instant] + case _ => classOf[String] + } + + println(s"Property for field: $fieldName and type $fieldType ( $clazz )") + val field = if (mgmt.containsPropertyKey(fieldName)) { + println(s"Load existing property: $fieldName") + mgmt.getPropertyKey(fieldName) + } else { + println(s"Create property: $fieldName") + val k = mgmt.makePropertyKey(fieldName).dataType(classOf[String]).make() + k + } + + val idxWithlLabel = s"index_${vertexLabel}_${fieldName}" + val idx1 = if (mgmt.containsGraphIndex(idxWithlLabel)) { + println(s"Load existing label constraint index: $idxWithlLabel") + mgmt.getGraphIndex(idxWithlLabel) + } else { + println(s"Create label constraint index: $idxWithlLabel") + val idx = mgmt + .buildIndex(idxWithlLabel, classOf[Vertex]) + .addKey(field) + .indexOnly(mgmt.getOrCreateVertexLabel(vertexLabel)) + .buildCompositeIndex() + idx + } + + val idxName = s"index_$fieldName" + val idx = if (mgmt.containsGraphIndex(idxName)) { + println(s"Load existing index: $idxName") + mgmt.getGraphIndex(idxName) + } else { + println(s"Create index: $idxName") + val idx = mgmt.buildIndex(idxName, classOf[Vertex]).addKey(field).buildCompositeIndex() + idx + } + } + + def main(args: Array[String]): Unit = { + if (args.size >= 1) { + val conf = args(0) + val graph: JanusGraph = JanusGraphFactory.open(conf) + GraphManager.initializeSchema(graph) + graph.close() + } else { + println("No configuration file for JanusGraph provided as CLI argument") + } + } +} diff --git a/src/main/scala/metron/graph/TestMain.scala b/src/main/scala/metron/graph/TestMain.scala new file mode 100644 index 0000000..5b54257 --- /dev/null +++ b/src/main/scala/metron/graph/TestMain.scala @@ -0,0 +1,36 @@ +package metron.graph + +import org.janusgraph.core.{JanusGraph, JanusGraphFactory} +import org.json.simple.JSONObject +import org.json.simple.parser.JSONParser + +import scala.io.Source + +object TestMain { + + def graphConnection(configFile: String) = { + JanusGraphFactory.open(configFile) + } + + def runAzureNsg(configFile: String) = { + val graph: JanusGraph = graphConnection(configFile) + + val allData = getClass.getResource("/azurensg_sample_record.json") + val parser = new JSONParser() + val in = Source.fromURL(allData).bufferedReader() + val jsonObject = parser.parse(in).asInstanceOf[JSONObject] + + GraphManager.initializeSchema(graph) + GraphImporter.importJsonObject(graph, jsonObject) + graph.close() + } + + def main(args: Array[String]): Unit = { + if (args.length >= 1) { + val configFile = args(0) + runAzureNsg(configFile) + } else { + println("Insufficient argument. Please provide janusgraph config file as CLI argument.") + } + } +}