Add socks proxy support and lab test environment (#17622)

* Add support for ssh tunnel

* Patched create_connection method

* Add proxy option to esxi check

* Update changelog

* Update config defaults and only run tests on correct environments

* Add unit tests for proxy

---------

Co-authored-by: Jose Manuel Almaza <josemanuel.almaza@datadoghq.com>

Author: sarah-witt

esxi/assets/configuration/spec.yaml

@@ -5,6 +5,13 @@ files:
   - template: init_config
     options:
     - template: init_config/default
+    - name: proxy
+      value:
+        example: socks5://<PROXY_SERVER>:<PORT>
+        display_default: ""
+        type: string
+      description: |
+        SOCKS Proxy to use for all instances.
   - template: instances
     options:
         - name: host
@@ -162,6 +169,13 @@ files:
             display_default: []
             example:
               - <HOST_TAG>
+        - name: proxy
+          value:
+            example: socks5://<PROXY_SERVER>:<PORT>
+            display_default: ""
+            type: string
+          description: |
+            SOCKS Proxy to use for this instance. Overrides the `proxy` setting in `init_config`.
         - name: ssl_verify
           description: Set to false to disable SSL verification when connecting to the ESXi host
           value:

esxi/changelog.d/17622.added

@@ -0,0 +1 @@
+Add socks proxy support.

esxi/datadog_checks/esxi/check.py

@@ -3,9 +3,12 @@
 # Licensed under a 3-clause BSD style license (see LICENSE)
 import logging
 import re
+import socket
 import ssl
 from collections import defaultdict
+from urllib.parse import urlparse
 
+import socks
 from pyVim import connect
 from pyVmomi import vim, vmodl
 from six import iteritems
@@ -36,6 +39,10 @@
 )
 
 
+def create_connection(address, timeout, source_address, proxy_host, proxy_port):
+    return socks.create_connection(address, timeout, source_address, socks.SOCKS5, proxy_host, proxy_port)
+
+
 class EsxiCheck(AgentCheck):
     __NAMESPACE__ = 'esxi'
 
@@ -56,6 +63,17 @@ def __init__(self, name, init_config, instances):
         self.ssl_capath = self.instance.get("ssl_capath")
         self.ssl_cafile = self.instance.get("ssl_cafile")
         self.tags = [f"esxi_url:{self.host}"]
+        self.proxy_host = None
+        self.proxy_port = None
+        proxy = self.instance.get('proxy', init_config.get('proxy'))
+        if proxy:
+            parsed_proxy = urlparse(proxy)
+            proxy_scheme = parsed_proxy.scheme
+            if proxy_scheme != 'socks5':
+                self.log.warning('Proxy scheme %s not supported; ignoring', proxy_scheme)
+            else:
+                self.proxy_host = parsed_proxy.hostname
+                self.proxy_port = parsed_proxy.port
 
     def _validate_excluded_host_tags(self, excluded_host_tags):
         valid_excluded_host_tags = []
@@ -344,7 +362,15 @@ def check(self, _):
             else:
                 context.load_default_certs(ssl.Purpose.SERVER_AUTH)
 
+            create_connection_method = socket.create_connection
+            if self.proxy_host:
+                socket.create_connection = lambda address, timeout, source_address, **kwargs: create_connection(
+                    address, timeout, source_address, self.proxy_host, self.proxy_port
+                )
+
             connection = connect.SmartConnect(host=self.host, user=self.username, pwd=self.password, sslContext=context)
+            socket.create_connection = create_connection_method
+
             self.conn = connection
             self.content = connection.content
 

esxi/datadog_checks/esxi/config_models/defaults.py

@@ -8,6 +8,10 @@
 #     ddev -x validate models -s <INTEGRATION_NAME>
 
 
+def shared_proxy():
+    return ''
+
+
 def instance_disable_generic_tags():
     return False
 
@@ -20,6 +24,10 @@ def instance_min_collection_interval():
     return 15
 
 
+def instance_proxy():
+    return ''
+
+
 def instance_ssl_verify():
     return True
 

esxi/datadog_checks/esxi/config_models/instance.py

@@ -72,6 +72,7 @@ class InstanceConfig(BaseModel):
     metric_patterns: Optional[MetricPatterns] = None
     min_collection_interval: Optional[float] = None
     password: str
+    proxy: Optional[str] = None
     resource_filters: Optional[tuple[ResourceFilter, ...]] = None
     service: Optional[str] = None
     ssl_cafile: Optional[str] = None

esxi/datadog_checks/esxi/config_models/shared.py

@@ -16,7 +16,7 @@
 from datadog_checks.base.utils.functions import identity
 from datadog_checks.base.utils.models import validation
 
-from . import validators
+from . import defaults, validators
 
 
 class SharedConfig(BaseModel):
@@ -25,6 +25,7 @@ class SharedConfig(BaseModel):
         arbitrary_types_allowed=True,
         frozen=True,
     )
+    proxy: Optional[str] = None
     service: Optional[str] = None
 
     @model_validator(mode='before')
@@ -37,6 +38,8 @@ def _validate(cls, value, info):
         field_name = field.alias or info.field_name
         if field_name in info.context['configured_fields']:
             value = getattr(validators, f'shared_{info.field_name}', identity)(value, field=field)
