Adding file_download and file_upload

pyntc/devices/ios_device.py

@@ -685,6 +685,109 @@ def file_copy_remote_exists(self, src, dest=None, file_system=None):
         log.debug("Host %s: File %s does not already exist on remote.", self.host, src)
         return False
 
+    def file_download(self, url, dest="", md5=None, file_system=None, read_timeout=2000):
+        """Download file from remote server.
+
+        Args:
+            url (str): URL of the file to download. (e.g., "scp://user:[email protected]/image.bin")
+            dest (str, optional): Destination name for file. Defaults to using the same name as the file on the remote server.
+            md5 (str, optional): Expected MD5 hash of the file to download. If provided, the file will be verified.
+            file_system (str, optional): File system to copy file to.
+            read_timeout (int, optional): Netmiko timeout when waiting for device prompt. Default 2000.
+
+        Raises:
+            SocketClosedError: Error raised if connection to device is closed.
+            FileTransferError: Error in transferring file.
+            FileTransferError: Error if unable to verify file was transferred successfully.
+        """
+        self.enable()
+        if file_system is None:
+            file_system = self._get_file_system()
+
+        filename = dest or os.path.basename(url)
+        try:
+            file_md5 = self.file_md5(filename, file_system)
+        except FileTransferError:
+            # File doesn't exist, download it
+            pass
+        else:
+            # File already exists
+            if not md5:
+                log.info("Host %s: File %s already exists on remote.", self.host, filename)
+                return True
+            if file_md5 != md5:
+                log.warning(
+                    "Host %s: File %s already exists on remote but MD5 hash mismatch. Downloading again.",
+                    self.host,
+                    filename,
+                )
+
+        copy_command = f"copy {url} {file_system}{dest}"
+
+        self.show(copy_command, read_timeout=read_timeout)
+
+        if not md5:
+            return True
+        return self.verify_md5(filename, md5, file_system)
+
+    def file_upload(self, filename, url, file_system=None, read_timeout=2000):
+        """Upload file to remote server.
+
+        Examples:
+            # Upload file to remote server
+            >>> device.file_upload("image.bin", "scp://user:[email protected]/image.bin")
+            # Upload running config to remote server (override default file system)
+            >>> device.file_upload("running-config", "scp://user:[email protected]/config.txt", file_system="")
+
+        Args:
+            filename (str): Name of the file to upload.
+            url (str): URL of the file to upload to. (e.g., "scp://user:[email protected]/image.bin")
+            file_system (str, optional): File system the file is on.
+            read_timeout (int, optional): Netmiko timeout when waiting for device prompt. Default 2000.
+        """
+        self.enable()
+        if file_system is None:
+            file_system = self._get_file_system()
+
+        copy_command = f"copy {file_system}{filename} {url}"
+        self.show(copy_command, read_timeout=read_timeout)
+
+    def file_md5(self, filename, file_system=None, read_timeout=1000):
+        """Return the MD5 hash of a file on the device.
+
+        Args:
+            filename (str): Name of the file to get the MD5 hash of.
+            file_system (str, optional): File system the file is on. Defaults to the default file system.
+            read_timeout (int, optional): Netmiko timeout when waiting for device prompt. Default 1000.
+
+        Returns:
+            str: MD5 hash of the file.
+        """
+        self.enable()
+        if file_system is None:
+            file_system = self._get_file_system()
+
+        command = f"verify /md5 {file_system}{filename}"
+        response = self.show(command, read_timeout=read_timeout)
+        md5_pattern = r"=\s+([a-f0-9]{32})"
+        match = re.search(md5_pattern, response)
+        if match:
+            return match.group(1)
+        raise FileTransferError(f"Unable to verify MD5 hash for file {filename}")
+
+    def verify_md5(self, filename, md5, file_system=None):
+        """Verify the MD5 hash of a file on the device.
+
+        Args:
+            filename (str): Name of the file to verify.
+            md5 (str): MD5 hash of the file.
+            file_system (str, optional): File system the file is on. Defaults to the default file system.
+
+        Returns:
+            bool: True if the MD5 hash matches, False otherwise.
+        """
+        return self.file_md5(filename, file_system) == md5
+
     def install_os(self, image_name, install_mode=False, read_timeout=2000, **vendor_specifics):
         """Installs the prescribed Network OS, which must be present before issuing this command.