Skip to content

Commit 4f79191

Browse files
committed
Fix lint
Signed-off-by: Tushar <tgupta3@users.noreply.github.com>
1 parent f5e5f1b commit 4f79191

File tree

1 file changed

+45
-30
lines changed

1 file changed

+45
-30
lines changed

suzieq/poller/controller/source/vcenter.py

Lines changed: 45 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,8 @@
88

99
import asyncio
1010
import logging
11-
from requests.auth import HTTPBasicAuth
12-
from requests.exceptions import RequestException
13-
from typing import Any, Dict, List, Optional, Tuple, Union
14-
from urllib.parse import urljoin, urlparse
11+
from typing import Dict, List, Optional, Union
12+
from urllib.parse import urlparse
1513
from pyVim.connect import SmartConnect, Disconnect
1614
from pyVmomi import vim, vmodl
1715
import ssl
@@ -84,6 +82,7 @@ def validate_password(cls, password):
8482
except SensitiveLoadError as e:
8583
raise ValueError(e)
8684

85+
8786
class Vcenter(Source, InventoryAsyncPlugin):
8887
def __init__(self, config_data: dict, validate: bool = True) -> None:
8988
self._status = 'init'
@@ -109,14 +108,17 @@ def _load(self, input_data):
109108
)
110109
self._server = self._data.server
111110
if not self._auth:
112-
raise InventorySourceError(f"{self.name} Vcenter must have an "
113-
"'auth' set in the 'namespaces' section"
114-
)
111+
raise InventorySourceError(
112+
f"{self.name} Vcenter must have an "
113+
"'auth' set in the 'namespaces' section")
115114

116115
def _init_session(self):
117116
"""Initialize the session property"""
118117
context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
119-
context.verify_mode = ssl.CERT_NONE if not self._data.ssl_verify else ssl.CERT_REQUIRED
118+
context.verify_mode = ssl.CERT_REQUIRED
119+
if not self._data.ssl_verify:
120+
context.verify_mode = ssl.CERT_NONE
121+
120122
try:
121123
self._session = SmartConnect(
122124
host=self._server.host,
@@ -127,47 +129,56 @@ def _init_session(self):
127129
)
128130
except Exception as e:
129131
self._session = None
130-
raise InventorySourceError(f"Failed to connect to VCenter: {str(e)}")
132+
raise InventorySourceError(
133+
f"Failed to connect to VCenter: {str(e)}")
131134

132135
def _get_custom_keys(self, content, attribute_names):
133136
"""Retrieve custom attribute keys based on their names."""
134-
all_custom_fields = {field.name: field.key for field in content.customFieldsManager.field}
135-
return [all_custom_fields[name] for name in attribute_names if name in all_custom_fields]
137+
all_custom_fields = {field.name: field.key
138+
for field in content.customFieldsManager.field}
139+
return [
140+
all_custom_fields[name]
141+
for name in attribute_names
142+
if name in all_custom_fields
143+
]
136144

137145
def _create_filter_spec(self, view):
138-
"""Create and return a FilterSpec based on provided view and attribute keys."""
146+
"""Return a FilterSpec based on provided view and attribute keys."""
139147
traversal_spec = vmodl.query.PropertyCollector.TraversalSpec(
140-
name='traverseEntities',
141-
path='view',
142-
skip=False,
148+
name='traverseEntities', path='view', skip=False,
143149
type=vim.view.ContainerView,
144-
selectSet=[vmodl.query.PropertyCollector.SelectionSpec(name='traverseEntities')]
145-
)
146-
prop_set = vmodl.query.PropertyCollector.PropertySpec(all=False, type=vim.VirtualMachine)
150+
selectSet=[vmodl.query.PropertyCollector.SelectionSpec(
151+
name='traverseEntities')])
152+
prop_set = vmodl.query.PropertyCollector.PropertySpec(
153+
all=False, type=vim.VirtualMachine)
147154
prop_set.pathSet = ['name', 'guest.ipAddress', 'customValue']
148-
obj_spec = vmodl.query.PropertyCollector.ObjectSpec(obj=view, selectSet=[traversal_spec])
155+
obj_spec = vmodl.query.PropertyCollector.ObjectSpec(
156+
obj=view, selectSet=[traversal_spec])
149157
filter_spec = vmodl.query.PropertyCollector.FilterSpec()
150158
filter_spec.objectSet = [obj_spec]
151159
filter_spec.propSet = [prop_set]
152160
return filter_spec
153161

