Add functions with arguments for constructing SSH command and config strings

Author: bennuttall

hostedpi/cli/options.py

@@ -23,7 +23,8 @@
     Option("--launchpad", "--lp", help="Launchpad usernames to source SSH keys from"),
 ]
 ipv6 = Annotated[bool, Option(help="Use the IPv6 connection method")]
-numeric = Annotated[bool, Option("-n", help="Use numeric IPv6 address in SSH command")]
+numeric = Annotated[bool, Option("--numeric", "-n", help="Use numeric IPv6 address in SSH command")]
+user = Annotated[str, Option("--username", "-u", help="Username for SSH connection")]
 yes = Annotated[bool, Option("--yes", "-y", help="Proceed without confirmation")]
 number = Annotated[Union[int, None], Option(help="Number of Raspberry Pi servers to create", min=1)]
 full_table = Annotated[bool, Option(help="Show full table of Raspberry Pi server info")]

hostedpi/cli/ssh.py

@@ -12,7 +12,10 @@
 
 @ssh_app.command("command")
 def do_command(
-    name: arguments.server_name, ipv6: options.ipv6 = False, numeric: options.numeric = False
+    name: arguments.server_name,
+    ipv6: options.ipv6 = False,
+    numeric: options.numeric = False,
+    user: options.user = "root",
 ):
     """
     Get the SSH command to connect to a Raspberry Pi server
@@ -23,14 +26,11 @@ def do_command(
         raise Exit(1)
     try:
         if ipv6:
-            if numeric:
-                print(pi.ipv6_ssh_command_numeric)
-            else:
-                print(pi.ipv6_ssh_command)
+            print(pi.get_ipv6_ssh_command(numeric=numeric, user=user))
         else:
-            print(pi.ipv4_ssh_command)
+            print(pi.get_ipv4_ssh_command(user=user))
     except HostedPiException as exc:
-        print(f"hostedpi error: {exc}")
+        print_error(f"hostedpi error: {exc}")
         raise Exit(1)
 
 
@@ -40,20 +40,21 @@ def do_config(
     filter: options.filter_pattern_pi = None,
     ipv6: options.ipv6 = False,
     numeric: options.numeric = False,
+    user: options.user = "root",
 ):
     """
     Get the SSH config to connect to one or more Raspberry Pi servers
     """
+    if numeric and not ipv6:
+        print_error("--numeric is only supported with --ipv6")
+
     pis = get_pis(names, filter)
     for pi in pis:
         try:
             if ipv6:
-                if numeric:
-                    print(pi.ipv6_ssh_config_numeric)
-                else:
-                    print(pi.ipv6_ssh_config)
+                print(pi.get_ipv6_ssh_config(numeric=numeric, user=user))
             else:
-                print(pi.ipv4_ssh_config)
+                print(pi.get_ipv4_ssh_config(user=user))
         except HostedPiException as exc:
-            print(f"hostedpi error: {exc}")
+            print_error(f"hostedpi error: {exc}")
             raise Exit(1)

hostedpi/pi.py

@@ -274,58 +274,32 @@ def ipv6_ssh_hostname(self) -> str:
     @property
     def ipv4_ssh_command(self) -> str:
         """
-        The SSH command required to connect to the Pi using SSH over IPv4
+        The default SSH command required to connect to the Pi using SSH over IPv4
         """
-        return f"ssh -p {self.ipv4_ssh_port} root@{self.ipv4_ssh_hostname}"
+        return self.get_ipv4_ssh_command()
 
     @property
     def ipv6_ssh_command(self) -> str:
         """
-        The SSH command required to connect to the Pi using SSH over IPv6
+        The default SSH command required to connect to the Pi using SSH over IPv6
         """
-        return f"ssh root@{self.ipv6_ssh_hostname}"
-
-    @property
-    def ipv6_ssh_command_numeric(self) -> str:
-        """
-        The SSH command required to connect to the Pi using SSH over IPv6 (using the IPv6 address
-        rather than hostname)
-        """
-        return f"ssh root@[{self.ipv6_address.compressed}]"
+        return self.get_ipv6_ssh_command()
 
     @property
     def ipv4_ssh_config(self) -> str:
         """
