fix: machine scp/ssh operations

juju/machine.py

@@ -3,14 +3,17 @@
 
 import ipaddress
 import logging
+import typing
 
 import pyrfc3339
 
-from . import model, tag, jasyncio
+from juju.utils import block_until, juju_ssh_key_paths
+
+from . import jasyncio, model, tag
 from .annotationhelper import _get_annotations, _set_annotations
 from .client import client
 from .errors import JujuError
-from juju.utils import juju_ssh_key_paths
+from .status import InstanceStatusT, MachineAgentStatusT
 
 log = logging.getLogger(__name__)
 
@@ -37,14 +40,14 @@ async def destroy(self, force=False):
         return await self.model._wait('machine', self.id, 'remove')
     remove = destroy
 
-    async def get_annotations(self):
+    async def get_annotations(self) -> typing.Dict[str, str]:
         """Get annotations on this machine.
 
         :return dict: The annotations for this application
         """
         return await _get_annotations(self.tag, self.connection)
 
-    async def set_annotations(self, annotations):
+    async def set_annotations(self, annotations: typing.Dict[str, str]):
         """Set annotations on this machine.
 
         :param annotations map[string]string: the annotations as key/value
@@ -53,7 +56,7 @@ async def set_annotations(self, annotations):
         """
         return await _set_annotations(self.tag, annotations, self.connection)
 
-    def _format_addr(self, addr):
+    def _format_addr(self, addr: str):
         """Validate and format IP address.
 
         :param addr: IPv6 or IPv4 address
@@ -69,8 +72,9 @@ def _format_addr(self, addr):
             fmt = '{}'
         return fmt.format(ipaddr)
 
-    async def scp_to(self, source, destination, user='ubuntu', proxy=False,
-                     scp_opts=''):
+    async def scp_to(self, source: str, destination: str, user: str = 'ubuntu', proxy: bool = False,
+                     scp_opts: typing.Union[str, typing.List[str]] = '',
+                     wait_for_active: bool = True, timeout: typing.Optional[int] = None):
         """Transfer files to this machine.
 
         :param str source: Local path of file(s) to transfer
