foss-vuln-tracker/molerat/src/main/java/it/unitn/molerat/data/db/MongoWrapper.java

246 lines
9.7 KiB
Java

package it.unitn.molerat.data.db;
import com.mongodb.BasicDBObject;
import com.mongodb.MongoClient;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import com.mongodb.client.model.Indexes;
import it.unitn.molerat.data.memory.AnalysisEntry;
import it.unitn.molerat.evidence.VulnerabilityEvidence;
import org.bson.Document;
import org.bson.types.ObjectId;
import java.util.*;
public class MongoWrapper {
private static final String projectCollection = "projects";
private static final String vulnsCollection = "vulns";
private static final String vulnEvdCollection = "vuln_evidences";
private MongoClient client = null;
private MongoDatabase db = null;
public MongoWrapper(String dbName) {
this.client = new MongoClient();
this.db = client.getDatabase(dbName);
}
private FindIterable<Document> getEntries(BasicDBObject query, String collection) {
MongoCollection<Document> docs = db.getCollection(collection);
return docs.find(query);
}
private Document getEntry(BasicDBObject query, String collection) {
return getEntries(query, collection).first();
}
private ObjectId getEntryId(Document doc) {
return (doc != null) ? doc.getObjectId("_id") : null;
}
private Document getProjectEntry(String name, String repoType, String repoPath) {
BasicDBObject query = new BasicDBObject();
query.put("name", name);
query.put("repo_type", repoType);
query.put("repo_path", repoPath);
return getEntry(query, projectCollection);
}
private Document getProjectEntry(ObjectId projectId) {
BasicDBObject query = new BasicDBObject();
query.put("_id", projectId);
return getEntry(query, projectCollection);
}
private ObjectId getProjectId(String name, String repoType, String repoPath) {
return getEntryId(getProjectEntry(name, repoType, repoPath));
}
private ObjectId getProjectId(ObjectId cveId) {
return getCveEntry(cveId).getObjectId("owner_id");
}
private Document getCveEntry(ObjectId projectId, String cve) {
BasicDBObject query = new BasicDBObject();
query.put("cve", cve);
query.put("owner_id", projectId);
return getEntry(query, vulnsCollection);
}
private Document getCveEntry(ObjectId cveId) {
BasicDBObject query = new BasicDBObject();
query.put("_id", cveId);
return getEntry(query, vulnsCollection);
}
private ObjectId getCveId(ObjectId projectId, String cve) {
return getEntryId(getCveEntry(projectId, cve));
}
public Set<ObjectId> getAllCveIds() {
Set<ObjectId> ids = new LinkedHashSet<>();
MongoCollection<Document> vulns = db.getCollection(vulnsCollection);
FindIterable<Document> docs = vulns.find();
for (Document doc : docs) {
ids.add(doc.getObjectId("_id"));
}
return ids;
}
private FindIterable<Document> getVulnEvidences(ObjectId cveId) {
BasicDBObject query = new BasicDBObject() ;
query.put("owner_id", cveId);
FindIterable<Document> result = getEntries(query, vulnEvdCollection);
result.sort(new BasicDBObject("order", -1));
return result;
}
private ObjectId insertProject(String name, String repoType, String repoPath) {
Document proj = getProjectEntry(name, repoType, repoPath);
if (proj == null) {
MongoCollection<Document> projects = db.getCollection(projectCollection);
proj = new Document("name", name)
.append("repo_type", repoType)
.append("repo_path", repoPath);
projects.insertOne(proj);
}
return getEntryId(proj);
}
private ObjectId insertCve(ObjectId projectId, String cve, String fixCommit) {
Document vuln = getCveEntry(projectId, cve);
if (vuln == null) {
MongoCollection<Document> vulns = db.getCollection(vulnsCollection);
vuln = new Document("cve", cve)
.append("fix_commit", fixCommit)
.append("owner_id", projectId)
.append("processed", false);
vulns.insertOne(vuln);
MongoCollection<Document> projects = db.getCollection(projectCollection);
projects.findOneAndUpdate(new BasicDBObject("_id", projectId), new BasicDBObject("$push",
new BasicDBObject("vulns", getEntryId(vuln))));
}
return getEntryId(vuln);
}
private boolean isVulnProcessed(ObjectId cveId) {
MongoCollection<Document> vulns = db.getCollection(vulnsCollection);
Document cve = vulns.find(new BasicDBObject("_id", cveId)).first();
return (Boolean)cve.get("processed");
}
public boolean insertAnalysisEntry(AnalysisEntry entry) {
// if there's no molerat.evidence for some reason, return false
if (entry.getVulnEvidences().isEmpty()) {
return false;
}
ObjectId projectId = insertProject(entry.getProjectName(), entry.getRepositoryType(), entry.getRepositoryPath());
ObjectId cveId = insertCve(projectId, entry.getCveName(), entry.getFixCommit());
MongoCollection<Document> vulns = db.getCollection(vulnsCollection);
MongoCollection<Document> evds = db.getCollection(vulnEvdCollection);
if (isVulnProcessed(cveId)) {
vulns.updateOne(
new BasicDBObject("_id", cveId),
new BasicDBObject("$set", new BasicDBObject("processed", false))
);
evds.deleteMany(new BasicDBObject("owner_id", cveId));
}
List<Document> records = new LinkedList<Document>();
int order= 0;
Set<String> commits = entry.getCommits();
for (String commit : commits) {
Set<VulnerabilityEvidence> evidences = entry.getVulnEvidences().get(commit);
if (evidences != null) {
for (VulnerabilityEvidence e : evidences) {
BasicDBObject record = new BasicDBObject();
record.put("owner_id", cveId);
record.put("revision", commit);
record.put("order", order);
record.put("file_path", e.getPath());
record.put("container", e.getContainer());
record.put("line_number", e.getLineNumber());
record.put("line_contents", e.getLineContents());
Document doc = new Document(record);
records.add(doc);
}
}
order--;
}
evds.insertMany(records);
vulns.updateOne(
new BasicDBObject("_id", cveId),
new BasicDBObject("$set", new BasicDBObject("processed", true))
);
//create index
evds.createIndex(Indexes.descending("order"));
return true;
}
public AnalysisEntry getAnalysisEntry(String projectName, String cveName, String repoType, String repoPath) {
ObjectId projectId = getProjectId(projectName, repoType, repoPath);
ObjectId cveId = getCveId(projectId, cveName);
return getAnalysisEntry(cveId);
}
public boolean analysisEntryExists(String projectName, String cveName, String repoType, String repoPath) {
ObjectId projectId = getProjectId(projectName, repoType, repoPath);
ObjectId cveId = getCveId(projectId, cveName);
return (projectId != null && cveId != null);
}
public Set<AnalysisEntry> getAnalysisEntries() {
Set<AnalysisEntry> entries = new LinkedHashSet<>();
FindIterable<Document> cves = db.getCollection(vulnsCollection).find();
for (Document cve : cves) {
ObjectId cveId = cve.getObjectId("_id");
AnalysisEntry entry = getAnalysisEntry(cveId);
entries.add(entry);
}
return entries;
}
public AnalysisEntry getAnalysisEntry(ObjectId cveId) {
if (!isVulnProcessed(cveId)) {
return null;
}
ObjectId projectId = getProjectId(cveId);
Document projectDoc = getProjectEntry(projectId);
String projectName = projectDoc.getString("name");
String repoType = projectDoc.getString("repo_type");
String repoPath = projectDoc.getString("repo_path");
Document cveDoc = getCveEntry(cveId);
String cveName = cveDoc.getString("cve");
String fixCommit = cveDoc.getString("fix_commit");
Map<String, Set<VulnerabilityEvidence>> vulnEvidences = new LinkedHashMap<>();
FindIterable<Document> docs = getVulnEvidences(cveId);
for (Document doc : docs) {
String filePath = doc.getString("file_path");
String container = doc.getString("container");
String revision = doc.getString("revision");
int lineNumber = doc.getInteger("line_number");
String lineContents = doc.getString("line_contents");
VulnerabilityEvidence evidence = new VulnerabilityEvidence(filePath, revision, container, lineNumber, lineContents);
if (vulnEvidences.containsKey(revision)) {
vulnEvidences.get(revision).add(evidence);
}
else {
Set<VulnerabilityEvidence> set = new LinkedHashSet<>();
set.add(evidence);
vulnEvidences.put(evidence.getCommit(), set);
}
}
/* TODO: add stuff for getting the changes molerat.evidence
Map<String, Set<ChangeEvidence>> chgEvidences = new TreeMap<>();
*/
return new AnalysisEntry(cveName, projectName, repoType, repoPath, fixCommit, vulnEvidences.keySet(), vulnEvidences, null);
}
}