Esempio n. 1
0
    def fetch(
        self,
        src: str,
        dest: str,
        configuration_name: str = DEFAULT_CONFIGURATION_NAME,
        expand_variables: bool = False,
    ) -> None:
        """
        Will fetch a single file from the remote Windows host and create a
        local copy. Like copy(), this can be slow when it comes to fetching
        large files due to the limitation of WinRM.

        This method will first store the file in a temporary location before
        creating or replacing the file at dest if the checksum is correct.

        :param src: The path to the file on the remote host to fetch
        :param dest: The path on the localhost host to store the file as
        :param configuration_name: The PowerShell configuration endpoint to
            use when fetching the file.
        :param expand_variables: Expand variables in path. Disabled by default
            Enable for cmd like expansion (for example %TMP% in path)
        """
        if expand_variables:
            dest = os.path.expanduser(os.path.expandvars(dest))
        log.info("Fetching '%s' to '%s'" % (src, dest))

        with RunspacePool(self.wsman,
                          configuration_name=configuration_name) as pool:
            script = get_pwsh_script("fetch.ps1")
            powershell = PowerShell(pool)
            powershell.add_script(script).add_argument(src).add_argument(
                expand_variables)

            log.debug("Starting remote process to output file data")
            powershell.invoke()
            _handle_powershell_error(powershell,
                                     "Failed to fetch file %s" % src)
            log.debug("Finished remote process to output file data")

            expected_hash = powershell.output[-1]

            temp_file, path = tempfile.mkstemp()
            try:
                file_bytes = base64.b64decode(powershell.output[0])
                os.write(temp_file, file_bytes)

                sha1 = hashlib.sha1()
                sha1.update(file_bytes)
                actual_hash = sha1.hexdigest()

                log.debug("Remote Hash: %s, Local Hash: %s" %
                          (expected_hash, actual_hash))
                if actual_hash != expected_hash:
                    raise WinRMError("Failed to fetch file %s, hash mismatch\n"
                                     "Source: %s\nFetched: %s" %
                                     (src, expected_hash, actual_hash))
                shutil.copy(path, dest)
            finally:
                os.close(temp_file)
                os.remove(path)
Esempio n. 2
0
    def copy(self,
             src,
             dest,
             configuration_name=DEFAULT_CONFIGURATION_NAME,
             expand_variables=False):
        """
        Copies a single file from the current host to the remote Windows host.
        This can be quite slow when it comes to large files due to the
        limitations of WinRM but it is designed to be as fast as it can be.
        During the copy process, the bytes will be stored in a temporary file
        before being copied.

        When copying it will replace the file at dest if one already exists. It
        also will verify the checksum of the copied file is the same as the
        actual file locally before copying the file to the path at dest.

        :param src: The path to the local file
        :param dest: The path to the destination file on the Windows host
        :param configuration_name: The PowerShell configuration endpoint to
            use when copying the file.
        :param expand_variables: Expand variables in path. Disabled by default
            Enable for cmd like expansion (for example %TMP% in path)
        :return: The absolute path of the file on the Windows host
        """
        def read_buffer(b_path, total_size, buffer_size):
            offset = 0
            sha1 = hashlib.sha1()

            with open(b_path, 'rb') as src_file:
                for data in iter((lambda: src_file.read(buffer_size)), b""):
                    log.debug(
                        "Reading data of file at offset=%d with size=%d" %
                        (offset, buffer_size))
                    offset += len(data)
                    sha1.update(data)
                    b64_data = base64.b64encode(data)

                    result = [to_unicode(b64_data)]
                    if offset == total_size:
                        result.append(
                            to_unicode(
                                base64.b64encode(to_bytes(sha1.hexdigest()))))

                    yield result

                # the file was empty, return empty buffer
                if offset == 0:
                    yield [
                        u"",
                        to_unicode(base64.b64encode(to_bytes(
                            sha1.hexdigest())))
                    ]

        if expand_variables:
            src = os.path.expanduser(os.path.expandvars(src))
        b_src = to_bytes(src)
        src_size = os.path.getsize(b_src)
        log.info("Copying '%s' to '%s' with a total size of %d" %
                 (src, dest, src_size))

        with RunspacePool(self.wsman,
                          configuration_name=configuration_name) as pool:
            # Get the buffer size of each fragment to send, subtract. Adjust to size of the base64 encoded bytes. Also
            # subtract 82 for the fragment, message, and other header info that PSRP adds.
            buffer_size = int((self.wsman.max_payload_size - 82) / 4 * 3)

            log.info("Creating file reader with a buffer size of %d" %
                     buffer_size)
            read_gen = read_buffer(b_src, src_size, buffer_size)

            command = get_pwsh_script('copy.ps1')
            log.debug("Starting to send file data to remote process")
            powershell = PowerShell(pool)
            powershell.add_script(command) \
                .add_argument(dest) \
                .add_argument(expand_variables)
            powershell.invoke(input=read_gen)
            _handle_powershell_error(powershell, "Failed to copy file")

        log.debug("Finished sending file data to remote process")
        for warning in powershell.streams.warning:
            warnings.warn(str(warning))

        output_file = to_unicode(powershell.output[-1]).strip()
        log.info("Completed file transfer of '%s' to '%s'" %
                 (src, output_file))
        return output_file