Skip to content

Commit 051148d

Browse files
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>
1 parent e35457c commit 051148d

File tree

14 files changed

+311
-22
lines changed

14 files changed

+311
-22
lines changed

esxi/assets/configuration/spec.yaml

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,13 @@ files:
55
- template: init_config
66
options:
77
- template: init_config/default
8+
- name: proxy
9+
value:
10+
example: socks5://<PROXY_SERVER>:<PORT>
11+
display_default: ""
12+
type: string
13+
description: |
14+
SOCKS Proxy to use for all instances.
815
- template: instances
916
options:
1017
- name: host
@@ -162,6 +169,13 @@ files:
162169
display_default: []
163170
example:
164171
- <HOST_TAG>
172+
- name: proxy
173+
value:
174+
example: socks5://<PROXY_SERVER>:<PORT>
175+
display_default: ""
176+
type: string
177+
description: |
178+
SOCKS Proxy to use for this instance. Overrides the `proxy` setting in `init_config`.
165179
- name: ssl_verify
166180
description: Set to false to disable SSL verification when connecting to the ESXi host
167181
value:

esxi/changelog.d/17622.added

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Add socks proxy support.

esxi/datadog_checks/esxi/check.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,12 @@
33
# Licensed under a 3-clause BSD style license (see LICENSE)
44
import logging
55
import re
6+
import socket
67
import ssl
78
from collections import defaultdict
9+
from urllib.parse import urlparse
810

11+
import socks
912
from pyVim import connect
1013
from pyVmomi import vim, vmodl
1114
from six import iteritems
@@ -36,6 +39,10 @@
3639
)
3740

3841

42+
def create_connection(address, timeout, source_address, proxy_host, proxy_port):
43+
return socks.create_connection(address, timeout, source_address, socks.SOCKS5, proxy_host, proxy_port)
44+
45+
3946
class EsxiCheck(AgentCheck):
4047
__NAMESPACE__ = 'esxi'
4148

@@ -56,6 +63,17 @@ def __init__(self, name, init_config, instances):
5663
self.ssl_capath = self.instance.get("ssl_capath")
5764
self.ssl_cafile = self.instance.get("ssl_cafile")
5865
self.tags = [f"esxi_url:{self.host}"]
66+
self.proxy_host = None
67+
self.proxy_port = None
68+
proxy = self.instance.get('proxy', init_config.get('proxy'))
69+
if proxy:
70+
parsed_proxy = urlparse(proxy)
71+
proxy_scheme = parsed_proxy.scheme
72+
if proxy_scheme != 'socks5':
73+
self.log.warning('Proxy scheme %s not supported; ignoring', proxy_scheme)
74+
else:
75+
self.proxy_host = parsed_proxy.hostname
76+
self.proxy_port = parsed_proxy.port
5977

6078
def _validate_excluded_host_tags(self, excluded_host_tags):
6179
valid_excluded_host_tags = []
@@ -344,7 +362,15 @@ def check(self, _):
344362
else:
345363
context.load_default_certs(ssl.Purpose.SERVER_AUTH)
346364

365+
create_connection_method = socket.create_connection
366+
if self.proxy_host:
367+
socket.create_connection = lambda address, timeout, source_address, **kwargs: create_connection(
368+
address, timeout, source_address, self.proxy_host, self.proxy_port
369+
)
370+
347371
connection = connect.SmartConnect(host=self.host, user=self.username, pwd=self.password, sslContext=context)
372+
socket.create_connection = create_connection_method
373+
348374
self.conn = connection
349375
self.content = connection.content
350376

esxi/datadog_checks/esxi/config_models/defaults.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,10 @@
88
# ddev -x validate models -s <INTEGRATION_NAME>
99

1010

11+
def shared_proxy():
12+
return ''
13+
14+
1115
def instance_disable_generic_tags():
1216
return False
1317

@@ -20,6 +24,10 @@ def instance_min_collection_interval():
2024
return 15
2125

2226

27+
def instance_proxy():
28+
return ''
29+
30+
2331
def instance_ssl_verify():
2432
return True
2533

esxi/datadog_checks/esxi/config_models/instance.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ class InstanceConfig(BaseModel):
7272
metric_patterns: Optional[MetricPatterns] = None
7373
min_collection_interval: Optional[float] = None
7474
password: str
75+
proxy: Optional[str] = None
7576
resource_filters: Optional[tuple[ResourceFilter, ...]] = None
7677
service: Optional[str] = None
7778
ssl_cafile: Optional[str] = None

esxi/datadog_checks/esxi/config_models/shared.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
from datadog_checks.base.utils.functions import identity
1717
from datadog_checks.base.utils.models import validation
1818

19-
from . import validators
19+
from . import defaults, validators
2020

2121

2222
class SharedConfig(BaseModel):
@@ -25,6 +25,7 @@ class SharedConfig(BaseModel):
2525
arbitrary_types_allowed=True,
2626
frozen=True,
2727
)
28+
proxy: Optional[str] = None
2829
service: Optional[str] = None
2930

3031
@model_validator(mode='before')
@@ -37,6 +38,8 @@ def _validate(cls, value, info):
3738
field_name = field.alias or info.field_name
3839
if field_name in info.context['configured_fields']:
3940
value = getattr(validators, f'shared_{info.field_name}', identity)(value, field=field)
41+
else:
42+
value = getattr(defaults, f'shared_{info.field_name}', lambda: value)()
4043

