import string
import logging
import pathlib
import multiprocessing
from shutil import move
from traceback import format_exc
from datetime import datetime

from man_spider.lib.smb import *
from man_spider.lib.file import *
from man_spider.lib.util import *
from man_spider.lib.errors import *
from man_spider.lib.processpool import *


log = logging.getLogger("manspider.spiderling")


class SpiderlingMessage:
    """
    Message which gets sent back to the parent through parent_queue
    """

    def __init__(self, message_type, target, content):
        """
        "message_type" is a string, and can be:
            "e" - error
            "a" - authentication failure
        """
        self.type = message_type
        self.target = target
        self.content = content


class Spiderling:
    """
    Enumerates SMB shares and spiders all possible directories/filenames up to maxdepth
    Designed to be threadable
    """

    # these extensions don't get parsed for content, unless explicitly specified
    dont_parse = [
        ".png",
        ".gif",
        ".tiff",
        ".msi",
        ".bmp",
        ".jpg",
        ".jpeg",
        ".zip",
        ".gz",
        ".bz2",
        ".7z",
        ".xz",
    ]

    def __init__(self, target, parent):

        try:
            self.parent = parent
            self.target = target

            # unless we're only searching local files, connect to target
            if type(self.target) == pathlib.PosixPath:
                self.local = True
                self.go()

            else:
                self.local = False

                self.smb_client = SMBClient(
                    target.host,
                    parent.username,
                    parent.password,
                    parent.domain,
                    parent.nthash,
                    parent.use_kerberos,
                    parent.aes_key,
                    parent.dc_ip,
                    port=target.port,
                )

                logon_result = self.smb_client.login()
                if logon_result not in [True, None]:
                    self.message_parent("a", logon_result)

                if logon_result is not None:
                    self.go()

            # file parsing parallelized one process at a time
            # allows file to be parsed while next one is being fetched
            self.parser_process = None

        except KeyboardInterrupt:
            log.critical("Spiderling Interrupted")

        # log all exceptions
        except Exception as e:
            if log.level <= logging.DEBUG:
                log.error(format_exc())
            else:
                log.error(f"Error in spiderling: {e}")

    def go(self):
        """
        go spider go spider go
        """

        # local files
        if self.local:
            if self.parent.parser.content_filters:
                self.parse_local_files(self.files)
            else:
                # just list the files
                list(self.files)

        else:
            # remote files
            for file in self.files:
                # if content searching is enabled, parse the file
                if self.parent.parser.content_filters:
                    try:
                        self.parser_process.join()
                    except AttributeError:
                        pass
                    self.parser_process = multiprocessing.Process(target=self.parse_file, args=(file,))
                    self.parser_process.start()

                # otherwise, just save it
                elif not self.local:
                    log.info(f"{self.target}: {file.share}\\{file.name} ({bytes_to_human(file.size)})")
                    if not self.parent.no_download:
                        self.save_file(file)

        log.info(f"Finished spidering {self.target}")

    @property
    def files(self):
        """
        Yields all files on the target to be parsed/downloaded
        Premptively download matching files into temp directory
        """

        if self.local:
            for file in list(list_files(self.target)):
                if self.extension_blacklisted(file):
                    log.debug(f"{self.target}: Skipping {file}: extension is blacklisted")
                    continue

                if self.parent.modified_after or self.parent.modified_before:
                    try:
                        mod_time = file.stat().st_mtime
                    except Exception:
                        mod_time = None

                    if not self.date_match(mod_time):
                        log.debug(f"Skipping {file}: does not match date filters")
                        continue

                if self.path_match(file) or (self.parent.or_logic and self.parent.parser.content_filters):
                    if self.path_match(file):
                        log.debug(pathlib.Path(file).relative_to(self.target))
                    if not self.is_binary_file(file):
                        yield file
                else:
                    log.debug(f"Skipping {file}: does not match filename/extension filters")

        else:
            for share in self.shares:
                for remote_file in self.list_files(share):
                    if not self.parent.no_download or self.parent.parser.content_filters:
                        self.get_file(remote_file)
                    yield remote_file

    def parse_file(self, file):
        """
        Simple wrapper around self.parent.parser.parse_file()
        For sole purpose of threading
        """

        try:
            if type(file) == RemoteFile:
                matches = self.parent.parser.parse_file(str(file.tmp_filename), pretty_filename=str(file))
                if matches and not self.parent.no_download:
                    self.save_file(file)
                else:
                    file.tmp_filename.unlink()

            else:
                log.debug(f"Found file: {file}")
                self.parent.parser.parse_file(file, file)

        # log all exceptions
        except Exception as e:
            if log.level <= logging.DEBUG:
                log.error(format_exc())
            else:
                log.error(f"Error parsing file {file}: {e}")

        except KeyboardInterrupt:
            log.critical("File parsing interrupted")

    @property
    def shares(self):
        """
        Lists all shares on single target
        Includes both enumerated shares and user-specified shares (which may be hidden from enumeration)
        """

        # Keep track of shares we've already yielded to avoid duplicates
        yielded_shares = set()

        # First, yield enumerated shares that match filters
        for share in self.smb_client.shares:
            if self.share_match(share):
                yielded_shares.add(share.lower())
                yield share

        # If user specified a share whitelist, also try those shares even if not enumerated
        # (some shares are hidden from enumeration but still accessible)
        if self.parent.share_whitelist:
            for share in self.parent.share_whitelist:
                if share.lower() not in yielded_shares:
                    # Check if it's blacklisted
                    if (not self.parent.share_blacklist) or (share.lower() not in self.parent.share_blacklist):
                        log.debug(f"{self.target}: Adding non-enumerated share from whitelist: {share}")
                        yielded_shares.add(share.lower())
                        yield share

    def list_files(self, share, path="", depth=0, tries=2):
        """
        List files inside a specific directory
        Only yield files which conform to all filters (except content)
        """

        if depth < self.parent.maxdepth and self.dir_match(path):
            files = []
            while tries > 0:
                try:
                    files = list(self.smb_client.ls(share, path))
                    break
                except FileListError as e:
                    if "ACCESS_DENIED" in str(e):
                        log.debug(f"{self.target}: Error listing files: {e}")
                        break
                    else:
                        tries -= 1

            if files:
                log.debug(f"{self.target}: {share}{path}: contains {len(files):,} items")

            for f in files:
                name = f.get_longname()
                full_path = f"{path}\\{name}"
                # if it's a directory, go deeper
                if f.is_directory():
                    for file in self.list_files(share, full_path, (depth + 1)):
                        yield file

                else:
                    # skip the file if it didn't match extension filters
                    if self.extension_blacklisted(name):
                        log.debug(f"{self.target}: Skipping {share}{full_path}: extension is blacklisted")
                        continue

                    if not self.path_match(name):
                        if not (
                            # all of these have to be true in order to get past this point
                            # "or logic" is enabled
                            self.parent.or_logic
                            and
                            # and file does not have a "don't parse" extension
                            (not self.is_binary_file(name))
                            and
                            # and content filters are enabled
                            self.parent.parser.content_filters
                        ):
                            log.debug(f"{self.target}: Skipping {share}{full_path}: filename/extensions do not match")
                            continue

                    # try to get the size of the file
                    try:
                        filesize = f.get_filesize()
                    except Exception as e:
                        self.smb_client.handle_impacket_error(e)
                        continue

                    if self.parent.modified_after or self.parent.modified_before:
                        try:
                            mod_time = f.get_mtime_epoch()
                        except Exception as e:
                            self.smb_client.handle_impacket_error(e)
                            mod_time = None

                        # check if file matches date filters
                        if not self.date_match(mod_time):
                            log.debug(f"{self.target}: Skipping {share}{full_path}: does not match date filters")
                            continue

                    # make the RemoteFile object (the file won't be read yet)
                    full_path_fixed = full_path.lstrip("\\")
                    remote_file = RemoteFile(full_path_fixed, share, self.target, size=filesize)

                    # if it's a non-empty file that's smaller than the size limit
                    if filesize > 0 and filesize < self.parent.max_filesize:
                        # if it matched filename/extension filters and we're downloading files
                        if (
                            self.parent.file_extensions or self.parent.filename_filters
                        ) and not self.parent.no_download:
                            # but the extension is marked as "don't parse"
                            if self.is_binary_file(name):
                                # don't parse it, instead save it and continue
                                log.info(f"{self.target}: {remote_file.share}\\{remote_file.name}")
                                if self.get_file(remote_file):
                                    self.save_file(remote_file)
                                    continue

                        # file is ready to be parsed
                        yield remote_file

                    else:
                        log.debug(f"{self.target}: {full_path} is either empty or too large")

    def path_match(self, file):
        """
        Based on whether "or" logic is enabled, return True or False
        if the filename + extension meets the requirements
        """
        filename_match = self.filename_match(file)
        extension_match = self.extension_whitelisted(file)
        if self.parent.or_logic:
            return (filename_match and self.parent.filename_filters) or (
                extension_match and self.parent.file_extensions
            )
        else:
            return filename_match and extension_match

    def share_match(self, share):
        """
        Return true if "share" matches any of the share filters
        """

        # if the share has been whitelisted
        if (not self.parent.share_whitelist) or (share.lower() in self.parent.share_whitelist):
            # and hasn't been blacklisted
            if (not self.parent.share_blacklist) or (share.lower() not in self.parent.share_blacklist):
                return True
            else:
                log.debug(f"{self.target}: Skipping blacklisted share: {share}")
        else:
            log.debug(f"{self.target}: Skipping share {share}: not in whitelist")

        return False

    def dir_match(self, path):
        """
        Return true if "path" matches any of the directory filters
        """

        # convert forward slashes to backwards
        dirname = str(path).lower().replace("/", "\\")

        # root path always passes
        if not path:
            return True

        # if whitelist check passes
        if (not self.parent.dir_whitelist) or any([k.lower() in dirname for k in self.parent.dir_whitelist]):
            # and blacklist check passes
            if (not self.parent.dir_blacklist) or not any([k.lower() in dirname for k in self.parent.dir_blacklist]):
                return True
            else:
                log.debug(f"{self.target}: Skipping blacklisted dir: {path}")
        else:
            log.debug(f"{self.target}: Skipping dir {path}: not in whitelist")

        return False

    def filename_match(self, filename):
        """
        Return true if "filename" matches any of the filename filters
        """

        if (not self.parent.filename_filters) or any(
            [f_regex.match(str(pathlib.Path(filename).stem)) for f_regex in self.parent.filename_filters]
        ):
            return True
        else:
            log.debug(f"{self.target}: {filename} does not match filename filters")

        return False

    def is_binary_file(self, filename):
        """
        Returns true if file is a bad extension type, e.g. encrypted or compressed
        """

        extension = "".join(pathlib.Path(filename).suffixes).lower()
        if any([extension.endswith(e.lower()) for e in self.dont_parse]):
            if extension not in self.parent.file_extensions:
                log.debug(f"{self.target}: Not parsing {filename} due to undesirable extension")
                return True
        return False

    def extension_blacklisted(self, filename):
        """
        Return True if folder, file name, or extension has been blacklisted
        """
        extension = "".join(pathlib.Path(filename).suffixes).lower()
        excluded_extensions = list(self.parent.extension_blacklist)

        if not excluded_extensions:
            return False

        if not any([extension.endswith(e) for e in excluded_extensions]):
            return False
        else:
            log.debug(f"{self.target}: Skipping file with blacklisted extension: {filename}")
            return True

    def extension_whitelisted(self, filename):
        """
        Return True if file extension has been whitelisted
        """
        # a .tar.gz file will match both filters ".gz" and ".tar.gz"
        extension = "".join(pathlib.Path(filename).suffixes).lower()
        extensions = list(self.parent.file_extensions)

        if not extensions:
            return True

        # if whitelist check passes
        if any([(extension.endswith(e) if e else extension == e) for e in extensions]):
            log.debug(f"{self.target}: {filename} matches extension filters")
            return True
        else:
            log.debug(f"{self.target}: Skipping file {filename}, does not match extension filters")
            return False

    def message_parent(self, message_type, content=""):
        """
        Send a message to the parent spider
        """

        self.parent.spiderling_queue.put(SpiderlingMessage(message_type, self.target, content))

    def parse_local_files(self, files):

        with ProcessPool(self.parent.threads) as pool:
            for r in pool.map(self.parse_file, files):
                pass

    def save_file(self, remote_file):
        """
        Moves a file from temp storage into the loot directory
        """

        allowed_chars = string.ascii_lowercase + string.ascii_uppercase + string.digits + "._ "

        # replace backslashes with underscores to preserve directory names
        loot_filename = str(remote_file).replace("\\", "_")
        # remove weird characters
        loot_filename = "".join([c for c in loot_filename if c in allowed_chars])
        loot_dest = self.parent.loot_dir / loot_filename
        try:
            move(str(remote_file.tmp_filename), str(loot_dest))
        except Exception:
            log.warning(f"Error saving {remote_file}")

    def get_file(self, remote_file):
        """
        Attempts to retrieve "remote_file" from share and returns True if successful
        """

        try:
            smb_client = self.parent.get_smb_client(self.target)
            log.debug(f"{self.target}: Downloading {remote_file.share}\\{remote_file.name}")
            remote_file.get(smb_client)
            return True
        except FileRetrievalError as e:
            log.debug(f"{self.target}: {e}")

        return False

    def date_match(self, file_time):
        """
        Return True if file modification time matches date filters
        file_time is a unix timestamp
        """

        if file_time is None:
            return True

        # Convert timestamp to datetime
        file_date = datetime.fromtimestamp(file_time)

        # Check modified_after
        if self.parent.modified_after:
            if file_date < self.parent.modified_after:
                log.debug(f"{self.target}: File too old: {file_date.strftime('%Y-%m-%d')}")
                return False

        # Check modified_before
        if self.parent.modified_before:
            if file_date > self.parent.modified_before:
                log.debug(f"{self.target}: File too new: {file_date.strftime('%Y-%m-%d')}")
                return False

        return True
