Server Message Block (SMB) v2

import uuid

import demistomock as demisto  # noqa: F401
from CommonServerPython import *  # noqa: F401

""" IMPORTS """

import smbclient
from smbclient import (
    mkdir,
    open_file,
    register_session,
    remove,
    rmdir,
    scandir,
)


def get_file_name(path):
    delimiter = "/" if "/" in path else "\\"
    return path.split(delimiter)[-1]


def handle_path(path):
    """
    Stripping the '\\' and '/' characters from a given path.
    Examples:
        \\Shared\\123.txt\\ ---> Shared\\123.txt
        /Shared/123.txt/ ---> Shared/123.txt
        \\Shared\\123.txt/\\ ---> Shared\\123.txt
    """
    return path.strip("\\/")


def create_share_path(hostname, path):
    r"""
    Create a path to the shared folder according to the smbprotocol convention: '\\server\share'.
    For reference see https://github.com/jborean93/smbprotocol/blob/master/examples/high-level/directory-management.py
    """
    return rf"\\{hostname}\{path}"


class SMBClient:
    def __init__(self, hostname, user, password, encrypt, port):
        self.hostname = hostname
        self._user = user
        self._password = password
        self._encrypt = encrypt
        self._port = port

    def create_session(
        self, hostname: str = None, user: str = None, password: str = None, encrypt: bool = False, port: int = None
    ):
        register_session(
            server=hostname or self.hostname,
            username=user or self._user,
            password=password or self._password,
            port=port or self._port,
            encrypt=encrypt or self._encrypt,
            auth_protocol="ntlm",
        )


def test_module(client: SMBClient):
    client.create_session()
    return "ok"


def smb_upload(client: SMBClient, args: dict):
    hostname = args.get("hostname")
    path = handle_path(args.get("file_path"))
    path = create_share_path(hostname or client.hostname, path)
    username = args.get("username")
    password = args.get("password")
    entryID = args.get("entryID")
    content = args.get("content")

    if not entryID and not content:
        raise DemistoException("You must provide a content to upload using one of the following arguments: content, entryID.")

    client.create_session(hostname, username, password)

    # For the content argument - the input is text.
    writing_mode = "w"
    if entryID:
        file = demisto.getFilePath(entryID)
        filePath = file["path"]
        writing_mode = "wb"

        with open(filePath, mode="rb") as f:
            content = f.read()

    with open_file(rf"{path}", mode=writing_mode) as file_obj:
        file_obj.write(content)
    return f"File {get_file_name(path)} was uploaded successfully"


def smb_download(client: SMBClient, args: dict):
    hostname = args.get("hostname")
    path = handle_path(args.get("file_path"))
    path = create_share_path(hostname or client.hostname, path)
    username = args.get("username")
    password = args.get("password")

    client.create_session(hostname, username, password)

    with open_file(rf"{path}", mode="rb") as fd:
        file_contents = fd.read()
        file_name = get_file_name(path)
        return fileResult(file_name, file_contents)


def smb_remove_file(client: SMBClient, args: dict):
    hostname = args.get("hostname")
    path = handle_path(args.get("file_path"))
    path = create_share_path(hostname or client.hostname, path)
    username = args.get("username")
    password = args.get("password")

    client.create_session(hostname, username, password)
    remove(path)
    file_name = get_file_name(path)
    return f"File {file_name} was deleted successfully"


def list_dir(client: SMBClient, args: dict):
    hostname = args.get("hostname")
    username = args.get("username")
    password = args.get("password")
    path = handle_path(args.get("path"))
    path = create_share_path(hostname or client.hostname, path)

    client.create_session(hostname, username, password)
    entries = list(scandir(path))

    files = []
    dirs = []

    for entry in entries:
        if entry.is_file():
            files.append(entry.name)
        if entry.is_dir():
            dirs.append(entry.name)

    context = {
        "SharedFolder": path,
        "Files": files,
        "Directories": dirs,
    }
    return CommandResults(
        outputs_prefix="SMB.Path",
        outputs_key_field="SharedFolder",
        outputs=context,
        readable_output=tableToMarkdown(f"List Of Entries for {path}", context),
    )


def smb_mkdir(client: SMBClient, args: dict):
    hostname = args.get("hostname")
    username = args.get("username")
    password = args.get("password")
    path = handle_path(args.get("path"))
    path = create_share_path(hostname or client.hostname, path)

    client.create_session(hostname, username, password)

    mkdir(path)

    return f"Directory: {path} was created successfully"


def smb_rmdir(client: SMBClient, args: dict):
    hostname = args.get("hostname")
    username = args.get("username")
    password = args.get("password")
    path = handle_path(args.get("path"))
    path = create_share_path(hostname or client.hostname, path)

    client.create_session(hostname, username, password)
    rmdir(path)

    return f"Directory: {path} was removed successfully"


def main():
    params = demisto.params()
    hostname = params["hostname"]
    port = int(params.get("port", "445"))
    user = params["credentials"]["identifier"]
    password = params["credentials"]["password"]
    encrypt = params.get("encrypt", False)
    dc = params.get("dc", None)
    verify = params.get("require_secure_negotiate", True)
    client_guid = params.get("client_guid", None)

    # Temporary workaround to an issue in the smbprotocol package.
    # Git issue: https://github.com/jborean93/smbprotocol/issues/109
    config = smbclient.ClientConfig(username=user, password=password, require_secure_negotiate=verify)
    config.domain_controller = dc

    if client_guid:
        try:
            client_guid = uuid.UUID(client_guid)
            config.client_guid = client_guid
        except ValueError:
            demisto.info(f"Failed to convert {client_guid} to a valid UUID string. Using a random generated UUID instead")

    client = SMBClient(hostname=hostname, user=user, password=password, encrypt=encrypt, port=port)

    demisto.info(f"Command being called is {demisto.command()}")

    try:
        if demisto.command() == "test-module":
            return_results(test_module(client))
        elif demisto.command() == "smb-download":
            return_results(smb_download(client, demisto.args()))
        elif demisto.command() == "smb-upload":
            return_results(smb_upload(client, demisto.args()))
        elif demisto.command() == "smb-directory-list":
            return_results(list_dir(client, demisto.args()))
        elif demisto.command() == "smb-file-remove":
            return_results(smb_remove_file(client, demisto.args()))
        elif demisto.command() == "smb-directory-create":
            return_results(smb_mkdir(client, demisto.args()))
        elif demisto.command() == "smb-directory-remove":
            return_results(smb_rmdir(client, demisto.args()))
    except Exception as e:
        return_error(f"Failed to execute {demisto.command()} command. Error: {e!s}")
    finally:
        smbclient.reset_connection_cache()


if __name__ in ("__main__", "__builtin__", "builtins"):
    main()