/**
* This file is part of Graylog.
*
* Graylog 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.
*
* Graylog 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 Graylog. If not, see <http://www.gnu.org/licenses/>.
*/
package org.graylog2.rules;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.google.common.io.Files;
import org.apache.commons.io.FilenameUtils;
import org.drools.compiler.kie.builder.impl.InternalKieModule;
import org.graylog2.plugin.Message;
import org.graylog2.plugin.RulesEngine;
import org.kie.api.KieServices;
import org.kie.api.builder.KieBuilder;
import org.kie.api.builder.KieFileSystem;
import org.kie.api.builder.KieModule;
import org.kie.api.builder.ReleaseId;
import org.kie.api.io.Resource;
import org.kie.api.io.ResourceType;
import org.kie.api.runtime.KieContainer;
import org.kie.api.runtime.KieSession;
import org.kie.api.runtime.rule.FactHandle;
import org.kie.internal.io.ResourceFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import javax.inject.Inject;
import javax.inject.Singleton;
import java.io.File;
import java.io.IOException;
import java.net.MalformedURLException;
import java.net.URI;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.concurrent.atomic.AtomicReference;
@Singleton
public class DroolsEngine implements RulesEngine {
private static final Logger LOG = LoggerFactory.getLogger(DroolsEngine.class);
private final KieServices kieServices;
private final Set<URI> builtinRuleUrls;
private KieContainer kieContainer;
private final AtomicReference<KieSession> session = new AtomicReference<>();
private final List<String> liveRules = Lists.newArrayList();
private int version = 0;
private ReleaseId currentReleaseId;
@Inject
public DroolsEngine(Set<URI> builtinRuleUrls) {
this.builtinRuleUrls = ImmutableSet.copyOf(builtinRuleUrls);
this.kieServices = KieServices.Factory.get();
this.liveRules.add("// placeholder rule");
commitRules();
}
public void stop() {
LOG.debug("Stopping drools session and removing all rules.");
final KieSession activeSession = session.getAndSet(null);
if (activeSession != null) {
activeSession.dispose();
}
if (currentReleaseId != null) {
kieServices.getRepository().removeKieModule(currentReleaseId);
}
}
@Override
public synchronized boolean addRule(String ruleSource) {
LOG.debug("Adding rule {}", ruleSource);
liveRules.add(ruleSource);
if (!commitRules()) {
// adding rule failed, remove the ruleSource from our list of liveRules again.
liveRules.remove(ruleSource);
return false;
}
return true;
}
@Override
public synchronized boolean addRulesFromFile(String rulesFile) {
LOG.debug("Adding drools rules from file {}", rulesFile);
try {
final String rulesSource = Files.toString(new File(rulesFile), StandardCharsets.UTF_8);
return addRule(rulesSource);
} catch (IOException e) {
LOG.warn("Could not read drools source file. Not loading rules: {}", e.getMessage());
}
return false;
}
@Override
public int evaluateInSharedSession(Message message) {
final KieSession kieSession = session.get();
if (kieSession != null) {
kieSession.insert(message);
return kieSession.fireAllRules();
}
return 0;
}
@Override
public RulesSession createPrivateSession() {
final KieSession kieSession = kieContainer.newKieSession();
// kieSession.addEventListener(new DebugRuleRuntimeEventListener());
// kieSession.addEventListener(new DebugAgendaEventListener());
kieSession.setGlobal("log", LOG);
return new DroolsRulesSession(kieSession);
}
@Override
public Object insertFact(Object fact) {
return session.get().insert(fact);
}
@Override
public boolean deleteFact(Object fact) {
final FactHandle factHandle = session.get().getFactHandle(fact);
if (factHandle == null) {
return false;
}
session.get().delete(factHandle);
return true;
}
private boolean commitRules() {
final ReleaseId previousReleaseId = currentReleaseId;
final ReleaseId newReleaseId = nextRulesPackageVersion();
LOG.debug("Committing rules as version {}", newReleaseId);
final boolean deployed = deployRules(newReleaseId);
if (deployed && previousReleaseId != null) {
kieServices.getRepository().removeKieModule(previousReleaseId);
}
return deployed;
}
private boolean deployRules(ReleaseId newReleaseId) {
try {
// add common header
// TODO this will go wrong at some point, use the DRL6Parser to figure what to add,
// and potentially also just modify the tree it generates to add the import and globals we want
final String[] drls = new String[liveRules.size()];
int i = 0;
for (String drl : liveRules) {
drls[i] = "package org.graylog2.rules\n" +
"import org.graylog2.plugin.*\n" +
"global org.slf4j.Logger log\n" +
"\n" + drl;
i++;
}
createAndDeployJar(kieServices,
newReleaseId,
drls);
if (kieContainer == null) {
kieContainer = kieServices.newKieContainer(newReleaseId);
final KieSession session = kieContainer.newKieSession();
this.session.set(session);
session.setGlobal("log", LOG);
}
kieContainer.updateToVersion(newReleaseId);
return true;
} catch (RulesCompilationException e) {
LOG.warn("Unable to add rules due to compilation errors.", e);
return false;
}
}
private ReleaseId nextRulesPackageVersion() {
currentReleaseId = kieServices.newReleaseId("org.graylog2", "dynamic-rules", Integer.toString(version++));
return currentReleaseId;
}
private KieModule createAndDeployJar(KieServices ks,
ReleaseId releaseId,
String... drls) throws RulesCompilationException {
byte[] jar = createKJar(ks, releaseId, null, drls);
return deployJar(ks, jar);
}
private byte[] createKJar(KieServices ks,
ReleaseId releaseId,
String pom,
String... drls) throws RulesCompilationException {
KieFileSystem kfs = ks.newKieFileSystem();
if (pom != null) {
kfs.write("pom.xml", pom);
} else {
kfs.generateAndWritePomXML(releaseId);
}
for (int i = 0; i < drls.length; i++) {
if (drls[i] != null) {
kfs.write("src/main/resources/r" + i + ".drl", drls[i]);
}
}
for (URI builtinRuleUrl : builtinRuleUrls) {
final String rulesFileName = FilenameUtils.getName(builtinRuleUrl.getPath());
final String path = "src/main/resources/" + rulesFileName;
final URL url;
try {
url = builtinRuleUrl.toURL();
} catch (MalformedURLException e) {
throw new RulesCompilationException(Collections.<org.kie.api.builder.Message>emptyList());
}
final Resource resource = ResourceFactory
.newUrlResource(url)
.setSourcePath(path)
.setResourceType(ResourceType.DRL);
kfs.write(resource);
}
final KieBuilder kb = ks.newKieBuilder(kfs).buildAll();
if (kb.getResults().hasMessages(org.kie.api.builder.Message.Level.ERROR)) {
throw new RulesCompilationException(kb.getResults().getMessages());
}
final InternalKieModule kieModule = (InternalKieModule) ks.getRepository().getKieModule(releaseId);
return kieModule.getBytes();
}
private KieModule deployJar(KieServices ks, byte[] jar) {
// Deploy jar into the repository
final Resource jarRes = ks.getResources().newByteArrayResource(jar);
return ks.getRepository().addKieModule(jarRes);
}
}