+        else:
+            value = getattr(defaults, f'shared_{info.field_name}', lambda: value)()
 
         return validation.utils.make_immutable(value)
 

esxi/datadog_checks/esxi/data/conf.yaml.example

@@ -9,6 +9,11 @@ init_config:
     #
     # service: <SERVICE>
 
+    ## @param proxy - string - optional - default: 
+    ## SOCKS Proxy to use for all instances.
+    #
+    # proxy: socks5://<PROXY_SERVER>:<PORT>
+
 ## Every instance is scheduled independently of the others.
 #
 instances:
@@ -128,6 +133,11 @@ instances:
     # excluded_host_tags:
     #   - <HOST_TAG>
 
+    ## @param proxy - string - optional - default: 
+    ## SOCKS Proxy to use for this instance. Overrides the `proxy` setting in `init_config`.
+    #
+    # proxy: socks5://<PROXY_SERVER>:<PORT>
+
     ## @param ssl_verify - boolean - optional - default: true
     ## Set to false to disable SSL verification when connecting to the ESXi host
     #

esxi/hatch.toml

@@ -2,8 +2,20 @@
 
 [[envs.default.matrix]]
 python = ["3.11"]
+setup = ["lab"]
+
+[[envs.default.matrix]]
+python = ["3.11"]
+setup = ["vcsim"]
 version = ["6.5", "7.0"]
-type = ["vcsim"]
 
 [envs.default.overrides]
-matrix.version.env-vars = "VSPHERE_VERSION"
+name."vcsim".e2e-env = { value = true }
+matrix.version.env-vars = "VSPHERE_VERSION"
+matrix.setup.e2e-env = { value = true, if = ["lab"], env = ["VSPHERE_LAB_IP", "VSPHERE_LAB_PRIVATE_KEY_PATH", "VSPHERE_LAB_USERNAME", "ESXI_HOST_IP", "ESXI_HOST_USERNAME", "ESXI_HOST_PASSWORD"] }
+matrix.setup.env-vars = [
+  { key = "USE_VSPHERE_LAB", value = "True", if = ["lab"] },
+]
+
+[envs.default]
+e2e-env = false

esxi/pyproject.toml

@@ -38,6 +38,7 @@ dynamic = [
 [project.optional-dependencies]
 deps = [
     "pyvmomi==8.0.2.0.1; python_version > '3.0'",
+    "pysocks==1.7.1"
 ]
 
 [project.urls]

esxi/tests/common.py

@@ -1,10 +1,14 @@
 # (C) Datadog, Inc. 2024-present
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
+import os
+
 from pyVmomi import vim, vmodl
 
 HOST = "127.0.0.1"
 PORT = 8989
+USE_VSPHERE_LAB = os.environ.get('USE_VSPHERE_LAB')
+
 VCSIM_INSTANCE = {
     'host': f"{HOST}:{str(PORT)}",
     'username': 'test',

esxi/tests/conftest.py

@@ -1,6 +1,7 @@
 # (C) Datadog, Inc. 2024-present
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
+import copy
 import os
 
 import pytest
@@ -19,19 +20,42 @@
     PERF_METRIC_ID,
     PORT,
     PROPERTIES_EX,
+    USE_VSPHERE_LAB,
     VCSIM_INSTANCE,
 )
+from .ssh_tunnel import socks_proxy
 
 
 @pytest.fixture(scope='session')
 def dd_environment():
-    compose_file = os.path.join(get_here(), 'docker', 'docker-compose.yaml')
-    conditions = [
-        WaitForPortListening(HOST, PORT, wait=10),
-        CheckDockerLogs(compose_file, 'export GOVC_URL', wait=10),
-    ]
-    with docker_run(compose_file, conditions=conditions):
-        yield VCSIM_INSTANCE
+    if not USE_VSPHERE_LAB:
+        compose_file = os.path.join(get_here(), 'docker', 'docker-compose.yaml')
+        conditions = [
+            WaitForPortListening(HOST, PORT, wait=10),
+            CheckDockerLogs(compose_file, 'export GOVC_URL', wait=10),
+        ]
+        with docker_run(compose_file, conditions=conditions):
+            yield VCSIM_INSTANCE
+    else:
+        lab_public_ip = os.environ.get('VSPHERE_LAB_IP')
+        lab_private_key = os.environ.get('VSPHERE_LAB_PRIVATE_KEY_PATH')
+        lab_username = os.environ.get('VSPHERE_LAB_USERNAME')
+        with socks_proxy(
+            lab_public_ip,
+            lab_username,
+            lab_private_key,
+        ) as socks:
+
+            socks_ip, socks_port = socks
+            print(f"socks: {socks}")
+            instance = copy.deepcopy(VCSIM_INSTANCE)
+
+            instance['host'] = os.environ.get('ESXI_HOST_IP')
+            instance['username'] = os.environ.get('ESXI_HOST_USERNAME')
+            instance['password'] = os.environ.get('ESXI_HOST_PASSWORD')
+            instance['ssl_verify'] = 'false'
+            instance['proxy'] = 'socks5://{}:{}'.format(socks_ip, socks_port)
+            yield instance
 
 
 @pytest.fixture