<commit_before>package alignment.alignment_v2;

import gov.ornl.stucco.DBClient.Constraint;
import gov.ornl.stucco.DBClient.DBConnection;
import gov.ornl.stucco.DBClient.Constraint.Condition;

import java.io.IOException;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.json.*;

import com.tinkerpop.rexster.client.RexProException;
import com.tinkerpop.rexster.client.RexsterClient;

/**
 * Connects to Graph DB, aligns and adds new incoming graph data, provides some misc. utility functions.
 *
 */
public class Align 
{

	private static final boolean SEARCH_FOR_DUPLICATES = false;
	private static final boolean ALIGN_VERT_PROPS = false;
	private RexsterClient client = null;
	private Logger logger = null;
	private ConfigFileLoader config = null;
	private DBConnection connection = null;

	public Align() throws IOException{
		this(new DBConnection());
	}

	public Align(DBConnection c) throws IOException {
		this.connection = c;
		client = connection.getClient(); //TODO shouldn't need this...

		logger = LoggerFactory.getLogger(Align.class);

		//loading configuration file into HashMap
		config = new ConfigFileLoader();

		connection.createIndices();
	}

	@Override
	protected void finalize() throws Throwable {
		DBConnection.closeClient(this.client);
		super.finalize();
	}



	//   public void printGraph()	{
	//   	int i = 1;
	//   	for (String key : vertIDCache.keySet())	{
	//  		System.out.println(i + " key = " + key + "; value = " + vertIDCache.get(key));
	//  		i++;
	//  	}
	//   }


