Code cleanup

This commit is contained in:
Dionysus 2024-03-11 22:33:18 -04:00
parent b018da4e4d
commit 87f2cf27ea
Signed by: acidvegas
GPG Key ID: EF4B922DB85DC9DE
3 changed files with 352 additions and 421 deletions

View File

@ -5,85 +5,30 @@
import asyncio
import json
import logging
import time
try:
import websockets
except ImportError:
raise ImportError('Missing required \'websockets\' library. (pip install websockets)')
# Set a default elasticsearch index if one is not provided
default_index = 'cert-stream'
def construct_map() -> dict:
'''Construct the Elasticsearch index mapping for Certstream records.'''
# Match on exact value or full text search
keyword_mapping = { 'type': 'text', 'fields': { 'keyword': { 'type': 'keyword', 'ignore_above': 256 } } }
# Construct the index mapping
mapping = {
'mappings': {
'properties' : {
'data': {
'properties': {
'cert_index': { 'type': 'integer' },
'cert_link' : { 'type': 'keyword' },
'leaf_cert' : {
'properties': {
'all_domains': { 'type': 'keyword' },
'extensions': {
'properties': {
'authorityInfoAccess' : { 'type': 'text' },
'authorityKeyIdentifier' : { 'type': 'text' },
'basicConstraints' : { 'type': 'text' },
'certificatePolicies' : { 'type': 'text' },
'crlDistributionPoints' : { 'type': 'text' },
'ctlPoisonByte' : { 'type': 'boolean' },
'extendedKeyUsage' : { 'type': 'text' },
'keyUsage' : { 'type': 'text' },
'subjectAltName' : { 'type': 'text' },
'subjectKeyIdentifier' : { 'type': 'text' }
}
},
'fingerprint': { 'type': 'keyword' },
'issuer': {
'properties': {
'C' : { 'type': 'keyword' },
'CN' : { 'type': 'text' },
'L' : { 'type': 'text' },
'O' : { 'type': 'text' },
'OU' : { 'type': 'text' },
'ST' : { 'type': 'text' },
'aggregated' : { 'type': 'text' },
'emailAddress' : { 'type': 'text' }
}
},
'not_after' : { 'type': 'integer' },
'not_before' : { 'type': 'integer' },
'serial_number' : { 'type': 'keyword' },
'signature_algorithm' : { 'type': 'text' },
'subject': {
'properties': {
'C' : { 'type': 'keyword' },
'CN' : { 'type': 'text' },
'L' : { 'type': 'text' },
'O' : { 'type': 'text' },
'OU' : { 'type': 'text' },
'ST' : { 'type': 'text' },
'aggregated' : { 'type': 'text' },
'emailAddress' : { 'type': 'text' }
}
}
}
},
'seen': { 'type': 'date', 'format': 'epoch_second' },
'source': {
'properties': {
'name' : { 'type': 'keyword' },
'url' : { 'type': 'keyword' }
}
},
'update_type': { 'type': 'keyword' }
}
},
'message_type': { 'type': 'keyword' }
'domain' : keyword_mapping,
'seen' : { 'type': 'date' }
}
}
}
@ -100,21 +45,30 @@ async def process_data(place_holder: str = None):
while True:
try:
async with websockets.connect('wss://certstream.calidog.io/') as websocket:
async with websockets.connect('wss://certstream.calidog.io') as websocket:
while True:
# Read a line from the websocket
line = await websocket.recv()
if line == '~eof': # Sentinel value to indicate the end of a process (Used with --watch with FIFO)
break
# Parse the JSON record
try:
record = json.loads(line)
except json.decoder.JSONDecodeError:
logging.error(f'Failed to parse JSON record from Certstream! ({line})')
input('Press Enter to continue...')
logging.error(f'Invalid line from the websocket: {line}')
continue
yield record
# Grab the unique domains from the record (excluding wildcards)
domains = record['data']['leaf_cert']['all_domains']
domains = set([domain[2:] if domain.startswith('*.') else domain for domain in domains])
# Construct the document
for domain in domains:
struct = {
'domain' : domain,
'seen' : time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())
}
yield {'_id': id, '_index': default_index, '_source': struct}
except websockets.ConnectionClosed:
logging.error('Connection to Certstream was closed. Attempting to reconnect...')
@ -122,40 +76,11 @@ async def process_data(place_holder: str = None):
except Exception as e:
logging.error(f'An error occurred while processing Certstream records! ({e})')
await asyncio.sleep(15)
async def strip_struct_empty(data: dict) -> dict:
'''
Recursively remove empty values from a nested dictionary or list.
:param data: The dictionary or list to clean.
'''
empties = [None, '', [], {}]
if isinstance(data, dict):
for key, value in list(data.items()):
if value in empties:
del data[key]
else:
cleaned_value = strip_struct_empty(value)
if cleaned_value in empties:
del data[key]
else:
data[key] = cleaned_value
return data
elif isinstance(data, list):
return [strip_struct_empty(item) for item in data if item not in empties and strip_struct_empty(item) not in empties]
else:
return data
break
async def test():
'''Test the Cert stream ingestion process'''
'''Test the ingestion process.'''
async for document in process_data():
print(document)
@ -163,14 +88,9 @@ async def test():
if __name__ == '__main__':
import argparse
import asyncio
parser = argparse.ArgumentParser(description='Certstream Ingestor for ERIS')
parser.add_argument('input_path', help='Path to the input file or directory')
args = parser.parse_args()
asyncio.run(test(args.input_path))
asyncio.run(test())

