[pve-devel] [PATCH v1 pve-esxi-import-tools 4/5] listvms: add arg parser, context manager for connections, fetch helper

Max Carrara m.carrara at proxmox.com
Tue Mar 19 16:32:49 CET 2024


In order to make the CLI interface more friendly to humans, Python's
`argparse` [0] module from the standard library is used to parse the
arguments provided to the script. Each option and positional argument
also contain a short help text that is shown when running the script
with either "-h" or "--help".

Additionally, this commit also adds a context manager [1] for
establishing connections to an ESXi host. The context manager ensures
that the connection is closed in its inner `finally` block.

The inner part of the VM-data-fetching loop in `main()` is factored
out into a separate helper function, which now raises a `RuntimeError`
if the datacenter of a VM cannot be looked up.

In general, should any exception be thrown inside the loop, its output
is subsequently logged to stderr. The loop then just continues like
before.

Any exception that is not caught inside of `main()` is now printed to
stderr, followed by exiting with `1`.

Overall, the script's behaviour and output on successful operations
remains the same, except regarding unsuccessful argument parsing and
displaying error messages. In other words, invocations prior to this
patch should result in the same JSON output (if successful).

This was tested by piping the outputs of this script before and after
this commit through `jq` and then comparing the outputs with `diff`.

[0]: https://docs.python.org/3.11/library/argparse.html
[1]: https://docs.python.org/3/library/contextlib.html#contextlib.contextmanager

Signed-off-by: Max Carrara <m.carrara at proxmox.com>
---
 listvms.py | 189 +++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 140 insertions(+), 49 deletions(-)

diff --git a/listvms.py b/listvms.py
index cdea95a..9f85bf9 100755
--- a/listvms.py
+++ b/listvms.py
@@ -1,10 +1,13 @@
 #!/usr/bin/python3
 
+import argparse
 import dataclasses
 import json
 import ssl
 import sys
+import textwrap
 
+from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
@@ -13,6 +16,97 @@ from pyVim.connect import SmartConnect, Disconnect
 from pyVmomi import vim
 
 
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        prog="listvms",
+        description="List VMs on an ESXi host.",
+    )
+
+    def _squeeze_and_wrap(text: str) -> str:
+        """Makes it easier to write help text using multiline strings."""
+
+        text = text.replace("\n", " ").strip()
+
+        # squeeze recurring spaces
+        while "  " in text:
+            text = text.replace("  ", " ")
+
+        return "\n".join(textwrap.wrap(text, 60, break_on_hyphens=False))
+
+    parser.add_argument(
+        "--skip-cert-verification",
+        help=_squeeze_and_wrap(
+            """Skip the verification of TLS certs, e.g. to allow self-signed
+            certs."""
+        ),
+        action="store_true",
+    )
+
+    parser.add_argument(
+        "hostname",
+        help=_squeeze_and_wrap("""The name or address of the ESXi host."""),
+    )
+
+    parser.add_argument(
+        "username",
+        help=_squeeze_and_wrap("""The name of the user to connect with."""),
+    )
+
+    parser.add_argument(
+        "password_file",
+        help=_squeeze_and_wrap(
+            """The file which contains the password for the provided
+            username."""
+        ),
+        type=Path,
+    )
+
+    return parser.parse_args()
+
+
+ at dataclass
+class EsxiConnectonArgs:
+    hostname: str
+    username: str
+    password_file: Path
+    skip_cert_verification: bool = False
+
+
+ at contextmanager
+def connect_to_esxi_host(args: EsxiConnectonArgs) -> vim.ServiceInstance:
+    """Opens a connection to an ESXi host with the given username and password
+    contained in the password file.
+    """
+    ssl_context = (
+        ssl._create_unverified_context()
+        if args.skip_cert_verification
+        else None
+    )
+
+    with open(args.password_file) as pw_file:
+        password = pw_file.read().strip()
+
+    connection = None
+
+    try:
+        connection = SmartConnect(
+            host=args.hostname,
+            user=args.username,
+            pwd=password,
+            sslContext=ssl_context,
+        )
+
+        yield connection
+
+    except OSError as err:
+        print(f"Failed to connect: {err}", file=sys.stderr)
+        raise
+
+    finally:
+        if connection is not None:
+            Disconnect(connection)
+
+
 @dataclass
 class VmVmxInfo:
     datastore: str