4144
return validation.utils.make_immutable(value)
4245

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,11 @@ init_config:
99
#
1010
# service: <SERVICE>
1111

12+
## @param proxy - string - optional - default:
13+
## SOCKS Proxy to use for all instances.
14+
#
15+
# proxy: socks5://<PROXY_SERVER>:<PORT>
16+
1217
## Every instance is scheduled independently of the others.
1318
#
1419
instances:
@@ -128,6 +133,11 @@ instances:
128133
# excluded_host_tags:
129134
# - <HOST_TAG>
130135

136+
## @param proxy - string - optional - default:
137+
## SOCKS Proxy to use for this instance. Overrides the `proxy` setting in `init_config`.
138+
#
139+
# proxy: socks5://<PROXY_SERVER>:<PORT>
140+
131141
## @param ssl_verify - boolean - optional - default: true
132142
## Set to false to disable SSL verification when connecting to the ESXi host
133143
#

esxi/hatch.toml

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,20 @@
22

33
[[envs.default.matrix]]
44
python = ["3.11"]
5+
setup = ["lab"]
6+
7+
[[envs.default.matrix]]
8+
python = ["3.11"]
9+
setup = ["vcsim"]
510
version = ["6.5", "7.0"]
6-
type = ["vcsim"]
711

812
[envs.default.overrides]
9-
matrix.version.env-vars = "VSPHERE_VERSION"
13+
name."vcsim".e2e-env = { value = true }
14+
matrix.version.env-vars = "VSPHERE_VERSION"
15+
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"] }
16+
matrix.setup.env-vars = [
17+
{ key = "USE_VSPHERE_LAB", value = "True", if = ["lab"] },
18+
]
19+
20+
[envs.default]
21+
e2e-env = false

esxi/pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ dynamic = [
3838
[project.optional-dependencies]
3939
deps = [
4040
"pyvmomi==8.0.2.0.1; python_version > '3.0'",
41+
"pysocks==1.7.1"
4142
]
4243

4344
[project.urls]

esxi/tests/common.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
11
# (C) Datadog, Inc. 2024-present
22
# All rights reserved
33
# Licensed under a 3-clause BSD style license (see LICENSE)
4+
import os
5+
46
from pyVmomi import vim, vmodl
57

68
HOST = "127.0.0.1"
79
PORT = 8989
10+
USE_VSPHERE_LAB = os.environ.get('USE_VSPHERE_LAB')
11+
812
VCSIM_INSTANCE = {
913
'host': f"{HOST}:{str(PORT)}",
1014
'username': 'test',

esxi/tests/conftest.py

Lines changed: 31 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# (C) Datadog, Inc. 2024-present
22
# All rights reserved
33
# Licensed under a 3-clause BSD style license (see LICENSE)
4+
import copy
45
import os
56

67
import pytest
@@ -19,19 +20,42 @@
1920
PERF_METRIC_ID,
2021
PORT,
2122
PROPERTIES_EX,
23+
USE_VSPHERE_LAB,
2224
VCSIM_INSTANCE,
2325
)
26+
from .ssh_tunnel import socks_proxy
2427

2528

2629
@pytest.fixture(scope='session')
2730
def dd_environment():
28-
compose_file = os.path.join(get_here(), 'docker', 'docker-compose.yaml')
29-
conditions = [
30-
WaitForPortListening(HOST, PORT, wait=10),
31-
CheckDockerLogs(compose_file, 'export GOVC_URL', wait=10),
32-
]
33-
with docker_run(compose_file, conditions=conditions):
34-
yield VCSIM_INSTANCE
31+
if not USE_VSPHERE_LAB:
32+
compose_file = os.path.join(get_here(), 'docker', 'docker-compose.yaml')
33+
conditions = [
34+
WaitForPortListening(HOST, PORT, wait=10),
35+
CheckDockerLogs(compose_file, 'export GOVC_URL', wait=10),
36+
]
37+
with docker_run(compose_file, conditions=conditions):
38+
yield VCSIM_INSTANCE
39+
else:
40+
lab_public_ip = os.environ.get('VSPHERE_LAB_IP')
41+
lab_private_key = os.environ.get('VSPHERE_LAB_PRIVATE_KEY_PATH')
42+
lab_username = os.environ.get('VSPHERE_LAB_USERNAME')
43+
with socks_proxy(
44+
lab_public_ip,
45+
lab_username,
46+
lab_private_key,
47+
) as socks:
48+
49+
socks_ip, socks_port = socks
50+
print(f"socks: {socks}")
51+
instance = copy.deepcopy(VCSIM_INSTANCE)
52+
53+
instance['host'] = os.environ.get('ESXI_HOST_IP')
54+
instance['username'] = os.environ.get('ESXI_HOST_USERNAME')
55+
instance['password'] = os.environ.get('ESXI_HOST_PASSWORD')
56+
instance['ssl_verify'] = 'false'
57+
instance['proxy'] = 'socks5://{}:{}'.format(socks_ip, socks_port)
58+
yield instance
3559

3660

3761
@pytest.fixture

0 commit comments

Comments
 (0)