jborean93/smb_b_open.py

## smb_b_open.py
# Copyright: (c) 2019, Jordan Borean (@jborean93) <jborean93@gmail.com>
# MIT License (see LICENSE or https://opensource.org/licenses/MIT)

import uuid

from contextlib import contextmanager
from io import BytesIO
from smbprotocol.connection import Connection
from smbprotocol.session import Session
from smbprotocol.open import CreateDisposition, FileAttributes, FilePipePrinterAccessMask, ImpersonationLevel, Open, \
    ShareAccess
from smbprotocol.tree import TreeConnect


@contextmanager
def smb_b_open(path, mode='r', share='r', username=None, password=None, encrypt=True):
    """
    Functions similar to the builtin open() method where it will create an open handle to a file over SMB. This can be
    used to read and/or write data to the file using the methods exposed by the Open() class in smbprotocol. Read and
    write operations only support bytes and not text strings.

    :param path: The full UNC path to the file to open. This should be '\\server\share\path.txt'.
    :param mode: The mode in which the file is to be opened, can be set to one of the following;
        'r': Opens the file for reading (default)
        'w': Opens the file for writing, truncating first
        'x': Create a new file and open it for writing, fail if the file already exists
    :param share: The SMB sharing mode to set for the opened file handle, can be set to one or more of the following:
        'r': Allows other handles to read from the file (default)
        'w': Allows other handles to write to the file
        'd': Allows other handles to delete the file
    :param username: Optional username to use for authentication, required if Kerberos is not used.
    :param password: Optional password to use for authentication, required if Kerberos is not used.
    :param enrypt: Whether to use encryption or not, Must be set to False if using an older SMB Dialect.
    :return: The opened smbprotocol Open() obj that has a read, write, and flush functions.
    """
    path_split = [e for e in path.split('\\') if e]

    if len(path_split) < 3:
        raise ValueError("Path should specify at least the server, share, and file to read.")

    server = path_split[0]
    share_name = path_split[1]
    file_path = "\\".join(path_split[2:])

    conn = Connection(uuid.uuid4(), server)
    conn.connect()

    try:
        session = Session(conn, username=username, password=password, require_encryption=encrypt)
        session.connect()
        try:
            tree = TreeConnect(session, r"\\%s\%s" % (server, share_name))
            tree.connect()
            try:
                if mode == 'r':
                    create_disposition = CreateDisposition.FILE_OPEN
                    access_mask = FilePipePrinterAccessMask.GENERIC_READ
                elif mode == 'w':
                    create_disposition = CreateDisposition.FILE_OVERWRITE_IF
                    access_mask = FilePipePrinterAccessMask.GENERIC_WRITE
                elif mode == 'x':
                    create_disposition = CreateDisposition.FILE_CREATE
                    access_mask = FilePipePrinterAccessMask.GENERIC_WRITE
                else:
                    raise ValueError("Invalid mode value specified.")

                share_map = {
                    'r': ShareAccess.FILE_SHARE_READ,
                    'w': ShareAccess.FILE_SHARE_WRITE,
                    'd': ShareAccess.FILE_SHARE_DELETE,
                }
                share_access = 0
                for s in share:
                    share_access |= share_map[s]

                obj = Open(tree, file_path)
                obj.create(
                    ImpersonationLevel.Impersonation,
                    access_mask,
                    FileAttributes.FILE_ATTRIBUTE_NORMAL,
                    share_access,
                    create_disposition,
                    0,
                )

                try:
                    yield obj
                finally:
                    obj.close()
            finally:
                tree.disconnect()
        finally:
            session.disconnect()
    finally:
        conn.disconnect()