@@ -112,62 +206,59 @@ def get_all_datacenters(service_instance: vim.ServiceInstance) -> list[vim.Datac
     dc_view.Destroy()
     return datacenters
 
+
+def fetch_and_update_vm_data(vm: vim.VirtualMachine, data: dict[Any, Any]):
+    """Fetches all required VM, datastore and datacenter information, and
+    then updates the given `dict`.
+
+    Raises:
+        RuntimeError: If looking up the datacenter for the given VM fails.
+    """
+    datacenter = get_datacenter_of_vm(vm)
+    if datacenter is None:
+        raise RuntimeError(f"Failed to lookup datacenter for VM {vm.name}")
+
+    data.setdefault(datacenter.name, {})
+
+    vms = data[datacenter.name].setdefault("vms", {})
+    datastores = data[datacenter.name].setdefault("datastores", {})
+
+    vms[vm.name] = VmInfo(
+        config=get_vm_vmx_info(vm),
+        disks=get_vm_disk_info(vm),
+        power=str(vm.runtime.powerState),
+    )
+
+    datastores.update({ds.name: ds.url for ds in vm.config.datastoreUrl})
+
+
 def main():
-    if sys.argv[1] == '--skip-cert-verification':
-        del sys.argv[1]
-        ssl_context = ssl._create_unverified_context()
-    else:
-        ssl_context = None
-
-    esxi_host = sys.argv[1]
-    esxi_user = sys.argv[2]
-    esxi_password_file = sys.argv[3]
-
-    esxi_password = ''
-    with open(esxi_password_file) as f:
-        esxi_password = f.read()
-        if esxi_password.endswith('\n'):
-            esxi_password = esxi_password[:-1]
+    args = parse_args()
 
-    try:
-        si = SmartConnect(
-            host=esxi_host,
-            user=esxi_user,
-            pwd=esxi_password,
-            sslContext=ssl_context,
-        )
-    except OSError as err:
-        print(f"failed to connect: {err}")
-        sys.exit(1)
+    connection_args = EsxiConnectonArgs(
+        hostname=args.hostname,
+        username=args.username,
+        password_file=args.password_file,
+        skip_cert_verification=args.skip_cert_verification,
+    )
 
-    try:
-        vms = list_vms(si)
+    with connect_to_esxi_host(connection_args) as connection:
         data = {}
-        for vm in vms:
-            name = 'vm ' + vm.name
+        for vm in list_vms(connection):
             try:
-                dc = get_datacenter_of_vm(vm)
-                if dc is None:
-                    print(
-                        f"Failed to get datacenter for {name}",
-                        file=sys.stderr
-                    )
-
-                vm_info = VmInfo(
-                    config=get_vm_vmx_info(vm),
-                    disks=get_vm_disk_info(vm),
-                    power=vm.runtime.powerState,
+                fetch_and_update_vm_data(vm, data)
+            except Exception as err:
+                print(
+                    f"Failed to get info for VM {vm.name}: {err}",
+                    file=sys.stderr,
                 )
 
-                datastore_info = {ds.name: ds.url for ds in vm.config.datastoreUrl}
-                data.setdefault(dc.name, {}).setdefault('vms', {})[vm.name] = vm_info
-                data.setdefault(dc.name, {}).setdefault('datastores', {}).update(datastore_info)
-            except Exception as err:
-                print("failed to get info for", name, ':', err, file=sys.stderr)
+    print(json.dumps(data, indent=2, default=json_dump_helper))
 
-        print(json.dumps(data, indent=2, default=json_dump_helper))
-    finally:
-        Disconnect(si)
 
 if __name__ == "__main__":
-    main()
+    try:
+        main()
+    except Exception as err:
+        print(f"Encountered unexpected error: {err}", file=sys.stderr)
+        sys.exit(1)
-- 
2.39.2





More information about the pve-devel mailing list