From 4595ddb9c1aa959dc97a0560f076ebf2c9b1c4c3 Mon Sep 17 00:00:00 2001 From: Kyle Maxwell Date: Tue, 17 Mar 2015 21:16:01 -0500 Subject: [PATCH 01/11] Whitespace fixes for PEP8 --- reaper.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/reaper.py b/reaper.py index a142586..7909c3b 100755 --- a/reaper.py +++ b/reaper.py @@ -8,9 +8,11 @@ logger = get_logger('reaper') + def exception_handler(request, exception): logger.error("Request %r failed: %r" % (request, exception)) + def reap(file_name): config = ConfigParser.SafeConfigParser(allow_no_value=False) cfg_success = config.read('combine.cfg') @@ -24,7 +26,7 @@ def reap(file_name): try: with open(inbound_url_file, 'rb') as f: - inbound_urls = [url.rstrip('\n') for url in f.readlines()] + inbound_urls = [url.rstrip('\n') for url in f.readlines()] except EnvironmentError as e: logger.error('Reaper: Error while opening "%s" - %s' % (inbound_url_file, e.strerror)) return @@ -40,7 +42,7 @@ def reap(file_name): headers = {'User-Agent': 'MLSecProject-Combine/0.1.2 (+https://github.com/mlsecproject/combine)'} logger.info('Fetching inbound URLs') - inbound_files=[] + inbound_files = [] for url in inbound_urls: if url.startswith('file://'): inbound_files.append(url.partition('://')[2]) @@ -50,14 +52,14 @@ def reap(file_name): inbound_harvest = [(response.url, response.status_code, response.text) for response in inbound_responses if response] for each in inbound_files: try: - with open(each,'rb') as f: - inbound_harvest.append(('file://'+each, 200, f.read())) + with open(each, 'rb') as f: + inbound_harvest.append(('file://' + each, 200, f.read())) except IOError as e: assert isinstance(logger, logging.Logger) logger.error('Reaper: Error while opening "%s" - %s' % (each, e.strerror)) logger.info('Fetching outbound URLs') - outbound_files=[] + outbound_files = [] for url in outbound_urls: if url.startswith('file://'): outbound_files.append(url.partition('://')[2]) @@ -67,8 +69,8 @@ def reap(file_name): outbound_harvest = [(response.url, response.status_code, response.text) for response in outbound_responses if response] for each in outbound_files: try: - with open(each,'rb') as f: - outbound_harvest.append(('file://'+each, 200, f.read())) + with open(each, 'rb') as f: + outbound_harvest.append(('file://' + each, 200, f.read())) except IOError as e: assert isinstance(logger, logging.Logger) logger.error('Reaper: Error while opening "%s" - %s' % (each, e.strerror)) From 73d6e112ddbdd50e5f203771ad3d07f3580ae583 Mon Sep 17 00:00:00 2001 From: Kyle Maxwell Date: Tue, 17 Mar 2015 22:20:05 -0500 Subject: [PATCH 02/11] Refactor to support testing --- combine.py | 94 ++++++++++++++++++++++++++++++++---------------------- 1 file changed, 56 insertions(+), 38 deletions(-) diff --git a/combine.py b/combine.py index c246b49..b36691b 100755 --- a/combine.py +++ b/combine.py @@ -14,41 +14,59 @@ logger = get_logger() -parser = argparse.ArgumentParser() -parser.add_argument('-t', '--type', help="Specify output type. Currently supported: CSV and exporting to CRITs") -parser.add_argument('-f', '--file', help="Specify output file. Defaults to harvest.FILETYPE") -parser.add_argument('-d', '--delete', help="Delete intermediate files", action="store_true") -parser.add_argument('-e', '--enrich', help="Enrich data", action="store_true") -parser.add_argument('--tiq-test', help="Output in tiq-test format", action="store_true") -args = parser.parse_args() - -possible_types = ['csv', 'json','crits'] - -if not args.type: - out_type = 'csv' -elif args.type.lower() not in possible_types: - sys.exit('Invalid file type specified. Possible types are: %s' % possible_types) -else: - out_type = args.type.lower() - -if args.file: - out_file = args.file -else: - out_file = 'harvest.'+out_type - -reap('harvest.json') -thresh('harvest.json', 'crop.json') -bale('crop.json', out_file, out_type, True) - -if args.enrich or args.tiq_test: - winnow('crop.json', 'crop.json', 'enrich.json') - bale('enrich.json', 'enriched.'+out_type, out_type, False) - -if args.tiq_test: - tiq_output('crop.json', 'enrich.json') - -if args.delete: - # be careful with this when we support a JSON output type - os.remove('harvest.json') - os.remove('crop.json') - os.remove('enrich.json') + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument('-t', '--type', help="Specify output type. Currently supported: CSV and exporting to CRITs") + parser.add_argument('-f', '--file', help="Specify output file. Defaults to harvest.FILETYPE") + parser.add_argument('-d', '--delete', help="Delete intermediate files", action="store_true") + parser.add_argument('-e', '--enrich', help="Enrich data", action="store_true") + parser.add_argument('--tiq-test', help="Output in tiq-test format", action="store_true") + return parser.parse_args() + + +def get_type(args): + possible_types = ['csv', 'json', 'crits'] + + if not args.type: + out_type = 'csv' + elif args.type.lower() not in possible_types: + sys.exit('Invalid file type specified. Possible types are: {}'.format(possible_types)) + else: + out_type = args.type.lower() + + return out_type + + +def get_file(args, out_type): + if args.file: + out_file = args.file + else: + out_file = 'harvest.' + out_type + + +def main(): + args = get_args() + out_type = get_type(args) + out_file = get_file(args) + + reap('harvest.json') + thresh('harvest.json', 'crop.json') + bale('crop.json', out_file, out_type, True) + + if args.enrich or args.tiq_test: + winnow('crop.json', 'crop.json', 'enrich.json') + bale('enrich.json', 'enriched.' + out_type, out_type, False) + + if args.tiq_test: + tiq_output('crop.json', 'enrich.json') + + if args.delete: + # be careful with this when we support a JSON output type + os.remove('harvest.json') + os.remove('crop.json') + os.remove('enrich.json') + + +if __name__ == "__main__": + main() From 18ae3ba17aa9a7602dce8a1c3da04718b663d561 Mon Sep 17 00:00:00 2001 From: Kyle Maxwell Date: Wed, 18 Mar 2015 14:30:24 -0500 Subject: [PATCH 03/11] Remove unused import and fix syntax for get_file() --- combine.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/combine.py b/combine.py index b36691b..c81ddbf 100755 --- a/combine.py +++ b/combine.py @@ -3,7 +3,6 @@ import argparse import os import sys -import logging # Combine components from logger import get_logger @@ -43,12 +42,13 @@ def get_file(args, out_type): out_file = args.file else: out_file = 'harvest.' + out_type + return out_file def main(): args = get_args() out_type = get_type(args) - out_file = get_file(args) + out_file = get_file(args, out_type) reap('harvest.json') thresh('harvest.json', 'crop.json') From 2c8972a36c74ae5aa568cf7b77497bd0e4950b3f Mon Sep 17 00:00:00 2001 From: Kyle Maxwell Date: Wed, 18 Mar 2015 14:35:43 -0500 Subject: [PATCH 04/11] Try to catch errors in major processes --- combine.py | 35 +++++++++++++++++++---------------- 1 file changed, 19 insertions(+), 16 deletions(-) diff --git a/combine.py b/combine.py index c81ddbf..a15844c 100755 --- a/combine.py +++ b/combine.py @@ -50,22 +50,25 @@ def main(): out_type = get_type(args) out_file = get_file(args, out_type) - reap('harvest.json') - thresh('harvest.json', 'crop.json') - bale('crop.json', out_file, out_type, True) - - if args.enrich or args.tiq_test: - winnow('crop.json', 'crop.json', 'enrich.json') - bale('enrich.json', 'enriched.' + out_type, out_type, False) - - if args.tiq_test: - tiq_output('crop.json', 'enrich.json') - - if args.delete: - # be careful with this when we support a JSON output type - os.remove('harvest.json') - os.remove('crop.json') - os.remove('enrich.json') + # TODO: possibly the wrong pattern here? + err = reap('harvest.json') + if not err: + err = thresh('harvest.json', 'crop.json') + if not err: + err = bale('crop.json', out_file, out_type, True) + if not err: + if args.enrich or args.tiq_test: + winnow('crop.json', 'crop.json', 'enrich.json') + bale('enrich.json', 'enriched.' + out_type, out_type, False) + + if args.tiq_test: + tiq_output('crop.json', 'enrich.json') + + if args.delete: + # be careful with this when we support a JSON output type + os.remove('harvest.json') + os.remove('crop.json') + os.remove('enrich.json') if __name__ == "__main__": From f466189b809ea4344fdd74524ae809563a87696f Mon Sep 17 00:00:00 2001 From: Kyle Maxwell Date: Wed, 18 Mar 2015 14:37:40 -0500 Subject: [PATCH 05/11] Remove another unused import --- reaper.py | 1 - 1 file changed, 1 deletion(-) diff --git a/reaper.py b/reaper.py index 7909c3b..4807932 100755 --- a/reaper.py +++ b/reaper.py @@ -1,7 +1,6 @@ import ConfigParser import grequests import json -import sys from logger import get_logger import logging From bce5f06aa4ba3abb853bf88d84d340585a60806e Mon Sep 17 00:00:00 2001 From: Kyle Maxwell Date: Wed, 18 Mar 2015 14:39:17 -0500 Subject: [PATCH 06/11] Remove unused import and specify regex strings --- thresher.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/thresher.py b/thresher.py index 81a95b4..d1c4155 100755 --- a/thresher.py +++ b/thresher.py @@ -4,17 +4,16 @@ import feedparser import json import re -import sys from logger import get_logger -import logging from csv import reader from itertools import ifilter logger = get_logger('thresher') + def indicator_type(indicator): - ip_regex = '^(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$' - domain_regex = '(www\.)?(?P
([\d\w.][-\d\w.]{0,253}[\d\w.]+\.)+(AC|AD|AE|AERO|AF|AG|AI|AL|AM|AN|AO|AQ|AR|ARPA|AS|ASIA|AT|AU|AW|AX|AZ|BA|BB|BD|BE|BF|BG|BH|BI|BIZ|BJ|BM|BN|BO|BR|BS|BT|BV|BW|BY|BZ|CA|CAT|CC|CD|CF|CG|CH|CI|CK|CL|CM|CN|COM|COOP|CR|CU|CV|CX|CY|CZ|DE|DJ|DK|DM|DO|DZ|EC|EDU|EE|EG|ER|ES|ET|EU|FI|FJ|FK|FM|FO|FR|GA|GB|GD|GE|GF|GG|GH|GI|GL|GM|GN|GOV|GP|GQ|GR|GS|GT|GU|GW|GY|HK|HM|HN|HR|HT|HU|ID|IE|IL|IM|IN|INFO|INT|IO|IQ|IR|IS|IT|JE|JM|JO|JOBS|JP|KE|KG|KH|KI|KM|KN|KP|KR|KW|KY|KZ|LA|LB|LC|LI|LK|LR|LS|LT|LU|LV|LY|MA|MC|MD|ME|MG|MH|MIL|MK|ML|MM|MN|MO|MOBI|MP|MQ|MR|MS|MT|MU|MUSEUM|MV|MW|MX|MY|MZ|NA|NAME|NC|NET|NF|NG|NI|NL|NO|NP|NR|NU|NZ|OM|ORG|PA|PE|PF|PG|PH|PK|PL|PM|PN|PR|PRO|PS|PT|PW|PY|QA|RE|RO|RS|RU|RW|SA|SB|SC|SD|SE|SG|SH|SI|SJ|SK|SL|SM|SN|SO|SR|ST|SU|SV|SY|SZ|TC|TD|TEL|TF|TG|TH|TJ|TK|TL|TM|TN|TO|TP|TR|TRAVEL|TT|TV|TW|TZ|UA|UG|UK|US|UY|UZ|VA|VC|VE|VG|VI|VN|VU|WF|WS|XN|XN|XN|XN|XN|XN|XN|XN|XN|XN|XN|YE|YT|YU|ZA|ZM|ZW))' + ip_regex = r'^(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$' + domain_regex = r'(www\.)?(?P
([\d\w.][-\d\w.]{0,253}[\d\w.]+\.)+(AC|AD|AE|AERO|AF|AG|AI|AL|AM|AN|AO|AQ|AR|ARPA|AS|ASIA|AT|AU|AW|AX|AZ|BA|BB|BD|BE|BF|BG|BH|BI|BIZ|BJ|BM|BN|BO|BR|BS|BT|BV|BW|BY|BZ|CA|CAT|CC|CD|CF|CG|CH|CI|CK|CL|CM|CN|COM|COOP|CR|CU|CV|CX|CY|CZ|DE|DJ|DK|DM|DO|DZ|EC|EDU|EE|EG|ER|ES|ET|EU|FI|FJ|FK|FM|FO|FR|GA|GB|GD|GE|GF|GG|GH|GI|GL|GM|GN|GOV|GP|GQ|GR|GS|GT|GU|GW|GY|HK|HM|HN|HR|HT|HU|ID|IE|IL|IM|IN|INFO|INT|IO|IQ|IR|IS|IT|JE|JM|JO|JOBS|JP|KE|KG|KH|KI|KM|KN|KP|KR|KW|KY|KZ|LA|LB|LC|LI|LK|LR|LS|LT|LU|LV|LY|MA|MC|MD|ME|MG|MH|MIL|MK|ML|MM|MN|MO|MOBI|MP|MQ|MR|MS|MT|MU|MUSEUM|MV|MW|MX|MY|MZ|NA|NAME|NC|NET|NF|NG|NI|NL|NO|NP|NR|NU|NZ|OM|ORG|PA|PE|PF|PG|PH|PK|PL|PM|PN|PR|PRO|PS|PT|PW|PY|QA|RE|RO|RS|RU|RW|SA|SB|SC|SD|SE|SG|SH|SI|SJ|SK|SL|SM|SN|SO|SR|ST|SU|SV|SY|SZ|TC|TD|TEL|TF|TG|TH|TJ|TK|TL|TM|TN|TO|TP|TR|TRAVEL|TT|TV|TW|TZ|UA|UG|UK|US|UY|UZ|VA|VC|VE|VG|VI|VN|VU|WF|WS|XN|XN|XN|XN|XN|XN|XN|XN|XN|XN|XN|YE|YT|YU|ZA|ZM|ZW))' if re.match(ip_regex, indicator): return "IPv4" @@ -40,7 +39,7 @@ def process_sans(response, source, direction): for line in response.splitlines(): if not line.startswith('#') and len(line) > 0: # Because SANS zero-pads their addresses - i = re.sub('\.0{1,2}', '.', line.split()[0].lstrip('0')) + i = re.sub(r'\.0{1,2}', '.', line.split()[0].lstrip('0')) date = line.split()[-1] data.append((i, indicator_type(i), direction, source, '', date)) return data @@ -185,7 +184,6 @@ def thresh(input_file, output_file): # When we have plugins, this hack won't be necessary for response in crop['inbound']: logger.info('Evaluating %s' % response[0]) - # TODO: logging if response[1] == 200: for site in thresher_map: if site in response[0]: From e15f153bf66296a1535ac2f844aa25d3d6886028 Mon Sep 17 00:00:00 2001 From: Kyle Maxwell Date: Wed, 18 Mar 2015 14:40:06 -0500 Subject: [PATCH 07/11] Remove unused import and specify regex strings --- winnower.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/winnower.py b/winnower.py index 60a30b7..2ee0eae 100755 --- a/winnower.py +++ b/winnower.py @@ -6,13 +6,10 @@ import json import pygeoip import re -import sys from netaddr import IPAddress, IPRange, IPSet from sortedcontainers import SortedDict - from logger import get_logger -import logging logger = get_logger('winnower') @@ -89,14 +86,14 @@ def reserved(address): def is_ipv4(address): - if re.match('(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$', address): + if re.match(r'(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$', address): return True else: return False def is_fqdn(address): - if re.match('(?=^.{4,255}$)(^((?!-)[a-zA-Z0-9-]{1,63}(? Date: Wed, 18 Mar 2015 14:41:15 -0500 Subject: [PATCH 08/11] Remove unused imports and unused parameter --- baler.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/baler.py b/baler.py index 5dde618..1f85a5a 100755 --- a/baler.py +++ b/baler.py @@ -2,11 +2,9 @@ import datetime as dt import gzip import json -import logging import os import re import requests -import sys import time import unicodecsv import threading @@ -139,7 +137,7 @@ def bale_CRITs_indicator(base_url, data, indicator_que): logger.info("don't yet know what to do with: %s[%s]" % (indicator[1], indicator[0])) -def bale_CRITs(harvest, filename): +def bale_CRITs(harvest): """ taking the output from combine and pushing it to the CRITs web API""" # checking the minimum requirements for parameters # it would be nice to have some metadata on the feeds that can be imported in the intel library: From 38652b535d5fd76c587f2f75707aa3e2cd008696 Mon Sep 17 00:00:00 2001 From: Kyle Maxwell Date: Sun, 26 Apr 2015 22:00:18 -0400 Subject: [PATCH 09/11] Add pre-commit --- .pre-commit-config.yaml | 17 +++++++++++++++++ requirements.txt | 1 + 2 files changed, 18 insertions(+) create mode 100644 .pre-commit-config.yaml diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..b887fa4 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,17 @@ +- repo: git://github.com/pre-commit/pre-commit-hooks + sha: master + hooks: + - id: autopep8-wrapper + args: ['-i', '--ignore=E501'] + - id: check-json + - id: check-yaml + - id: end-of-file-fixer + - id: flake8 + args: [--max-line-length=255] + - id: trailing-whitespace + +- repo: git://github.com/ivanlei/pre-commit-python-sorter + sha: master + hooks: + - id: python-import-sorter + args: [--silent-overwrite, --force_single_line] diff --git a/requirements.txt b/requirements.txt index b372c08..1437159 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,3 +11,4 @@ requests>=2.3.0,<2.6.0 sortedcontainers==0.9.4 wsgiref==0.1.2 unicodecsv==0.9.4 +pre-commit From c746e2a0a11b5dca580f2217298987608f04a840 Mon Sep 17 00:00:00 2001 From: Kyle Maxwell Date: Sun, 26 Apr 2015 22:13:45 -0400 Subject: [PATCH 10/11] Move enrichment checks to functions --- winnower.py | 53 ++++++++++++++++++++++++++++++++--------------------- 1 file changed, 32 insertions(+), 21 deletions(-) diff --git a/winnower.py b/winnower.py index c0b52f0..a46a974 100755 --- a/winnower.py +++ b/winnower.py @@ -2,14 +2,16 @@ import ConfigParser import csv import datetime as dt -import dnsdb_query import json -import pygeoip import re -from netaddr import IPAddress, IPRange, IPSet -from sortedcontainers import SortedDict +import dnsdb_query +import pygeoip from logger import get_logger +from netaddr import IPAddress +from netaddr import IPRange +from netaddr import IPSet +from sortedcontainers import SortedDict logger = get_logger('winnower') @@ -21,13 +23,13 @@ def load_gi_org(filename): + # no return function because gi_org is scoped to the module + # ugly hack with open(filename, 'rb') as f: org_reader = csv.DictReader(f, fieldnames=['start', 'end', 'org']) for row in org_reader: gi_org[row['start']] = (IPRange(row['start'], row['end']), unicode(row['org'], errors='replace')) - return gi_org - def org_by_addr(address): as_num = None @@ -44,7 +46,6 @@ def maxhits(dns_records): hmax = 0 hostname = None for record in dns_records: - #logger.info("examining %s" % record) if record['count'] > hmax: hmax = record['count'] hostname = record['rrname'].rstrip('.') @@ -118,16 +119,7 @@ def is_fqdn(address): return False -def winnow(in_file, out_file, enr_file): - config = ConfigParser.SafeConfigParser(allow_no_value=True) - cfg_success = config.read('combine.cfg') - if not cfg_success: - logger.error('Winnower: Could not read combine.cfg.') - logger.error('HINT: edit combine-example.cfg and save as combine.cfg.') - return - - server = config.get('Winnower', 'dnsdb_server') - api = config.get('Winnower', 'dnsdb_api') +def check_enrich_ip(config): enrich_ip = config.get('Winnower', 'enrich_ip') if enrich_ip == '1' or enrich_ip == 'True': enrich_ip = True @@ -135,7 +127,10 @@ def winnow(in_file, out_file, enr_file): else: enrich_ip = False logger.info('Enriching IPv4 indicators: FALSE') + return enrich_ip + +def check_enrich_dns(config): enrich_dns = config.get('Winnower', 'enrich_dns') if enrich_dns == '1' or enrich_dns == 'True': enrich_dns = True @@ -143,9 +138,27 @@ def winnow(in_file, out_file, enr_file): else: enrich_dns = False logger.info('Enriching DNS indicators: FALSE') + return enrich_dns - logger.info('Setting up DNSDB client') +def winnow(in_file, out_file, enr_file): + config = ConfigParser.SafeConfigParser(allow_no_value=True) + cfg_success = config.read('combine.cfg') + if not cfg_success: + logger.error('Winnower: Could not read combine.cfg.') + logger.error('HINT: edit combine-example.cfg and save as combine.cfg.') + return + + enrich_ip = check_enrich_ip(config) + enrich_dns = check_enrich_dns(config) + if enrich_dns: + server = config.get('Winnower', 'dnsdb_server') + api = config.get('Winnower', 'dnsdb_api') + else: + server = None + api = None + + logger.info('Setting up DNSDB client') # handle the case where we aren't using DNSDB dnsdb = dnsdb_query.DnsdbClient(server, api) if api == 'YOUR_API_KEY_HERE' or len(dnsdb.query_rdata_name('google.com')) == 0: @@ -157,7 +170,7 @@ def winnow(in_file, out_file, enr_file): # TODO: make these locations configurable? logger.info('Loading GeoIP data') - gi_org = load_gi_org('data/GeoIPASNum2.csv') + load_gi_org('data/GeoIPASNum2.csv') wheat = [] enriched = [] @@ -167,7 +180,6 @@ def winnow(in_file, out_file, enr_file): (addr, addr_type, direction, source, note, date) = each # this should be refactored into appropriate functions if addr_type == 'IPv4' and is_ipv4(addr): - #logger.info('Enriching %s' % addr) ipaddr = IPAddress(addr) if not reserved(ipaddr): wheat.append(each) @@ -180,7 +192,6 @@ def winnow(in_file, out_file, enr_file): else: logger.error('Found invalid address: %s from: %s' % (addr, source)) elif addr_type == 'FQDN' and is_fqdn(addr): - #logger.info('Enriching %s' % addr) wheat.append(each) if enrich_dns and dnsdb: # print "Enriching %s" % addr From 3499f2e29a1eeb1a48c5576653f3ffbbee8949aa Mon Sep 17 00:00:00 2001 From: Kyle Maxwell Date: Sun, 26 Apr 2015 22:16:45 -0400 Subject: [PATCH 11/11] Move DNSDB setup to function --- winnower.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/winnower.py b/winnower.py index a46a974..4162887 100755 --- a/winnower.py +++ b/winnower.py @@ -141,6 +141,16 @@ def check_enrich_dns(config): return enrich_dns +def setup_dnsdb(server, api): + logger.info('Setting up DNSDB client') + # handle the case where we aren't using DNSDB + dnsdb = dnsdb_query.DnsdbClient(server, api) + if api == 'YOUR_API_KEY_HERE' or len(dnsdb.query_rdata_name('google.com')) == 0: + dnsdb = None + logger.info('Invalid DNSDB configuration found') + return dnsdb + + def winnow(in_file, out_file, enr_file): config = ConfigParser.SafeConfigParser(allow_no_value=True) cfg_success = config.read('combine.cfg') @@ -151,19 +161,17 @@ def winnow(in_file, out_file, enr_file): enrich_ip = check_enrich_ip(config) enrich_dns = check_enrich_dns(config) - if enrich_dns: + if enrich_dns or enrich_ip: server = config.get('Winnower', 'dnsdb_server') api = config.get('Winnower', 'dnsdb_api') else: server = None api = None - logger.info('Setting up DNSDB client') - # handle the case where we aren't using DNSDB - dnsdb = dnsdb_query.DnsdbClient(server, api) - if api == 'YOUR_API_KEY_HERE' or len(dnsdb.query_rdata_name('google.com')) == 0: + if server and api: + dnsdb = setup_dnsdb(server, api) + else: dnsdb = None - logger.info('Invalid DNSDB configuration found') with open(in_file, 'rb') as f: crop = json.load(f)