	public boolean load(String newGraphSection){

		boolean ret = true;
		//do all the json obj parsing up front, in case you need to panic & leave early.
		List<JSONObject> verts = new ArrayList<JSONObject>();
		List<JSONObject> edges = new ArrayList<JSONObject>();
		List<JSONObject> vertsToRetry = new ArrayList<JSONObject>();
		List<JSONObject> edgesToRetry = new ArrayList<JSONObject>();
		try{
			JSONObject graphson = new JSONObject(newGraphSection);
			JSONArray json_verts = graphson.optJSONArray("vertices");

			if(json_verts != null){			//if there are vertices...
				int vertCount = json_verts.length();
				for(int i=0; i<vertCount; i++){		//add each one to verts list
					JSONObject vert = (JSONObject)json_verts.get(i);
					if(!vert.has("name"))
						vert.put("name", vert.get("_id"));	//add "name" field equals its ID if needed
					verts.add(vert);	//place vertex json object
				}
			}
			//...and likewise for edges
			JSONArray json_edges = graphson.optJSONArray("edges");
			if(json_edges != null){
				int edgeCount = json_edges.length();
				for(int i=0; i<edgeCount; i++){
					JSONObject edge = (JSONObject)json_edges.get(i);
					edges.add(edge);
				}
			}
		}catch(Exception e){ 
			//we want *any* graphson problems to end up here
			//being noisy when these fail is probably ok, we shouldn't really ever fail here except when testing, etc.
			logger.error("Error parsing GraphSON in load()!");
			logger.error("The graphson was: " + newGraphSection);
			logger.error("Exception!",e);
			return false;
		}

		//for *vertices*, you have a json object that you can load.
		for(JSONObject vert : verts){
			try{
				String vert_name = vert.getString("name");
				String vert_id = vert.optString("_id");
				if(vert_name == null || vert_name == ""){
					vert_name = vert_id;
					vert.put("name", vert_name);
				}
				boolean new_vert = false;
				String otherVertID = connection.findVertId(vert_name);
				new_vert = (otherVertID == null);
				
				if(!new_vert && SEARCH_FOR_DUPLICATES){
					Map<String, Object> vertMap = null;
					vertMap = jsonVertToMap(vert);
					otherVertID = findDuplicateVertex(vertMap);
					new_vert = (otherVertID == null);
				}
				if(new_vert){ //only add new...
					loadNewJSONVert(vert);
					
					List<JSONObject> newEdges = findNewEdges(vert);
					for(JSONObject edge: newEdges){
						edges.add(edge);
					}
				}else{
					if(ALIGN_VERT_PROPS){
						Map<String, Object> vertMap = null;
						if(vertMap == null) vertMap = jsonVertToMap(vert); //might have this from above already, or might not
						alignVertProps(otherVertID, vertMap);
					}
					else{
						logger.debug("Attempted to add vertex when duplicate exists.  ALIGN_VERT_PROPS is false, so ignoring new vert.  vert was: " + vert);
					}
				}
			}catch(Exception e){
				//TODO warn
				vertsToRetry.add(vert);
			}
		}

		//for *edges*, you can't really do that, so find IDs and build a map of needed properties instead.
		for(JSONObject edge : edges){
			try{
				boolean edgeResult = loadJSONEdge(edge);
				if( edgeResult == false) edgesToRetry.add(edge);  //this can happen if the edge is missing one of its verts, which could be in the retry queue as well.
			}catch(Exception e){
				//TODO warn
				edgesToRetry.add(edge);
			}
		}

		//retry verts as needed.
		for(JSONObject vert : vertsToRetry){
			try{
				String vert_name = vert.getString("name");
				boolean new_vert = false;
				String otherVertID = connection.findVertId(vert_name);
				new_vert = (otherVertID == null);
				
				if(!new_vert && SEARCH_FOR_DUPLICATES){
					Map<String, Object> vertMap = null;
					vertMap = jsonVertToMap(vert);
					otherVertID = findDuplicateVertex(vertMap);
					new_vert = (otherVertID == null);
				}
				if(new_vert){ //only add new...
					loadNewJSONVert(vert);
					
					List<JSONObject> newEdges = findNewEdges(vert);
					for(JSONObject edge: newEdges){
						edges.add(edge);
					}
				}else{
					if(ALIGN_VERT_PROPS){
						Map<String, Object> vertMap = null;
						if(vertMap == null) vertMap = jsonVertToMap(vert); //might have this from above already, or might not
						alignVertProps(otherVertID, vertMap);
					}
					else{
						logger.debug("Attempted to add vertex when duplicate exists.  ALIGN_VERT_PROPS is false, so ignoring new vert.  vert was: " + vert);
					}
				}
			}catch(Exception e){
				//TODO warn
				ret = false;
			}
		}

		//for *edges*, you can't really do that, so find IDs and build a map of needed properties instead.
		for(JSONObject edge : edgesToRetry){
			try{
				loadJSONEdge(edge); //TODO: unused return code - is it useful here?
			}catch(Exception e){
				//TODO warn
				ret = false;
			}
		}
		
		return ret;//TODO currently this is not idempotent, but after docIDs are added to the metadata (and used) they will be.
	}

	private void loadNewJSONVert(JSONObject vert) throws RexProException, IOException{
		Map<String, Object> vertMap = null;
		String type = null;
		type = (String)vert.opt("vertexType");
		if(type.equals("IP")){ 
			//if its an ip vert, and doesn't have an ip int, just add that here.
			//TODO: the extractors should really be doing this, but some aren't
			long ipInt = vert.optLong("ipInt");
			if(ipInt == 0){
				String ipString = vert.getString("name");
				ipInt = getIpInt(ipString);
				vert.put("ipInt", ipInt);
			}
		}
		vertMap = jsonVertToMap(vert);
		//connection.addVertexFromJSON(vert);
		connection.addVertexFromMap(vertMap);
	}
	
