Merge pull request #5 from infrasonar/unifi-svc

Unifi Device Service support
infrasonar · Mar 14, 2024 · 18ef9ce · 18ef9ce
2 parents 989e809 + 871051e
commit 18ef9ce
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 28 deletions.
diff --git a/README.md b/README.md
@@ -83,7 +83,7 @@ infrasonar get-assets 123 -o yaml
 
 ## VMware guests
 
-Generate YAML (or JSON) with VMware Guests which are found on ESX or vCenter but wherefore no asset with the `vmwareguest` collector is found. This YAML can then be used to install the VMware Guest collector with a single command.
+Generate YAML (or JSON) with VMware Guests which are found on ESX or vCenter but wherefore no asset with the `vmwareguest` collector is found. This YAML can then be used to install the VMware Guest collector with a single command. The default is to use the VCenter. For ESX the `-c esx` argument.
 
 Example:  _(in the example below, 123 is a container Id)_
 
@@ -119,7 +119,7 @@ infrasonar apply-assets missing.yaml -a -d
 
 ## UniFi devices
 
-Generate YAML (or JSON) with UniFi devices which are found on a UniFi Controller asset(s) but wherefore no asset with the `unifidevice` collector is found. This YAML can then be used to install the UniFi Device collector with a single command.
+Generate YAML (or JSON) with UniFi devices which are found on a UniFi Controller asset(s) but wherefore no asset with the `unifidevice` or `unifidevicesvc` collector is found. This YAML can then be used to install the UniFi Device collector with a single command. The default is to use the `unificontroller`. For using the _service_ variant, use the `-c unificontrollersvc` argument.
 
 Example:  _(in the example below, 123 is a container Id)_
 

diff --git a/bin/infrasonar b/bin/infrasonar
@@ -3,8 +3,7 @@ import argparse
 import asyncio
 import getpass
 import json
-import os
-import pprint
+
 import re
 import sys
 import yaml
@@ -15,7 +14,7 @@ from dataclasses import dataclass
 from setproctitle import setproctitle
 from typing import Any, Optional, Dict, List, Tuple
 
