eris

- Elasticsearch Recon Ingestion Scripts (ERIS) 🔎
git clone git://git.acid.vegas/-c.git
sniff_patch.py (3708B)
      1 #!/usr/bin/env python
      2 # Elasticsearch Recon Ingestion Scripts (ERIS) - Developed by Acidvegas (https://git.acid.vegas/eris)
      3 
      4 # Note:
      5 #   This is a patch for the elasticsearch 8.x client to fix the sniff_* options.
      6 #   This patch is only needed if you use the sniff_* options and only works with basic auth.
      7 #   Call init_elasticsearch() with normal Elasticsearch params.
      8 #
      9 # Source:
     10 #   - https://github.com/elastic/elasticsearch-py/issues/2005#issuecomment-1645641960
     11 
     12 import base64
     13 
     14 import elasticsearch._sync.client as client
     15 from elasticsearch.exceptions import SerializationError, ConnectionError
     16 
     17 
     18 def init_elasticsearch(*args, **kwargs):
     19     '''
     20     Initialize the Elasticsearch client with the sniff patch.
     21     
     22     :param args: Elasticsearch positional arguments.
     23     :param kwargs: Elasticsearch keyword arguments.
     24     '''
     25     client.default_sniff_callback = _override_sniff_callback(kwargs['basic_auth'])
     26 
     27     return client.Elasticsearch(*args, **kwargs)
     28 
     29 
     30 def _override_sniff_callback(basic_auth):
     31     '''
     32     Taken from https://github.com/elastic/elasticsearch-py/blob/8.8/elasticsearch/_sync/client/_base.py#L166
     33     Completely unmodified except for adding the auth header to the elastic request.
     34     Allows us to continue using the sniff_* options while this is broken in the library.
     35 
     36     TODO: Remove this when this issue is patched:
     37         - https://github.com/elastic/elasticsearch-py/issues/2005
     38     '''
     39     auth_str = base64.b64encode(':'.join(basic_auth).encode()).decode()
     40     sniffed_node_callback = client._base._default_sniffed_node_callback
     41 
     42     def modified_sniff_callback(transport, sniff_options):
     43         for _ in transport.node_pool.all():
     44             try:
     45                 meta, node_infos = transport.perform_request(
     46                     'GET',
     47                     '/_nodes/_all/http',
     48                     headers = {
     49                         'accept': 'application/vnd.elasticsearch+json; compatible-with=8',
     50                         'authorization': f'Basic {auth_str}' # This auth header is missing in 8.x releases of the client, and causes 401s
     51                     },
     52                     request_timeout = (
     53                         sniff_options.sniff_timeout
     54                         if not sniff_options.is_initial_sniff
     55                         else None
     56                     ),
     57                 )
     58             except (SerializationError, ConnectionError):
     59                 continue
     60 
     61             if not 200 <= meta.status <= 299:
     62                 continue
     63 
     64             node_configs = []
     65             for node_info in node_infos.get('nodes', {}).values():
     66                 address = node_info.get('http', {}).get('publish_address')
     67                 if not address or ':' not in address:
     68                     continue
     69 
     70                 if '/' in address:
     71                     # Support 7.x host/ip:port behavior where http.publish_host has been set.
     72                     fqdn, ipaddress = address.split('/', 1)
     73                     host = fqdn
     74                     _, port_str = ipaddress.rsplit(':', 1)
     75                     port = int(port_str)
     76                 else:
     77                     host, port_str = address.rsplit(':', 1)
     78                     port = int(port_str)
     79 
     80                 assert sniffed_node_callback is not None
     81                 sniffed_node = sniffed_node_callback(
     82                     node_info, meta.node.replace(host=host, port=port)
     83                 )
     84                 if sniffed_node is None:
     85                     continue
     86 
     87                 # Use the node which was able to make the request as a base.
     88                 node_configs.append(sniffed_node)
     89 
     90             if node_configs:
     91                 return node_configs
     92 
     93         return []
     94 
     95     return modified_sniff_callback