def node_exec_command(node_id, command, username, hostname, port, ssh_key=None, password=None, container_name=None, timeout=None): try: client = connect(hostname=hostname, port=port, username=username, password=password, pkey=ssh_key, timeout=timeout) except AztkError as e: return NodeOutput(node_id, None, e) if container_name: cmd = "sudo docker exec 2>&1 -t {0} /bin/bash -c 'set -e; set -o pipefail; {1}; wait'".format( container_name, command) else: cmd = "/bin/bash 2>&1 -c 'set -e; set -o pipefail; {0}; wait'".format( command) _, stdout, _ = client.exec_command(cmd, get_pty=True) output = stdout.read().decode("utf-8") client.close() return NodeOutput(node_id, output, None)
def node_copy( node_id, source_path, destination_path, username, hostname, port, ssh_key=None, password=None, container_name=None, timeout=None, ): try: client = connect(hostname=hostname, port=port, username=username, password=password, pkey=ssh_key, timeout=timeout) except AztkError as e: return NodeOutput(node_id, None, e) sftp_client = client.open_sftp() try: if container_name: # put the file in /tmp on the host tmp_file = "/tmp/" + os.path.basename(source_path) sftp_client.put(source_path, tmp_file) # move to correct destination on container docker_command = "sudo docker cp {0} {1}:{2}".format( tmp_file, container_name, destination_path) _, stdout, _ = client.exec_command(docker_command, get_pty=True) output = stdout.read().decode("utf-8") # clean up sftp_client.remove(tmp_file) return NodeOutput(node_id, output, None) else: output = sftp_client.put(source_path, destination_path).__str__() return NodeOutput(node_id, output, None) except (IOError, PermissionError) as e: return NodeOutput(node_id, None, e) finally: sftp_client.close() client.close()
def copy_from_node( node_id, source_path, destination_path, username, hostname, port, ssh_key=None, password=None, container_name=None, timeout=None, ): try: client = connect(hostname=hostname, port=port, username=username, password=password, pkey=ssh_key, timeout=timeout) except AztkError as e: return NodeOutput(node_id, False, e) sftp_client = client.open_sftp() try: if destination_path: destination_path = os.path.join(os.path.dirname(destination_path), node_id, os.path.basename(destination_path)) os.makedirs(os.path.dirname(destination_path), exist_ok=True) with open(destination_path, "wb") as f: sftp_client.getfo(source_path, f) return NodeOutput(node_id, f, None) else: import tempfile # create 2mb temporary file f = tempfile.SpooledTemporaryFile(2 * 1024**3) sftp_client.getfo(source_path, f) return NodeOutput(node_id, f, None) except OSError as e: return NodeOutput(node_id, None, e) finally: sftp_client.close() client.close()