96 lines
3.8 KiB
Python
96 lines
3.8 KiB
Python
#!/usr/bin/env python
|
|
# Elasticsearch Recon Ingestion Scripts (ERIS) - Developed by Acidvegas (https://git.acid.vegas/eris)
|
|
# sniff_patch.py [asyncronous developement]
|
|
|
|
# Note:
|
|
# This is a patch for the elasticsearch 8.x client to fix the sniff_* options.
|
|
# This patch is only needed if you use the sniff_* options and only works with basic auth.
|
|
# Call init_elasticsearch() with normal Elasticsearch params.
|
|
#
|
|
# Source:
|
|
# - https://github.com/elastic/elasticsearch-py/issues/2005#issuecomment-1645641960
|
|
|
|
import base64
|
|
|
|
import elasticsearch._async.client as async_client
|
|
from elasticsearch.exceptions import SerializationError, ConnectionError
|
|
|
|
|
|
async def init_elasticsearch_async(*args, **kwargs):
|
|
'''
|
|
Initialize the Async Elasticsearch client with the sniff patch.
|
|
|
|
:param args: Async Elasticsearch positional arguments.
|
|
:param kwargs: Async Elasticsearch keyword arguments.
|
|
'''
|
|
async_client.default_sniff_callback = _override_async_sniff_callback(kwargs['basic_auth'])
|
|
|
|
return async_client.AsyncElasticsearch(*args, **kwargs)
|
|
|
|
|
|
def _override_async_sniff_callback(basic_auth):
|
|
'''
|
|
Taken from https://github.com/elastic/elasticsearch-py/blob/8.8/elasticsearch/_sync/client/_base.py#L166
|
|
Completely unmodified except for adding the auth header to the elastic request.
|
|
Allows us to continue using the sniff_* options while this is broken in the library.
|
|
|
|
TODO: Remove this when this issue is patched:
|
|
- https://github.com/elastic/elasticsearch-py/issues/2005
|
|
'''
|
|
auth_str = base64.b64encode(':'.join(basic_auth).encode()).decode()
|
|
sniffed_node_callback = async_client._base._default_sniffed_node_callback
|
|
|
|
async def modified_async_sniff_callback(transport, sniff_options):
|
|
for _ in transport.node_pool.all():
|
|
try:
|
|
meta, node_infos = await transport.perform_request(
|
|
'GET',
|
|
'/_nodes/_all/http',
|
|
headers={
|
|
'accept': 'application/vnd.elasticsearch+json; compatible-with=8',
|
|
'authorization': f'Basic {auth_str}' # This auth header is missing in 8.x releases of the client, and causes 401s
|
|
},
|
|
request_timeout=(
|
|
sniff_options.sniff_timeout
|
|
if not sniff_options.is_initial_sniff
|
|
else None
|
|
),
|
|
)
|
|
except (SerializationError, ConnectionError):
|
|
continue
|
|
|
|
if not 200 <= meta.status <= 299:
|
|
continue
|
|
|
|
node_configs = []
|
|
for node_info in node_infos.get('nodes', {}).values():
|
|
address = node_info.get('http', {}).get('publish_address')
|
|
if not address or ':' not in address:
|
|
continue
|
|
|
|
if '/' in address:
|
|
# Support 7.x host/ip:port behavior where http.publish_host has been set.
|
|
fqdn, ipaddress = address.split('/', 1)
|
|
host = fqdn
|
|
_, port_str = ipaddress.rsplit(':', 1)
|
|
port = int(port_str)
|
|
else:
|
|
host, port_str = address.rsplit(':', 1)
|
|
port = int(port_str)
|
|
|
|
assert sniffed_node_callback is not None
|
|
sniffed_node = await sniffed_node_callback(
|
|
node_info, meta.node.replace(host=host, port=port)
|
|
)
|
|
if sniffed_node is None:
|
|
continue
|
|
|
|
# Use the node which was able to make the request as a base.
|
|
node_configs.append(sniffed_node)
|
|
|
|
if node_configs:
|
|
return node_configs
|
|
|
|
return []
|
|
|
|
return modified_async_sniff_callback |