# Reads a file and store it in a BytesIO buffer.
b_io = BytesIO()
with smb_b_open(r'\\server.domain.local\c$\temp\file.xlsx') as file_obj:
    offset = 0
    while offset < file_obj.end_of_file:
        length = file_obj.connection.max_read_size
        data = file_obj.read(offset, length)
        b_io.write(data)
        offset += length

    b_io.seek(0)
	# Copyright: (c) 2019, Jordan Borean (@jborean93) <jborean93@gmail.com>
	# MIT License (see LICENSE or https://opensource.org/licenses/MIT)

	import uuid

	from contextlib import contextmanager
	from io import BytesIO
	from smbprotocol.connection import Connection
	from smbprotocol.session import Session
	from smbprotocol.open import CreateDisposition, FileAttributes, FilePipePrinterAccessMask, ImpersonationLevel, Open, \
	ShareAccess
	from smbprotocol.tree import TreeConnect


	@contextmanager
	def smb_b_open(path, mode='r', share='r', username=None, password=None, encrypt=True):
	"""
	Functions similar to the builtin open() method where it will create an open handle to a file over SMB. This can be
	used to read and/or write data to the file using the methods exposed by the Open() class in smbprotocol. Read and
	write operations only support bytes and not text strings.

	:param path: The full UNC path to the file to open. This should be '\\server\share\path.txt'.
	:param mode: The mode in which the file is to be opened, can be set to one of the following;
	'r': Opens the file for reading (default)
	'w': Opens the file for writing, truncating first
	'x': Create a new file and open it for writing, fail if the file already exists
	:param share: The SMB sharing mode to set for the opened file handle, can be set to one or more of the following:
	'r': Allows other handles to read from the file (default)
	'w': Allows other handles to write to the file
	'd': Allows other handles to delete the file
	:param username: Optional username to use for authentication, required if Kerberos is not used.
	:param password: Optional password to use for authentication, required if Kerberos is not used.
	:param enrypt: Whether to use encryption or not, Must be set to False if using an older SMB Dialect.
	:return: The opened smbprotocol Open() obj that has a read, write, and flush functions.
	"""
	path_split = [e for e in path.split('\\') if e]

	if len(path_split) < 3:
	raise ValueError("Path should specify at least the server, share, and file to read.")

	server = path_split[0]
	share_name = path_split[1]
	file_path = "\\".join(path_split[2:])

	conn = Connection(uuid.uuid4(), server)
	conn.connect()

	try:
	session = Session(conn, username=username, password=password, require_encryption=encrypt)
	session.connect()
	try:
	tree = TreeConnect(session, r"\\%s\%s" % (server, share_name))
	tree.connect()
	try:
	if mode == 'r':
	create_disposition = CreateDisposition.FILE_OPEN
	access_mask = FilePipePrinterAccessMask.GENERIC_READ
	elif mode == 'w':
	create_disposition = CreateDisposition.FILE_OVERWRITE_IF
	access_mask = FilePipePrinterAccessMask.GENERIC_WRITE
	elif mode == 'x':
	create_disposition = CreateDisposition.FILE_CREATE
	access_mask = FilePipePrinterAccessMask.GENERIC_WRITE
	else:
	raise ValueError("Invalid mode value specified.")

	share_map = {
	'r': ShareAccess.FILE_SHARE_READ,
	'w': ShareAccess.FILE_SHARE_WRITE,
	'd': ShareAccess.FILE_SHARE_DELETE,
	}
	share_access = 0
	for s in share:
	share_access \|= share_map[s]

	obj = Open(tree, file_path)
	obj.create(
	ImpersonationLevel.Impersonation,
	access_mask,
	FileAttributes.FILE_ATTRIBUTE_NORMAL,
	share_access,
	create_disposition,
	0,
	)

	try:
	yield obj
	finally:
	obj.close()
	finally:
	tree.disconnect()
	finally:
	session.disconnect()
	finally:
	conn.disconnect()


	# Reads a file and store it in a BytesIO buffer.
	b_io = BytesIO()
	with smb_b_open(r'\\server.domain.local\c$\temp\file.xlsx') as file_obj:
	offset = 0
	while offset < file_obj.end_of_file:
	length = file_obj.connection.max_read_size
	data = file_obj.read(offset, length)
	b_io.write(data)
	offset += length

	b_io.seek(0)