-__version__ = '0.1.14'  # Update version in setup as well
+__version__ = '0.1.15'  # Update version in setup as well
 
 
 _labels_example = """
@@ -1313,14 +1312,9 @@ def hyperv_guests(container_id: int, api: str, verify_ssl: bool, output: str,
         print(_apply_warning)
 
 
-def unifi_devices(container_id: int, api: str, verify_ssl: bool, output: str,
-                  token: Optional[str], include: Optional[List[int]]):
-    if token is None:
-        try:
-            token = getpass.getpass('Enter user or container token: ')
-        except KeyboardInterrupt:
-            sys.exit('Cancelled')
-
+def _unifi_devices(container_id: int, api: str, verify_ssl: bool, token: str,
+                   include: Optional[List[int]], ucontroller: str,
+                   udevice: str):
     try:
         assets = asyncio.run(aget_assets(
             api=api,
@@ -1332,7 +1326,7 @@ def unifi_devices(container_id: int, api: str, verify_ssl: bool, output: str,
             not_kind=None,
             mode='normal',
             not_mode=None,
-            collector='unificontroller',
+            collector=ucontroller,
             not_collector=None,
             label=None,
             not_label=None))
@@ -1343,7 +1337,7 @@ def unifi_devices(container_id: int, api: str, verify_ssl: bool, output: str,
         sys.exit(msg)
 
     if not assets:
-        sys.exit('No assets with the unificontroller collector found')
+        return None, 0
 
     asyncio.run(asyncio.sleep(0.5))
 
@@ -1353,7 +1347,7 @@ def unifi_devices(container_id: int, api: str, verify_ssl: bool, output: str,
             continue
         config = next((
             x for x in asset['collectors']
-            if x['key'] == 'unificontroller'
+            if x['key'] == ucontroller
         ))['config'] or None
 
         address = config.get('address', '') or asset['name']
@@ -1362,7 +1356,7 @@ def unifi_devices(container_id: int, api: str, verify_ssl: bool, output: str,
         vssl = config.get('ssl', False)
 
         data = asyncio.run(async_get_check_data(api, verify_ssl, asset['id'],
-                                                token, 'unificontroller',
+                                                token, ucontroller,
                                                 'device'))
         for device in data['data']['device']:
             device['_parent'] = (address, port, site, vssl)
@@ -1381,7 +1375,7 @@ def unifi_devices(container_id: int, api: str, verify_ssl: bool, output: str,
             not_kind=None,
             mode=None,
             not_mode=None,
-            collector='unifidevice',
+            collector=udevice,
             not_collector=None,
             label=None,
             not_label=None))
@@ -1405,7 +1399,7 @@ def unifi_devices(container_id: int, api: str, verify_ssl: bool, output: str,
             mode=None,
             not_mode=None,
             collector=None,
-            not_collector='unifidevice',
+            not_collector=udevice,
             label=None,
             not_label=None))
     except KeyboardInterrupt:
@@ -1421,13 +1415,11 @@ def unifi_devices(container_id: int, api: str, verify_ssl: bool, output: str,
     for asset in installed:
         mac = next((
             x for x in asset['collectors']
-            if x['key'] == 'unifidevice'))['config']['mac'].upper()
+            if x['key'] == udevice))['config']['mac'].upper()
         device_map.pop(mac, None)
 
     if not device_map:
-        print('No missing UniFi devices are found '
-              f'(installed devices: {total})', file=sys.stderr)
-        sys.exit(0)
+        return [], total
 
     assets = []
     for device in device_map.values():
@@ -1439,7 +1431,7 @@ def unifi_devices(container_id: int, api: str, verify_ssl: bool, output: str,
 
         asset['name'] = name
         asset['collectors'] = [{
-            'key': 'unifidevice',
+            'key': udevice,
             'config': {
                 'controller': device['_parent'][0],
                 'port': device['_parent'][1],
@@ -1449,6 +1441,34 @@ def unifi_devices(container_id: int, api: str, verify_ssl: bool, output: str,
             },
         }]
         assets.append(asset)
+    return assets, total
+
+
+def unifi_devices(container_id: int, api: str, verify_ssl: bool, output: str,
+                  token: Optional[str], ucontroller: str,
+                  include: Optional[List[int]]):
+    if token is None:
+        try:
+            token = getpass.getpass('Enter user or container token: ')
+        except KeyboardInterrupt:
+            sys.exit('Cancelled')
+
+    assets, total = _unifi_devices(
+        container_id=container_id, api=api, verify_ssl=verify_ssl,
+        token=token, include=include, ucontroller=ucontroller,
+        udevice=(
+            'unifidevice' if ucontroller == 'unificontroller'
+            else 'unifidevicesvc')
+        )
+
+    if assets is None:
+        sys.exit('No assets with the unificontroller or unificontrollersvc '
+                 'collector found')
+
+    if not assets:
+        print('No missing UniFi devices are found '
+              f'(installed devices: {total})', file=sys.stderr)
+        sys.exit(0)
 
     data = OrderedDict()
     data['container'] = container_id
@@ -1474,7 +1494,7 @@ if __name__ == '__main__':
     parser.add_argument(
         '--api',
         default='https://api.infrasonar.com',
-        help='URL for the API')
+        help='URL for the API (default: https://api.infrasonar.com)')
     parser.add_argument(
         '--skip-verify-ssl',
         action='store_true',
@@ -1652,6 +1672,14 @@ if __name__ == '__main__':
                                         'remains empty, all UniFi Controller '
                                         'assets will be used'
                                       ))
+    action_unifi_devices.add_argument('-c', '--collector',
+                                      choices=[
+                                        'unificontroller',
+                                        'unificontrollersvc'],
+                                      default='unificontroller',
+                                      help=(
+                                        'Use either vCenter or ESX for '
+                                        'getting the guests list'))
     action_unifi_devices.add_argument('--token',
                                       default=None,
                                       help='Token for authentication')
@@ -1723,6 +1751,7 @@ if __name__ == '__main__':
             not args.skip_verify_ssl,
             args.output,
             args.token,
+            args.collector,
             args.include or None,
         )
     else:

diff --git a/setup.py b/setup.py
@@ -1,8 +1,7 @@
 """
-locan installation: pip install -e .
-
 Upload to PyPI
 
+pip -v install --use-pep517 -e .
 python setup.py sdist
 twine upload --repository pypitest dist/infrasonar-X.X.X.tar.gz
 twine upload --repository pypi dist/infrasonar-X.X.X.tar.gz
@@ -17,7 +16,7 @@
 
 setup(
     name='infrasonar',
-    version='0.1.14',  # Update version in infrasonar as well
+    version='0.1.15',  # Update version in infrasonar as well
     description='InfraSonar Toolkit',
     url='https://github.com/infrasonar/toolkit',
     long_description=long_description,
@@ -46,6 +45,7 @@
         'Programming Language :: Python :: 3.9',
         'Programming Language :: Python :: 3.10',
     ],
+    packages=find_packages(),
     install_requires=[
         'aiohttp',
         'pyyaml',