@@ -79,10 +83,13 @@ async def scp_to(self, source, destination, user='ubuntu', proxy=False,
         :param bool proxy: Proxy through the Juju API server
         :param scp_opts: Additional options to the `scp` command
         :type scp_opts: str or list
+        :param bool wait_for_active: Wait until the machine is ready to take in ssh commands.
+        :param int timeout: Time in seconds to wait until the machine becomes ready.
         """
         if proxy:
             raise NotImplementedError('proxy option is not implemented')
-
+        if wait_for_active:
+            await block_until(lambda: self.addresses, timeout=timeout)
         try:
             # if dns_name is an IP address format it appropriately
             address = self._format_addr(self.dns_name)
@@ -92,8 +99,9 @@ async def scp_to(self, source, destination, user='ubuntu', proxy=False,
         destination = '{}@{}:{}'.format(user, address, destination)
         await self._scp(source, destination, scp_opts)
 
-    async def scp_from(self, source, destination, user='ubuntu', proxy=False,
-                       scp_opts=''):
+    async def scp_from(self, source: str, destination: str, user: str = 'ubuntu',
+                       proxy: bool = False, scp_opts: typing.Union[str, typing.List[str]] = '',
+                       wait_for_active: bool = True, timeout: typing.Optional[int] = None):
         """Transfer files from this machine.
 
         :param str source: Remote path of file(s) to transfer
@@ -102,10 +110,13 @@ async def scp_from(self, source, destination, user='ubuntu', proxy=False,
         :param bool proxy: Proxy through the Juju API server
         :param scp_opts: Additional options to the `scp` command
         :type scp_opts: str or list
+        :param bool wait_for_active: Wait until the machine is ready to take in ssh commands.
+        :param int timeout: Time in seconds to wait until the machine becomes ready.
         """
         if proxy:
             raise NotImplementedError('proxy option is not implemented')
-
+        if wait_for_active:
+            await block_until(lambda: self.addresses, timeout=timeout)
         try:
             # if dns_name is an IP address format it appropriately
             address = self._format_addr(self.dns_name)
@@ -115,7 +126,7 @@ async def scp_from(self, source, destination, user='ubuntu', proxy=False,
         source = '{}@{}:{}'.format(user, address, source)
         await self._scp(source, destination, scp_opts)
 
-    async def _scp(self, source, destination, scp_opts):
+    async def _scp(self, source: str, destination: str, scp_opts: typing.Union[str, typing.List[str]]):
         """ Execute an scp command. Requires a fully qualified source and
         destination.
         """
@@ -135,18 +146,23 @@ async def _scp(self, source, destination, scp_opts):
             raise JujuError("command failed: %s" % cmd)
 
     async def ssh(
-            self, command, user='ubuntu', proxy=False, ssh_opts=None):
+            self, command: str, user: str = 'ubuntu', proxy: bool = False,
+            ssh_opts: typing.Optional[typing.Union[str, typing.List[str]]] = None,
+            wait_for_active: bool = True, timeout: typing.Optional[int] = None):
         """Execute a command over SSH on this machine.
 
         :param str command: Command to execute
         :param str user: Remote username
         :param bool proxy: Proxy through the Juju API server
         :param str ssh_opts: Additional options to the `ssh` command
-
+        :param bool wait_for_active: Wait until the machine is ready to take in ssh commands.
+        :param int timeout: Time in seconds to wait until the machine becomes ready.
         """
         if proxy:
             raise NotImplementedError('proxy option is not implemented')
         address = self.dns_name
+        if wait_for_active:
+            await block_until(lambda: self.addresses, timeout=timeout)
         destination = "{}@{}".format(user, address)
         _, id_path = juju_ssh_key_paths()
         cmd = [
@@ -167,8 +183,17 @@ async def ssh(
         # stdout is a bytes-like object, returning a string might be more useful
         return stdout.decode()
 
+
+    @property
+    def addresses(self) -> typing.List[str]:
+        """Returns the machine addresses.
+
+        """
+        return self.safe_data['addresses'] or []
+
+
     @property
-    def agent_status(self):
+    def agent_status(self) -> MachineAgentStatusT:
         """Returns the current Juju agent status string.
 
         """
@@ -182,7 +207,7 @@ def agent_status_since(self):
         return pyrfc3339.parse(self.safe_data['agent-status']['since'])
 
     @property
-    def agent_version(self):
+    def agent_version(self) -> str:
         """Get the version of the Juju machine agent.
 
         May return None if the agent is not yet available.
@@ -194,7 +219,7 @@ def agent_version(self):
             return None
 
     @property
-    def status(self):
+    def status(self) -> InstanceStatusT:
         """Returns the current machine provisioning status string.
 
         """
@@ -215,17 +240,16 @@ def status_since(self):
         return pyrfc3339.parse(self.safe_data['instance-status']['since'])
 
     @property
-    def dns_name(self):
+    def dns_name(self) -> typing.Optional[str]:
         """Get the DNS name for this machine. This is a best guess based on the
         addresses available in current data.
 
         May return None if no suitable address is found.
         """
-        addresses = self.safe_data['addresses'] or []
         ordered_addresses = []
         ordered_scopes = ['public', 'local-cloud', 'local-fan']
         for scope in ordered_scopes:
-            for address in addresses:
+            for address in self.addresses:
                 if scope == address['scope']:
                     ordered_addresses.append(address)
         for address in ordered_addresses:
@@ -236,7 +260,7 @@ def dns_name(self):
         return None
 
     @property
-    def hostname(self):
+    def hostname(self) -> typing.Optional[str]:
         """Get the hostname for this machine as reported by the machine agent
         running on it. This is only supported on 2.8.10+ controllers.