-        A string containing the IPv4 SSH config for the Pi. The contents could be added to an SSH
-        config file for easy access to the Pi.
+        A string containing the default IPv4 SSH config for the Pi. The contents could be added to
+        an SSH config file for easy access to the Pi.
         """
-        return f"""Host {self.name}
-    user root
-    port {self.ipv4_ssh_port}
-    hostname {self.ipv4_ssh_hostname}
-        """.strip()
+        return self.get_ipv4_ssh_config()
 
     @property
     def ipv6_ssh_config(self) -> str:
         """
-        A string containing the IPv6 SSH config for the Pi. The contents could be added to an SSH
-        config file for easy access to the Pi.
+        A string containing the default IPv6 SSH config for the Pi. The contents could be added to
+        an SSH config file for easy access to the Pi.
         """
-        return f"""Host {self.name}
-    user root
-    hostname {self.ipv6_ssh_hostname}
-        """.strip()
-
-    @property
-    def ipv6_ssh_config_numeric(self) -> str:
-        """
-        A string containing the IPv6 SSH config for the Pi. The contents could be added to an SSH
-        config file for easy access to the Pi.
-        """
-        return f"""Host {self.name}
-    user root
-    hostname {self.ipv6_address.compressed}
-        """.strip()
+        return self.get_ipv6_ssh_config()
 
     @property
     def url(self) -> str:
@@ -523,6 +497,48 @@ def cancel(self):
 
         self._cancelled = True
 
+    def get_ipv4_ssh_command(self, user: str = "root") -> str:
+        """
+        Construct an SSH command required to connect to the Pi using SSH over IPv4
+        """
+        return f"ssh -p {self.ipv4_ssh_port} {user}@{self.ipv4_ssh_hostname}"
+
+    def get_ipv6_ssh_command(self, *, user: str = "root", numeric: bool = False) -> str:
+        """
+        Construct an SSH command required to connect to the Pi using SSH over IPv6
+        """
+        if numeric:
+            return f"ssh {user}@[{self.ipv6_address.compressed}]"
+        else:
+            return f"ssh {user}@{self.ipv6_ssh_hostname}"
+
+    def get_ipv4_ssh_config(self, user: str = "root") -> str:
+        """
+        Construct a string containing the IPv4 SSH config for the Pi. The contents could be added to
+        an SSH config file for easy access to the Pi.
+        """
+        return f"""Host {self.name}
+    user {user}
+    port {self.ipv4_ssh_port}
+    hostname {self.ipv4_ssh_hostname}
+        """.strip()
+
+    def get_ipv6_ssh_config(
+        self,
+        *,
+        user: str = "root",
+        numeric: bool = False,
+    ) -> str:
+        """
+        Construct a string containing the SSH config for the Pi. The contents could be added to an
+        SSH config file for easy access to the Pi.
+        """
+        hostname = self.ipv6_address.compressed if numeric else self.ipv6_ssh_hostname
+        return f"""Host {self.name}
+    user {user}
+    hostname {hostname}
+        """.strip()
+
     def add_ssh_keys(self, ssh_keys: SSHKeySources) -> set[str]:
         """
         Add SSH keys to the Pi from the specified sources.

hostedpi/picloud.py

@@ -96,18 +96,34 @@ def pis(self) -> dict[str, Pi]:
     @property
     def ipv4_ssh_config(self) -> str:
         """
-        A string containing the IPv4 SSH config for all Pis within the account. The contents could
-        be added to an SSH config file for easy access to the Pis in the account.
+        A string containing the default IPv4 SSH config for all Pis within the account. The contents
+        could be added to an SSH config file for easy access to the Pis in the account.
         """
-        return "\n".join(pi.ipv4_ssh_config for pi in self.pis.values())
+        return self.get_ipv4_ssh_config()
 
     @property
     def ipv6_ssh_config(self) -> str:
         """
-        A string containing the IPv6 SSH config for all Pis within the account. The contents could
-        be added to an SSH config file for easy access to the Pis in the account.
+        A string containing the default IPv6 SSH config for all Pis within the account. The contents
+        could be added to an SSH config file for easy access to the Pis in the account.
         """
