[ https://issues.apache.org/jira/browse/NIFI-1965?page=com.atlassian.jira.plugin.system.issuetabpanels:comment-tabpanel&focusedCommentId=15397362#comment-15397362 ]
ASF GitHub Bot commented on NIFI-1965: -------------------------------------- Github user trixpan commented on a diff in the pull request: https://github.com/apache/nifi/pull/496#discussion_r72596906 --- Diff: nifi-nar-bundles/nifi-enrich-bundle/nifi-enrich-processors/src/main/java/org/apache/nifi/processors/enrich/QueryDNS.java --- @@ -0,0 +1,269 @@ +/* + * 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. + */ + +package org.apache.nifi.processors.enrich; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.Hashtable; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; +import javax.naming.Context; +import javax.naming.NameNotFoundException; +import javax.naming.NamingEnumeration; +import javax.naming.NamingException; +import javax.naming.directory.Attributes; +import javax.naming.directory.BasicAttributes; +import javax.naming.directory.DirContext; +import javax.naming.directory.InitialDirContext; + +import org.apache.commons.lang3.StringUtils; + +import org.apache.nifi.annotation.behavior.EventDriven; +import org.apache.nifi.annotation.behavior.InputRequirement; +import org.apache.nifi.annotation.behavior.SideEffectFree; +import org.apache.nifi.annotation.behavior.SupportsBatching; +import org.apache.nifi.annotation.behavior.WritesAttribute; +import org.apache.nifi.annotation.behavior.WritesAttributes; +import org.apache.nifi.annotation.documentation.CapabilityDescription; +import org.apache.nifi.annotation.documentation.Tags; +import org.apache.nifi.annotation.lifecycle.OnScheduled; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.Relationship; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processor.util.StandardValidators; + + +@EventDriven +@SideEffectFree +@SupportsBatching +@Tags({"dns", "enrich", "ip"}) +@InputRequirement(InputRequirement.Requirement.INPUT_REQUIRED) +@CapabilityDescription("A powerful DNS query processor primary designed to enrich DataFlows with DNS based APIs " + + "(e.g. RBLs, ShadowServer's ASN lookup) but that can be also used to perform regular DNS lookups.") +@WritesAttributes({ + @WritesAttribute(attribute = "enrich.dns.record*.group*", description = "The captured fields of the DNS query response for each of the records received"), +}) +public class QueryDNS extends AbstractEnrichProcessor { + + public static final PropertyDescriptor DNS_QUERY_TYPE = new PropertyDescriptor.Builder() + .name("DNS Query Type") + .description("The DNS query type to be used by the processor (e.g. TXT, A)") + .required(true) + .defaultValue("TXT") + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .build(); + + public static final PropertyDescriptor DNS_SERVER = new PropertyDescriptor.Builder() + .name("DNS Servers") + .description("A comma separated list of DNS servers to be used. (Defaults to system wide if none is used)") + .required(false) + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .build(); + + public static final PropertyDescriptor DNS_TIMEOUT = new PropertyDescriptor.Builder() + .name("DNS Query Timeout") + .description("The amount of milliseconds to wait until considering a query as failed") + .required(true) + .defaultValue("1500") + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .build(); + + public static final PropertyDescriptor DNS_RETRIES = new PropertyDescriptor.Builder() + .name("DNS Query Retries") + .description("The number of attempts before giving up and moving on") + .required(true) + .defaultValue("1") + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .build(); + + + private final static List<PropertyDescriptor> propertyDescriptors; + private final static Set<Relationship> relationships; + + private DirContext ictx; + + // Assign the default and generally used contextFactory value + private String contextFactory = com.sun.jndi.dns.DnsContextFactory.class.getName();; + + static { + List<PropertyDescriptor> props = new ArrayList<>(); + props.add(QUERY_INPUT); + props.add(QUERY_PARSER); + props.add(QUERY_PARSER_INPUT); + props.add(DNS_RETRIES); + props.add(DNS_TIMEOUT); + props.add(DNS_SERVER); + props.add(DNS_QUERY_TYPE); + propertyDescriptors = Collections.unmodifiableList(props); + + Set<Relationship> rels = new HashSet<>(); + rels.add(REL_FOUND); + rels.add(REL_NOT_FOUND); + relationships = Collections.unmodifiableSet(rels); + } + + private AtomicBoolean initialized; + + @Override + protected List<PropertyDescriptor> getSupportedPropertyDescriptors() { + return propertyDescriptors; + } + + @Override + public Set<Relationship> getRelationships() { + return relationships; + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException { + try { + initializeResolver(context); + } catch (Exception e) { + context.yield(); + throw new ProcessException("Failed to initialize the JNDI DNS resolver server", e); + } + + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + + final String queryType = context.getProperty(DNS_QUERY_TYPE).getValue(); + final String queryInput = context.getProperty(QUERY_INPUT).evaluateAttributeExpressions(flowFile).getValue(); + final String queryParser = context.getProperty(QUERY_PARSER).getValue(); + final String queryRegex = context.getProperty(QUERY_PARSER_INPUT).getValue(); + final String dnsServer = context.getProperty(DNS_SERVER).getValue(); + + boolean found = false; + try { + String initialResult = ""; + Attributes results = doLookup(queryInput, queryType); + // NOERROR & NODATA seem to return empty Attributes handled bellow + // but defaulting to not found in any case + if (results.size() < 1) { + found = false; + + + } else { + int recordNumber = 0; + NamingEnumeration<?> dnsEntryIterator = results.get(queryType).getAll(); + while (dnsEntryIterator.hasMoreElements()) { + String dnsRecord = dnsEntryIterator.next().toString(); + // While NXDOMAIN is being generated by doLookup catch + if (dnsRecord != queryType + "NXDOMAIN") { + Map parsedResults = parseResponse(recordNumber, dnsRecord, queryParser, queryRegex, "dns"); + flowFile = session.putAllAttributes(flowFile, parsedResults); + found = true; + + + } else { + // Otherwise treat as not found + found = false; + + + } + // Increase the counter and iterate over next record.... + recordNumber++; + } + } + } catch (NamingException e) { + session.rollback(true); + getLogger().error("Something went wrong while processing the field. Please review your configuration"); + } + + + // Finally prepare to send the data down the pipeline + if (found) { + // Sending the resulting flowfile (with attributes) to REL_FOUND + session.transfer(flowFile, REL_FOUND); + } else { + // NXDOMAIN received, accepting the fate but forwarding + // to REL_NOT_FOUND + session.transfer(flowFile, REL_NOT_FOUND); + } + } + + @OnScheduled + public void onScheduled(ProcessContext context) { + initialized.set(false); + } + + protected void initializeResolver(final ProcessContext context ) { + if (initialized.get()) { + return; + } + + final String dnsTimeout = context.getProperty(DNS_TIMEOUT).getValue(); + final String dnsServer = context.getProperty(DNS_SERVER).getValue(); + final String dnsRetries = context.getProperty(DNS_RETRIES).getValue(); + + + String finalServer = ""; + Hashtable env = new Hashtable(); + env.put("java.naming.factory.initial", contextFactory); + env.put("com.sun.jndi.dns.timeout.initial", dnsTimeout); + env.put("com.sun.jndi.dns.timeout.retries", dnsRetries); + if (StringUtils.isNotEmpty(dnsServer)) { + for (String server : dnsServer.split(",")) { + finalServer = finalServer + "dns://" + server + "/. "; + } + env.put(Context.PROVIDER_URL, finalServer); + } + + try { + initializeContext(env); + } catch (NamingException e) { + getLogger().error("Could not initialize JNDI context"); + } + } + + + /** + * This method performs a simple DNS lookup using JNDI + * @param queryInput String containing the query body itself (e.g. 4.3.3.1.in-addr.arpa); + * @param queryType String containign the query type (e.g. TXT); + */ + protected Attributes doLookup(String queryInput, String queryType) throws NamingException { + // This is a simple DNS lookup attempt + + Attributes attrs; + + try { + // Uses pre-existing context to resolve + attrs = ictx.getAttributes(queryInput, new String[]{queryType}); + return attrs; + } catch ( NameNotFoundException e) { + getLogger().debug("Resolution for domain {} failed due to {}", new Object[]{queryInput, e}); + attrs = new BasicAttributes(queryType, "NXDOMAIN",true); + return attrs; + } + } + + protected void initializeContext(Hashtable env) throws NamingException { --- End diff -- addressed > Create a QueryDNS processor > --------------------------- > > Key: NIFI-1965 > URL: https://issues.apache.org/jira/browse/NIFI-1965 > Project: Apache NiFi > Issue Type: Bug > Reporter: Andre > > As part of a data pipeline security teams frequently must enrich data using > DNS enabled APIs such as: > ShadowServer BGP and ASN lookup via DNS > https://www.shadowserver.org/wiki/pmwiki.php/Services/IP-BGP#toc7 > Team Cymru Malware Hash Registry > http://www.team-cymru.org/MHR.html > Spamhaus (SBL, XBL, etc) > and others > QueryDNS will use an expression language enabled property to run a query > against DNS and add the raw result to an attribute (for later processing if > necessary). -- This message was sent by Atlassian JIRA (v6.3.4#6332)