154162
async def get_inventory_list(self) -> List:
155163
"""
156-
Retrieve VMs that have any of a list of specified custom attribute names using the Property Collector.
164+
Retrieve VMs that have any specified custom attribute names.
157165
158-
This method uses vSphere's Property Collector to fetch only properties that are required.
159-
This is a lot faster than fetching the entire inventory and filtering on attributes.
166+
This method uses vSphere's Property Collector to fetch only
167+
properties that are required. This is a lot faster than
168+
fetching the entire inventory and filtering on attributes.
160169
"""
161170
if not self._session:
162171
self._init_session()
163172

164173
content = self._session.RetrieveContent()
165-
view = content.viewManager.CreateContainerView(content.rootFolder, [vim.VirtualMachine], True)
174+
view = content.viewManager.CreateContainerView(
175+
content.rootFolder, [vim.VirtualMachine], True)
166176
attribute_keys = self._get_custom_keys(content, self._data.attributes)
167177

168178
filter_spec = self._create_filter_spec(view)
169179
retrieve_options = vmodl.query.PropertyCollector.RetrieveOptions()
170-
result = content.propertyCollector.RetrievePropertiesEx([filter_spec], retrieve_options)
180+
result = content.propertyCollector.RetrievePropertiesEx(
181+
[filter_spec], retrieve_options)
171182
vms_with_ip = {}
172183
while result:
173184
for obj in result.objects:
@@ -180,20 +191,23 @@ async def get_inventory_list(self) -> List:
180191
elif prop.name == 'guest.ipAddress' and prop.val:
181192
vm_ip = prop.val
182193
elif prop.name == 'customValue':
183-
has_custom_attr = any(cv.key in attribute_keys for cv in prop.val)
194+
has_custom_attr = any(
195+
cv.key in attribute_keys for cv in prop.val)
184196
if has_custom_attr and vm_ip:
185197
vms_with_ip[vm_name] = vm_ip
186198

187199
if hasattr(result, 'token') and result.token:
188-
result = content.propertyCollector.ContinueRetrievePropertiesEx(token=result.token)
200+
property_collector = content.propertyCollector
201+
result = property_collector.ContinueRetrievePropertiesEx(
202+
token=result.token)
189203
else:
190204
break
191205

192206
view.Destroy()
193-
logger.info(f'Vcenter: Retrieved {len(vms_with_ip)} VMs with IPs that have any of the specified attribute names')
207+
logger.info(
208+
f'Vcenter: Retrieved {len(vms_with_ip)} VMs with IPs')
194209
return vms_with_ip
195210

196-
197211
def parse_inventory(self, inventory_list: list) -> Dict:
198212
inventory = {}
199213
for name, ip in inventory_list.items():
@@ -203,7 +217,8 @@ def parse_inventory(self, inventory_list: list) -> Dict:
203217
'namespace': namespace,
204218
'hostname': name,
205219
}
206-
logger.info(f'Vcenter: Acting on inventory of {len(inventory)} devices')
220+
logger.info(
221+
f'Vcenter: Acting on inventory of {len(inventory)} devices')
207222
return inventory
208223

209224
async def _execute(self):
@@ -217,4 +232,4 @@ async def _execute(self):
217232

218233
async def _stop(self):
219234
if self._session:
220-
Disconnect(self._session)
235+
Disconnect(self._session)

0 commit comments

Comments
 (0)