View File

@ -11,15 +11,20 @@ except ImportError:
raise ImportError('Missing required \'aiofiles\' library. (pip install aiofiles)')
# Set a default elasticsearch index if one is not provided
default_index = 'dns-zones'
# Known DNS record types found in zone files
record_types = ('a','aaaa','caa','cdnskey','cds','cname','dnskey','ds','mx','naptr','ns','nsec','nsec3','nsec3param','ptr','rrsig','rp','sshfp','soa','srv','txt','type65534')
def construct_map() -> dict:
'''Construct the Elasticsearch index mapping for zone file records.'''
# Match on exact value or full text search
keyword_mapping = { 'type': 'text', 'fields': { 'keyword': { 'type': 'keyword', 'ignore_above': 256 } } }
# Construct the index mapping
mapping = {
'mappings': {
'properties': {
@ -31,18 +36,11 @@ def construct_map() -> dict:
}
# Add record types to mapping dynamically to not clutter the code
for item in record_types:
if item in ('a','aaaa'):
mapping['mappings']['properties']['records']['properties'][item] = {
for record_type in record_types:
if record_type in ('a','aaaa'):
mapping['mappings']['properties']['records']['properties'][record_type] = {
'properties': {
'data': { 'type': 'ip' },
'ttl': { 'type': 'integer' }
}
}
else:
mapping['mappings']['properties']['records']['properties'][item] = {
'properties': {
'data': keyword_mapping,
'data': { 'type': 'ip' if record_type in ('a','aaaa') else keyword_mapping},
'ttl': { 'type': 'integer' }
}
}
@ -52,43 +50,53 @@ def construct_map() -> dict:
async def process_data(file_path: str):
'''
Read and process zone file records.
Read and process the input file
:param file_path: Path to the zone file
:param input_path: Path to the input file
'''
async with aiofiles.open(file_path) as input_file:
# Initialize the cache
last = None
# Read the input file line by line
async for line in input_file:
line = line.strip()
if line == '~eof': # Sentinel value to indicate the end of a process (Used with --watch with FIFO)
return last
# Sentinel value to indicate the end of a process (for closing out a FIFO stream)
if line == '~eof':
yield last
break
# Skip empty lines and comments
if not line or line.startswith(';'):
continue
# Split the line into its parts
parts = line.split()
# Ensure the line has at least 3 parts
if len(parts) < 5:
logging.warning(f'Invalid line: {line}')
continue
# Split the record into its parts
domain, ttl, record_class, record_type, data = parts[0].rstrip('.').lower(), parts[1], parts[2].lower(), parts[3].lower(), ' '.join(parts[4:])
# Ensure the TTL is a number
if not ttl.isdigit():
logging.warning(f'Invalid TTL: {ttl} with line: {line}')
continue
else:
ttl = int(ttl)
# Anomaly...Doubtful any CHAOS/HESIOD records will be found in zone files
# Do not index other record classes (doubtful any CHAOS/HESIOD records will be found in zone files)
if record_class != 'in':
logging.warning(f'Unsupported record class: {record_class} with line: {line}')
continue
# We do not want to collide with our current mapping (Again, this is an anomaly)
# Do not index other record types
if record_type not in record_types:
logging.warning(f'Unsupported record type: {record_type} with line: {line}')
continue
@ -101,16 +109,18 @@ async def process_data(file_path: str):
elif data.endswith('.'):
data = data.rstrip('.')
# Check if we are still processing the same domain
if last:
if domain == last['domain']:
if domain == last['domain']: # This record is for the same domain as the cached document
if record_type in last['_doc']['records']:
last['_doc']['records'][record_type].append({'ttl': ttl, 'data': data}) # Do we need to check for duplicate records?
else:
last['_doc']['records'][record_type] = [{'ttl': ttl, 'data': data}]
continue
else:
yield last
yield last # Return the last document and start a new one
# Cache the document
last = {
'_op_type' : 'update',
'_id' : domain,
@ -126,10 +136,11 @@ async def process_data(file_path: str):
async def test(input_path: str):
'''
Test the Zone file ingestion process
Test the ingestion process
:param input_path: Path to the MassDNS log file
:param input_path: Path to the input file
'''
async for document in process_data(input_path):
print(document)
@ -139,7 +150,7 @@ if __name__ == '__main__':
import argparse
import asyncio
parser = argparse.ArgumentParser(description='Zone file Ingestor for ERIS')
parser = argparse.ArgumentParser(description='Ingestor for ERIS')
parser.add_argument('input_path', help='Path to the input file or directory')
args = parser.parse_args()