Merge pull request #265 from hemanthnakkina/register-remote-controller

Add command to register private controller

sunbeam-python/sunbeam/commands/juju.py

@@ -785,8 +785,8 @@ def __init__(
         self.registration_token = None
         self.juju_account = None
 
-        home = os.environ.get("SNAP_REAL_HOME")
-        os.environ["JUJU_DATA"] = f"{home}/.local/share/juju"
+        self.home = os.environ.get("SNAP_REAL_HOME")
+        os.environ["JUJU_DATA"] = f"{self.home}/.local/share/juju"
 
     def is_skip(self, status: Optional["Status"] = None) -> Result:
         """Determines if the step should be skipped or not.
@@ -800,6 +800,7 @@ def is_skip(self, status: Optional["Status"] = None) -> Result:
         except JujuAccountNotFound as e:
             LOG.warning(e)
             return Result(ResultType.FAILED, "Account was not registered locally")
+
         try:
             user = self._juju_cmd("show-user")
             LOG.debug(f"Found user: {user['user-name']}")
@@ -861,9 +862,10 @@ def run(self, status: Optional["Status"] = None) -> Result:
         # client need to login/logout?
         # Does saving the password in $HOME/.local/share/juju/accounts.yaml
         # avoids login/logout?
-        register_args = ["register", self.registration_token]
+        register_args = ["--debug", "register", self.registration_token]
         if self.replace:
             register_args.append("--replace")
+        LOG.debug(f"User registration args: {register_args}")
 
         try:
             child = pexpect.spawn(
@@ -885,7 +887,15 @@ def run(self, status: Optional["Status"] = None) -> Result:
                     if index in (0, 1, 3):
                         child.sendline(self.juju_account.password)
                     elif index == 2:
-                        child.sendline(self.controller)
+                        result = child.before.decode()
+                        # If controller already exists, the command keeps on asking
+                        # controller name, so change the controller name to dummy.
+                        # The command errors out at the next stage that controller
+                        # is already registered.
+                        if f'Controller "{self.controller}" already exists' in result:
+                            child.sendline("dummy")
+                        else:
+                            child.sendline(self.controller)
                     elif index == 4:
                         result = child.before.decode()
                         if "ERROR" in result:
@@ -902,6 +912,136 @@ def run(self, status: Optional["Status"] = None) -> Result:
         return Result(ResultType.COMPLETED)
 
 
+class RegisterRemoteJujuUserStep(RegisterJujuUserStep):
+    """Register remote user/controller in juju."""
+
+    def __init__(
+        self,
+        client: Client,
+        token: str,
+        controller: str,
+        data_location: Path,
+        replace: bool = False,
+    ):
+        # User name not required to register a user. Pass empty string to
+        # base class as user name
+        super().__init__(client, "", controller, data_location, replace)
+        self.registration_token = token
+        self.account_file = f"{self.controller}.yaml"
+
+    def is_skip(self, status: Optional["Status"] = None) -> Result:
+        """Determines if the step should be skipped or not.
+
+        :return: ResultType.SKIPPED if the Step should be skipped,
+                 ResultType.COMPLETED or ResultType.FAILED otherwise
+        """
+        try:
+            self.juju_account = JujuAccount.load(self.data_location, self.account_file)
+            LOG.debug(f"Local account found for {self.controller}")
+        except JujuAccountNotFound:
+            password = pwgen.pwgen(12)
+            self.juju_account = JujuAccount(user="REPLACE_USER", password=password)
+            LOG.debug(f"Writing to file {self.juju_account}")
+            self.juju_account.write(self.data_location, self.account_file)
+
+        return Result(ResultType.COMPLETED)
+
+    def _get_user_from_local_juju(self, controller: str) -> str | None:
+        """Get user name from local juju accounts file."""
+        try:
+            with open(f"{self.home}/.local/share/juju/accounts.yaml") as f:
+                accounts = yaml.safe_load(f)
+                user = (
+                    accounts.get("controllers", {}).get(self.controller, {}).get("user")
+                )
+                LOG.debug(f"Found user from accounts.yaml for {controller}: {user}")
+        except FileNotFoundError as e:
+            LOG.debug(f"Error in retrieving local user: {str(e)}")
+            user = None
+
+        return user
+
+    def run(self, status: Optional["Status"] = None) -> Result:
+        """Run the step to completion.
+
+        Invoked when the step is run and returns a ResultType to indicate
+
+        :return:
+        """
+        result = super().run()
+        if result.result_type != ResultType.COMPLETED:
+            # Delete the account file created in skip step
+            account_file = self.data_location / self.account_file
+            account_file.unlink(missing_ok=True)
+            return result
+
+        # Update user name from local juju accounts.yaml
+        user = self._get_user_from_local_juju(self.controller)
+        if not user:
+            return Result(ResultType.FAILED, "User not updated in local juju client")
+
+        if self.juju_account.user != user:
+            self.juju_account.user = user
+            LOG.debug(f"Updating user in {self.juju_account} file")
+            self.juju_account.write(self.data_location, self.account_file)
+
+        return Result(ResultType.COMPLETED)
+
+
+class UnregisterJujuController(BaseStep, JujuStepHelper):
+    """Unregister an external Juju controller."""
+
+    def __init__(self, controller: str, data_location: Path):
+        super().__init__(
+            "Unregister Juju controller", f"Unregistering juju controller {controller}"
+        )
+        self.controller = controller
+        self.account_file = data_location / f"{self.controller}.yaml"
+
+    def is_skip(self, status: Optional["Status"] = None) -> Result:
+        """Determines if the step should be skipped or not.
+
+        :return: ResultType.SKIPPED if the Step should be skipped,
+                 ResultType.COMPLETED or ResultType.FAILED otherwise
+        """
+        try:
+            self.get_controller(self.controller)
+        except ControllerNotFoundException:
+            self.account_file.unlink(missing_ok=True)
+            LOG.warning(
+                f"Controller {self.controller} not found, skipping unregsiter "
+                "controller"
+            )
+            return Result(ResultType.SKIPPED)
+
+        return Result(ResultType.COMPLETED)
+
+    def run(self, status: Optional["Status"] = None) -> Result:
+        """Run the step to completion.
+
+        Invoked when the step is run and returns a ResultType to indicate
+
+        :return:
+        """
+        try:
+            cmd = [
+                self._get_juju_binary(),
+                "unregister",
+                self.controller,
+                "--no-prompt",
+            ]
+            LOG.debug(f'Running command {" ".join(cmd)}')
+            process = subprocess.run(cmd, capture_output=True, text=True, check=True)
+            LOG.debug(
+                f"Command finished. stdout={process.stdout}, stderr={process.stderr}"
+            )
+            self.account_file.unlink(missing_ok=True)
+        except subprocess.CalledProcessError as e:
+            return Result(ResultType.FAILED, str(e))
+
+        return Result(ResultType.COMPLETED)
+
+
 class AddJujuMachineStep(BaseStep, JujuStepHelper):
     """Add machine in juju."""
 
@@ -1747,3 +1887,31 @@ def run(self, status: Status | None) -> Result:
                     ResultType.FAILED, f"Machine ID not found for node {node['name']}"
                 )
         return Result(ResultType.COMPLETED)
+
+
+class SwitchToController(BaseStep, JujuStepHelper):
+    """Switch to controller in juju."""
+
+    def __init__(
+        self,
+        controller: str,
+    ):
+        super().__init__(
+            "Switch to juju controller", f"Switching to juju controller {controller}"
+        )
+        self.controller = controller
+
+    def run(self, status: Optional["Status"] = None) -> Result:
+        """Switch to juju controller."""
+        try:
+            cmd = [self._get_juju_binary(), "switch", self.controller]
+            LOG.debug(f'Running command {" ".join(cmd)}')
+            process = subprocess.run(cmd, capture_output=True, text=True, check=True)
+            LOG.debug(
+                f"Command finished. stdout={process.stdout}, stderr={process.stderr}"
+            )
+        except subprocess.CalledProcessError as e:
+            LOG.exception(f"Error in switching the controller to {self.controller}")
+            return Result(ResultType.FAILED, str(e))
+
+        return Result(ResultType.COMPLETED)

sunbeam-python/sunbeam/commands/juju_utils.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2024 Canonical Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+import click
+from rich.console import Console
+from snaphelpers import Snap
+
+from sunbeam.commands.juju import RegisterRemoteJujuUserStep, UnregisterJujuController
+from sunbeam.jobs.common import run_plan
+from sunbeam.jobs.deployment import Deployment
+
+LOG = logging.getLogger(__name__)
+console = Console()
+
+
+@click.command()
+@click.option(
+    "-f",
+    "--force",
+    is_flag=True,
+    help="Force replacement if controller already exists with the same name",
+)
+@click.argument("name", type=str)
+@click.argument("token", type=str)
+@click.pass_context
+def register_controller(ctx: click.Context, name: str, token: str, force: bool) -> None:
+    """Register existing Juju controller."""
+    deployment: Deployment = ctx.obj
+    client = deployment.get_client()
+    data_location = Snap().paths.user_data
+
+    plan = [
+        RegisterRemoteJujuUserStep(client, token, name, data_location, replace=force)
+    ]
+
+    run_plan(plan, console)
+    console.print(f"Controller {name} registered")
+
+
+@click.command()
+@click.argument("name", type=str)
+@click.pass_context
+def unregister_controller(ctx: click.Context, name: str) -> None:
+    """Unregister external Juju controller."""
+    data_location = Snap().paths.user_data
+
+    plan = [UnregisterJujuController(name, data_location)]
+
+    run_plan(plan, console)
+    console.print(f"Controller {name} unregistered")

sunbeam-python/sunbeam/jobs/juju.py

@@ -176,9 +176,11 @@ def to_dict(self):
         return self.model_dump()
 
     @classmethod
-    def load(cls, data_location: Path) -> "JujuAccount":
+    def load(
+        cls, data_location: Path, account_file: str = ACCOUNT_FILE
+    ) -> "JujuAccount":
         """Load account from file."""
-        data_file = data_location / ACCOUNT_FILE
+        data_file = data_location / account_file
         try:
             with data_file.open() as file:
                 return JujuAccount(**yaml.safe_load(file))
@@ -188,9 +190,9 @@ def load(cls, data_location: Path) -> "JujuAccount":
                 f"cluster yet? {data_file}"
             ) from e
 
-    def write(self, data_location: Path):
+    def write(self, data_location: Path, account_file: str = ACCOUNT_FILE):
         """Dump self to file."""
-        data_file = data_location / ACCOUNT_FILE
+        data_file = data_location / account_file
         if not data_file.exists():
             data_file.touch()
         data_file.chmod(0o660)

sunbeam-python/sunbeam/main.py

@@ -25,6 +25,7 @@
 from sunbeam.commands import dashboard_url as dasboard_url_cmds
 from sunbeam.commands import generate_cloud_config as generate_cloud_config_cmds
 from sunbeam.commands import inspect as inspect_cmds
+from sunbeam.commands import juju_utils as juju_cmds
 from sunbeam.commands import launch as launch_cmds
 from sunbeam.commands import manifest as manifest_cmds
 from sunbeam.commands import openrc as openrc_cmds
@@ -92,6 +93,12 @@ def utils(ctx):
     """Utilities for debugging and managing sunbeam."""
 
 
+@click.group("juju", context_settings=CONTEXT_SETTINGS, cls=CatchGroup)
+@click.pass_context
+def juju(ctx):
+    """Utilities for managing juju."""
+
+
 def main():
     snap = Snap()
     logfile = log.prepare_logfile(snap.paths.user_common / "logs", "sunbeam")
@@ -130,6 +137,10 @@ def main():
     cli.add_command(utils)
     utils.add_command(utils_cmds.juju_login)
 
+    cli.add_command(juju)
+    juju.add_command(juju_cmds.register_controller)
+    juju.add_command(juju_cmds.unregister_controller)
+
     # Register the features after all groups,commands are registered
     FeatureManager.register(deployment, cli)
 

sunbeam-python/tests/unit/sunbeam/commands/test_juju.py

@@ -515,3 +515,37 @@ def test_run_when_merge_bindings_fails(
         step._bindings = {"endpoint1": "test-space", "endpoint2": "test-space"}
         result = step.run()
         assert result.result_type == ResultType.FAILED
+
+
+class TestUnregisterJujuControllerStep:
+    def test_is_skip(self, mocker, tmp_path):
+        step = juju.UnregisterJujuController("testcontroller", tmp_path)
+        mocker.patch.object(step, "get_controller", return_value={"testcontroller"})
+        result = step.is_skip()
+        assert result.result_type == ResultType.COMPLETED
+
+    def test_is_skip_controller_not_registered(self, mocker, tmp_path):
+        step = juju.UnregisterJujuController("testcontroller", tmp_path)
+        mocker.patch.object(
+            step,
+            "get_controller",
+            side_effect=juju.ControllerNotFoundException("Controller not found"),
+        )
+        result = step.is_skip()
+        assert result.result_type == ResultType.SKIPPED
+
+    @patch("subprocess.run")
+    def test_run(self, mock_run, mocker, tmp_path):
+        step = juju.UnregisterJujuController("testcontroller", tmp_path)
+        mocker.patch.object(step, "_get_juju_binary", return_value="/juju-mock")
+        result = step.run()
+        assert mock_run.call_count == 1
+        assert result.result_type == ResultType.COMPLETED
+
+    @patch("subprocess.run", side_effect=subprocess.CalledProcessError(1, "command"))
+    def test_run_unregister_failed(self, mock_run, mocker, tmp_path):
+        step = juju.UnregisterJujuController("testcontroller", tmp_path)
+        mocker.patch.object(step, "_get_juju_binary", return_value="/juju-mock")
+        result = step.run()
+        assert mock_run.call_count == 1
+        assert result.result_type == ResultType.FAILED