	//return true if succeeded
	//return false if edge cannot be added
	private boolean loadJSONEdge(JSONObject edge) throws JSONException, IOException, RexProException{
		String outv_id = connection.findVertId(edge.getString("_outV"));
		String inv_id = connection.findVertId(edge.getString("_inV"));
		String edgeName = edge.getString("_id");
		//String edgeID = findEdgeId(edgeName);
		if(outv_id == null){
			logger.error("Could not find out_v for edge: " + edge);
			return false;
		}
		if(inv_id == null){
			logger.error("Could not find in_v for edge: " + edge);
			return false;
		}
		String label = edge.optString("_label");
		if(connection.edgeExists(inv_id, outv_id, label)){
			//TODO need to merge edge props for this case, like verts above...
			logger.debug("Attempted to add duplicate edge.  ignoring it.  edge was: " + edge);
			return false;
		}
		connection.addEdgeFromJSON(edge); //TODO unused return code from this - is it even useful?
		return true;
	}
	
	private List<JSONObject> findNewEdges(JSONObject vert) {
		List<JSONObject> edges = new ArrayList<JSONObject>();
		
		String vert_name = vert.getString("name");
		String type = null;
		type = (String)vert.opt("vertexType");
		if(type == null){
			logger.warn("no vertex type specified for vertex:" + vert.toString());
		}else{
			if(type.equals("addressRange")){
				long endIpInt = vert.optLong("endIPInt");
				long startIpInt = vert.optLong("startIPInt");
				if( endIpInt !=0 && startIpInt != 0){
					List<Constraint> constraints = new ArrayList<Constraint>();
					Constraint c = new Constraint("vertexType", Condition.eq, "IP");
					constraints.add(c);
					c = new Constraint("ipInt", Condition.lte, endIpInt);
					constraints.add(c);
					c = new Constraint("ipInt", Condition.gte, startIpInt);
					constraints.add(c);<commit_msg>reduce redundant log output when adding edges<commit_after>package alignment.alignment_v2;

import gov.ornl.stucco.DBClient.Constraint;
import gov.ornl.stucco.DBClient.DBConnection;
import gov.ornl.stucco.DBClient.Constraint.Condition;

import java.io.IOException;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.json.*;

import com.tinkerpop.rexster.client.RexProException;
import com.tinkerpop.rexster.client.RexsterClient;

/**
 * Connects to Graph DB, aligns and adds new incoming graph data, provides some misc. utility functions.
 *
 */
public class Align 
{

	private static final boolean SEARCH_FOR_DUPLICATES = false;
	private static final boolean ALIGN_VERT_PROPS = false;
	private RexsterClient client = null;
	private Logger logger = null;
	private ConfigFileLoader config = null;
	private DBConnection connection = null;

	public Align() throws IOException{
		this(new DBConnection());
	}

	public Align(DBConnection c) throws IOException {
		this.connection = c;
		client = connection.getClient(); //TODO shouldn't need this...

		logger = LoggerFactory.getLogger(Align.class);

		//loading configuration file into HashMap
		config = new ConfigFileLoader();

		connection.createIndices();
	}

	@Override
	protected void finalize() throws Throwable {
		DBConnection.closeClient(this.client);
		super.finalize();
	}



	//   public void printGraph()	{
	//   	int i = 1;
	//   	for (String key : vertIDCache.keySet())	{
	//  		System.out.println(i + " key = " + key + "; value = " + vertIDCache.get(key));
	//  		i++;
	//  	}
	//   }


	public boolean load(String newGraphSection){

		boolean ret = true;
		//do all the json obj parsing up front, in case you need to panic & leave early.
		List<JSONObject> verts = new ArrayList<JSONObject>();
		List<JSONObject> edges = new ArrayList<JSONObject>();
		List<JSONObject> vertsToRetry = new ArrayList<JSONObject>();
		List<JSONObject> edgesToRetry = new ArrayList<JSONObject>();
		try{
			JSONObject graphson = new JSONObject(newGraphSection);
			JSONArray json_verts = graphson.optJSONArray("vertices");

			if(json_verts != null){			//if there are vertices...
				int vertCount = json_verts.length();
				for(int i=0; i<vertCount; i++){		//add each one to verts list
					JSONObject vert = (JSONObject)json_verts.get(i);
					if(!vert.has("name"))
						vert.put("name", vert.get("_id"));	//add "name" field equals its ID if needed
					verts.add(vert);	//place vertex json object
				}
			}
			//...and likewise for edges
			JSONArray json_edges = graphson.optJSONArray("edges");
			if(json_edges != null){
				int edgeCount = json_edges.length();
				for(int i=0; i<edgeCount; i++){
					JSONObject edge = (JSONObject)json_edges.get(i);
					edges.add(edge);
				}
			}
		}catch(Exception e){ 
			//we want *any* graphson problems to end up here
			//being noisy when these fail is probably ok, we shouldn't really ever fail here except when testing, etc.
			logger.error("Error parsing GraphSON in load()!");
			logger.error("The graphson was: " + newGraphSection);
			logger.error("Exception!",e);
			return false;
		}

		//for *vertices*, you have a json object that you can load.
		for(JSONObject vert : verts){
			try{
				String vert_name = vert.getString("name");
				String vert_id = vert.optString("_id");
				if(vert_name == null || vert_name == ""){
					vert_name = vert_id;
					vert.put("name", vert_name);
				}
				boolean new_vert = false;
				String otherVertID = connection.findVertId(vert_name);
				new_vert = (otherVertID == null);
				
				if(!new_vert && SEARCH_FOR_DUPLICATES){
					Map<String, Object> vertMap = null;
					vertMap = jsonVertToMap(vert);
					otherVertID = findDuplicateVertex(vertMap);
					new_vert = (otherVertID == null);
				}
				if(new_vert){ //only add new...
					loadNewJSONVert(vert);
					
					List<JSONObject> newEdges = findNewEdges(vert);
					for(JSONObject edge: newEdges){
						edges.add(edge);
					}
				}else{
					if(ALIGN_VERT_PROPS){
						Map<String, Object> vertMap = null;
						if(vertMap == null) vertMap = jsonVertToMap(vert); //might have this from above already, or might not
						alignVertProps(otherVertID, vertMap);
					}
					else{
						logger.debug("Attempted to add vertex when duplicate exists.  ALIGN_VERT_PROPS is false, so ignoring new vert.  vert was: " + vert);
					}
				}
			}catch(Exception e){
				//TODO warn
				vertsToRetry.add(vert);
			}
		}

		//for *edges*, you can't really do that, so find IDs and build a map of needed properties instead.
		for(JSONObject edge : edges){
			try{
				boolean edgeResult = loadJSONEdge(edge, false);
				if( edgeResult == false) edgesToRetry.add(edge);  //this can happen if the edge is missing one of its verts, which could be in the retry queue as well.
			}catch(Exception e){
				//TODO warn
				edgesToRetry.add(edge);
			}
		}

		//retry verts as needed.
		for(JSONObject vert : vertsToRetry){
			try{
				String vert_name = vert.getString("name");
				boolean new_vert = false;
				String otherVertID = connection.findVertId(vert_name);
				new_vert = (otherVertID == null);
				
				if(!new_vert && SEARCH_FOR_DUPLICATES){
					Map<String, Object> vertMap = null;
					vertMap = jsonVertToMap(vert);
					otherVertID = findDuplicateVertex(vertMap);
					new_vert = (otherVertID == null);
				}
				if(new_vert){ //only add new...
					loadNewJSONVert(vert);
					
					List<JSONObject> newEdges = findNewEdges(vert);
					for(JSONObject edge: newEdges){
						edges.add(edge);
					}
				}else{
					if(ALIGN_VERT_PROPS){
						Map<String, Object> vertMap = null;
						if(vertMap == null) vertMap = jsonVertToMap(vert); //might have this from above already, or might not
						alignVertProps(otherVertID, vertMap);
					}
					else{
						logger.debug("Attempted to add vertex when duplicate exists.  ALIGN_VERT_PROPS is false, so ignoring new vert.  vert was: " + vert);
					}
				}
			}catch(Exception e){
				//TODO warn
				ret = false;
			}
		}

		//for *edges*, you can't really do that, so find IDs and build a map of needed properties instead.
		for(JSONObject edge : edgesToRetry){
			try{
				loadJSONEdge(edge, true); //TODO: unused return code - is it useful here?
			}catch(Exception e){
				//TODO warn
				ret = false;
			}
		}
		
		return ret;//TODO currently this is not idempotent, but after docIDs are added to the metadata (and used) they will be.
	}

	private void loadNewJSONVert(JSONObject vert) throws RexProException, IOException{
		Map<String, Object> vertMap = null;
		String type = null;
		type = (String)vert.opt("vertexType");
		if(type.equals("IP")){ 
			//if its an ip vert, and doesn't have an ip int, just add that here.
			//TODO: the extractors should really be doing this, but some aren't
			long ipInt = vert.optLong("ipInt");
			if(ipInt == 0){
				String ipString = vert.getString("name");
				ipInt = getIpInt(ipString);
				vert.put("ipInt", ipInt);
			}
		}
		vertMap = jsonVertToMap(vert);
		//connection.addVertexFromJSON(vert);
		connection.addVertexFromMap(vertMap);
	}
	
	private boolean loadJSONEdge(JSONObject edge) throws JSONException, IOException, RexProException{
		return loadJSONEdge(edge, true);
	}
	
	//return true if succeeded
	//return false if edge cannot be added
	private boolean loadJSONEdge(JSONObject edge, boolean lastTry) throws JSONException, IOException, RexProException{
		String outv_id = connection.findVertId(edge.getString("_outV"));
		String inv_id = connection.findVertId(edge.getString("_inV"));
		String edgeName = edge.getString("_id");
		//String edgeID = findEdgeId(edgeName);
		if(outv_id == null){
			if(lastTry) logger.warn("Could not find out_v for edge: " + edge);
			return false;
		}
		if(inv_id == null){
			if(lastTry) logger.warn("Could not find in_v for edge: " + edge);
			return false;
		}
		String label = edge.optString("_label");
		if(connection.getEdgeCount(inv_id, outv_id, label) >= 1){
			//TODO need to merge edge props for this case, like verts above...
			logger.debug("Attempted to add duplicate edge.  ignoring it.  edge was: " + edge);
			return false;
		}
		connection.addEdgeFromJSON(edge); //TODO unused return code from this - is it even useful?
		return true;
	}
	
	private List<JSONObject> findNewEdges(JSONObject vert) {
		List<JSONObject> edges = new ArrayList<JSONObject>();
		
		String vert_name = vert.getString("name");
		String type = null;
		type = (String)vert.opt("vertexType");
		if(type == null){
			logger.warn("no vertex type specified for vertex:" + vert.toString());
		}else{
			if(type.equals("addressRange")){
				long endIpInt = vert.optLong("endIPInt");
				long startIpInt = vert.optLong("startIPInt");
				if( endIpInt !=0 && startIpInt != 0){
					List<Constraint> constraints = new ArrayList<Constraint>();
					Constraint c = new Constraint("vertexType", Condition.eq, "IP");
					constraints.add(c);
					c = new Constraint("ipInt", Condition.lte, endIpInt);
					constraints.add(c);
					c = new Constraint("ipInt", Condition.gte, startIpInt);
					constraints.add(c);
{'fraction_non_alphanumeric': 0.12202478885598332, 'fraction_numerical': 0.001042009432927498, 'mean_word_length': 3.5169680455783996, 'pattern_counts': {'":': 0, '<': 47, '<?xml version=': 0, '>': 44, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4562', 'n_tokens_mistral': 6850, 'n_tokens_neox': 5756, 'n_words': 1635}