-        return "\n".join(pi.ipv6_ssh_config for pi in self.pis.values())
+        return self.get_ipv6_ssh_config()
+
+    def get_ipv4_ssh_config(self, user: str = "root") -> str:
+        """
+        Construct a string containing the IPv4 SSH config for all Pis within the account. The
+        contents could be added to an SSH config file for easy access to the Pis in the account.
+        """
+        return "\n".join(pi.get_ipv4_ssh_config(user=user) for pi in self.pis.values())
+
+    def get_ipv6_ssh_config(self, user: str = "root", numeric: bool = False) -> str:
+        """
+        Construct a string containing the IPv6 SSH config for all Pis within the account. The
+        contents could be added to an SSH config file for easy access to the Pis in the account.
+        """
+        return "\n".join(
+            pi.get_ipv6_ssh_config(user=user, numeric=numeric) for pi in self.pis.values()
+        )
 
     def create_pi(
         self,

tests/test_pi.py

@@ -707,3 +707,56 @@ def test_get_pi_info_error_500(pi_name, pi_info_basic, auth, error_500):
     with pytest.raises(HostedPiServerError):
         pi.info
     assert auth._api_session.get.call_count == 1
+
+
+def test_pi_get_ipv4_ssh_command(pi_name, pi_info_basic, auth, pi_info_response):
+    pi = Pi(name=pi_name, info=pi_info_basic, auth=auth)
+    auth._api_session.get.return_value = pi_info_response
+    assert pi.get_ipv4_ssh_command() == "ssh -p 5100 [email protected]"
+    assert pi.get_ipv4_ssh_command(user="pi") == "ssh -p 5100 [email protected]"
+
+
+def test_pi_get_ipv6_ssh_command(pi_name, pi_info_basic, auth, pi_info_response):
+    pi = Pi(name=pi_name, info=pi_info_basic, auth=auth)
+    auth._api_session.get.return_value = pi_info_response
+    assert pi.get_ipv6_ssh_command() == "ssh [email protected]"
+    assert pi.get_ipv6_ssh_command(user="pi") == "ssh [email protected]"
+    assert pi.get_ipv6_ssh_command(numeric=True) == "ssh root@[2a00:1098:8:64::1]"
+    assert pi.get_ipv6_ssh_command(user="pi", numeric=True) == "ssh pi@[2a00:1098:8:64::1]"
+
+
+def test_pi_get_ipv4_ssh_config(pi_name, pi_info_basic, auth, pi_info_response):
+    pi = Pi(name=pi_name, info=pi_info_basic, auth=auth)
+    auth._api_session.get.return_value = pi_info_response
+
+    config = pi.get_ipv4_ssh_config()
+    assert "Host test-pi" in config
+    assert "user root" in config
+    assert "port 5100" in config
+    assert "hostname ssh.test-pi.hostedpi.com" in config
+
+    config = pi.get_ipv4_ssh_config(user="pi")
+    assert "Host test-pi" in config
+    assert "user pi" in config
+    assert "port 5100" in config
+    assert "hostname ssh.test-pi.hostedpi.com" in config
+
+
+def test_pi_get_ipv6_ssh_config(pi_name, pi_info_basic, auth, pi_info_response):
+    pi = Pi(name=pi_name, info=pi_info_basic, auth=auth)
+    auth._api_session.get.return_value = pi_info_response
+
+    config = pi.get_ipv6_ssh_config()
+    assert "Host test-pi" in config
+    assert "user root" in config
+    assert "hostname test-pi.hostedpi.com" in config
+
+    config = pi.get_ipv6_ssh_config(user="pi")
+    assert "Host test-pi" in config
+    assert "user pi" in config
+    assert "hostname test-pi.hostedpi.com" in config
+
+    config = pi.get_ipv6_ssh_config(numeric=True)
+    assert "Host test-pi" in config
+    assert "user root" in config
+    assert "hostname 2a00:1098:8:64::1" in config