From 8ca54797c8c30be0ddc5f707d5eb532b1eedd653 Mon Sep 17 00:00:00 2001
From: Michael Thamm <mike.thamm@canonical.com>
Date: Wed, 9 Oct 2024 01:48:24 -0400
Subject: [PATCH] [juju] Add plugin option for Juju state reporting

Signed-off-by: Michael Thamm <mike.thamm@canonical.com>
---
 sos/report/plugins/juju.py | 124 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 122 insertions(+), 2 deletions(-)

diff --git a/sos/report/plugins/juju.py b/sos/report/plugins/juju.py
index 8a194df615..6f8a138b27 100644
--- a/sos/report/plugins/juju.py
+++ b/sos/report/plugins/juju.py
@@ -8,20 +8,74 @@
 #
 # See the LICENSE file in the source distribution for further information.
 
-from sos.report.plugins import Plugin, UbuntuPlugin
+import pwd
+import json
+from sos.report.plugins import Plugin, UbuntuPlugin, PluginOpt
 
 
 class Juju(Plugin, UbuntuPlugin):
+    """The Juju plugin is aimed at collecting Juju-related logs,
+    configurations, and controller/model state(s).
+
+    Logs and agent configuration information (/var/log/juju and /var/lib/juju)
+    is collected by default since these are useful for troubleshooting.
+
+    The Juju state collection is disabled by default and can be enabled with
+    the 'juju-state' option. Collecting Juju state is safe in theory, but it
+    does act on the live controller(s)/model(s) and is therefore optional.
+
+    The default Juju state collection collects all controllers and models that
+    the 'juju-user' (default=ubuntu) has access to.
+
+    Specific controllers or models can be collected using the 'controllers'
+    and 'models' options.
+
+    Important: the string list is whitespace delimited, not colon delimited
+    (sos plugin standard). This is due to the underlying Juju CLI accepting
+    specific models in the format 'controller:model' and whitespaces are not
+    allowed in either controller and model names.
+
+    Example: models="controller_a:model_x controller_b:model_y"
+    """
 
     short_desc = 'Juju orchestration tool'
 
     plugin_name = 'juju'
-    profiles = ('virt', 'sysmgmt')
+    profiles = ('virt', 'sysmgmt',)
 
     # Using files instead of packages here because there is no identifying
     # package on a juju machine.
     files = ('/var/log/juju',)
 
+    option_list = [
+        PluginOpt(
+            "juju-state",
+            default=False,
+            val_type=bool,
+            desc="Include Juju state in the report",
+        ),
+        PluginOpt(
+            "juju-user",
+            default="ubuntu",
+            val_type=str,
+            desc="Juju client user.",
+        ),
+        PluginOpt(
+            "controllers",
+            default="",
+            val_type=str,
+            desc="Collect Juju state for specified controllers. Uses a \
+            whitespace delimited list.",
+        ),
+        PluginOpt(
+            "models",
+            default="",
+            val_type=str,
+            desc="Collect Juju state for specified models. Uses a whitespace \
+            delimited list.",
+        ),
+    ]
+
     def setup(self):
         # Juju service names are not consistent through deployments,
         # so we need to use a wildcard to get the correct service names.
@@ -53,6 +107,71 @@ def setup(self):
             # logs in the directory.
             self.add_copy_spec("/var/log/juju/*.log")
 
+        # Only include the Juju state report if this plugin option is set
+        if not self.get_option("juju-state"):
+            return
+
+        juju_user = self.get_option("juju-user")
+        try:
+            pwd.getpwnam(juju_user)
+        except KeyError:
+            self._log_warn(
+                f'User "{juju_user}" does not exist, '
+                "will not collect Juju information."
+            )
+            return
+
+        if self.get_option("controllers") and self.get_option("models"):
+            self._log_warn(
+                "Options: controllers, models are mutually exclusive. "
+                "Will not collect Juju information."
+            )
+            return
+
+        controllers_json = self.collect_cmd_output(
+            "juju controllers --format=json", runas=juju_user
+        )
+        if controllers_json["status"] == 0:
+            desired_controllers = set(
+                self.get_option("controllers").split(" ")
+            )
+            # If a controller option is supplied, use it. Otherwise, get all
+            #  controllers
+            if desired_controllers and desired_controllers != {""}:
+                controllers = desired_controllers
+            else:
+                controllers = set(
+                    json.loads(controllers_json["output"])[
+                        "controllers"
+                    ].keys()
+                )
+        else:
+            controllers = {}
+
+        # Specific models
+        if self.get_option("models"):
+            for model in self.get_option("models").split(" "):
+                command = f"juju status -m {model} --format=json"
+                self.add_cmd_output(command, runas=juju_user)
+
+        # All controllers and all models OR specific controllers and all
+        #  models for each
+        else:
+            for controller in controllers:
+                models_json = self.exec_cmd(
+                    f"juju models --all -c {controller} --format=json",
+                    runas=juju_user,
+                )
+                if models_json["status"] == 0:
+                    models = json.loads(models_json["output"])["models"]
+                    for model in models:
+                        short_name = model["short-name"]
+                        command = (
+                            f"juju status -m {controller}:{short_name} "
+                            f"--format=json"
+                        )
+                        self.add_cmd_output(command, runas=juju_user)
+
     def postproc(self):
         agents_path = "/var/lib/juju/agents/*"
         protect_keys = [
@@ -68,5 +187,6 @@ def postproc(self):
         self.do_path_regex_sub(agents_path, keys_regex, sub_regex)
         # Redact certificates
         self.do_file_private_sub(agents_path)
+        self.do_cmd_private_sub('juju controllers')
 
 # vim: set et ts=4 sw=4 :