# -*- coding: utf-8 -*-
# Copyright © 2023 Contrast Security, Inc.
# See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
import os
import time
import threading

from pip._vendor import pkg_resources

from importlib.util import find_spec

from contrast.utils import get_installed_distributions
from contrast.utils.library_reader.utils import (
    normalize_file_name,
    get_top_level_directories,
    get_data,
    get_hash,
    CONTRAST_AGENT_DIST,
    PY_SUFFIX,
    SO_SUFFIX,
    append_files_loaded_to_activity,
)

from contrast.agent import scope
from contrast.api import Library
from contrast.api.dtm_pb2 import Activity
from contrast.utils.patch_utils import get_loaded_modules
from contrast.extern import structlog as logging

logger = logging.getLogger("contrast")


def get_active_library_names_from_pkg():
    library_names = ["contrast"]
    libraries = list(pkg_resources.working_set)
    for library in libraries:
        library_names.append(library.key)
    return library_names


class LibraryReader(object):
    def __init__(self, settings, send_sr_message_func, send_ts_message_func):
        self.settings = settings
        # TODO: PYT-2140 remove send sr msg func once we deprecate SR
        self.send_sr_message_func = send_sr_message_func
        self.send_ts_message_func = send_ts_message_func
        self.analysis_thread = None
        self.installed_distribution_keys = []

        # All packages installed in the site-packages/ directory
        self.installed_dists = get_installed_distributions()

    def start_library_analysis_thread(self, daemon=True):
        self.analysis_thread = threading.Thread(target=self._read_libraries)
        self.analysis_thread.daemon = daemon

        self.analysis_thread.start()

    def join_library_analysis_thread(self):
        if self.analysis_thread:
            self.analysis_thread.join()

    def _read_libraries(self):
        """
        Looks at every library installed in self.installed_dists, then calls search_dist
        on each dist that has a file top_level.txt, which is autogenerated by installation,
        specifying which files/directories are associated with the package

        :exception: IOError if top_level.txt, METADATA, or RECORD cannot be found
        :return: None
        """
        with scope.contrast_scope():
            logger.debug("Analyzing libraries...")

            all_dists = read_dists(self.installed_dists)
            reportable_dists = []

            for dist in all_dists:
                library_name = dist.get("file_path", "")
                if library_name and library_name != CONTRAST_AGENT_DIST:
                    reportable_dists.append(dist)

            self._send_analysis_results_appstart(reportable_dists)
            self._send_analysis_results_files_loaded(reportable_dists)

    def _send_analysis_results_appstart(self, reportable_dists):
        """
        Send library discovery in ApplicationUpdate message.
        :param reportable_dists: A list of dictionaries containing information about
            each distribution.
        """
        if not reportable_dists:
            return

        libraries = [Library(dist) for dist in reportable_dists]

        logger.debug(
            "Discovered libraries: %s",
            [lib.file_path for lib in libraries],
        )

        logger.debug("Sending ApplicationUpdate message with library analysis results")

        # Import here to prevent circular import
        from contrast.reporting import teamserver_messages

        app_update_msg = teamserver_messages.ApplicationUpdate(libraries)
        self.send_ts_message_func(app_update_msg)

    def _send_analysis_results_files_loaded(self, reportable_dists):
        if not reportable_dists:
            return

        if self.settings.config.is_service_bypassed:
            from contrast.reporting import teamserver_messages

            msg = teamserver_messages.ObservedLibraryUsage(reportable_dists)

            logger.debug(
                "Sending list of modules loaded to Contrast", direct_to_teamserver=1
            )

            self.send_ts_message_func(msg)
        else:
            total_loaded_files_cnt = 0
            activity_dtm = Activity()
            send_msg = False

            for dist in reportable_dists:
                if len(dist["used_files"]) > 0:
                    send_msg = True
                    dist_hash = dist["hash_code"]

                    total_loaded_files_cnt += append_files_loaded_to_activity(
                        activity_dtm, dist["used_files"], dist_hash
                    )

            if send_msg:
                self.send_sr_message_func([activity_dtm])

                logger.debug(
                    "Sending Activity message with %s files loaded",
                    total_loaded_files_cnt,
                )


def read_dists(installed_dists):
    results = []

    for dist in installed_dists:
        analysis_result = search_dist(dist)

        if analysis_result:
            results.append(analysis_result)

    return results


def search_dist(dist):
    """
    Searches directories related to dist, gathering relevant statistics and metadata for
    LibraryDTM. Then, assuming library was loaded, appends that metadata to the self.libraries
    in a LibraryDTM

    Created package is added to the output for the process

    :param dist: contains a pkg_resources.DistInfoDistribution object
    referencing the current library being looked for
    :return: A dictionary containing contrast-relevant information about the current distribution
    """
    # __init__.py is required in order to be a package. If there are multiple subpackages
    # we want to make sure we include all __init__.py files as part of the count
    total_files = list()
    total_used_files = set()
    top_level_dirs = get_top_level_directories(dist)

    for directory in top_level_dirs:
        package_path = os.sep.join([dist.location, directory])

        all_dir_files = get_all_files(package_path)
        dir_used_files = search_modules(package_path, directory)

        total_files.extend(all_dir_files)
        total_used_files.update(dir_used_files)

    file_count = len(total_files)
    # Set the package path to be the first directory mentioned in top_level.txt
    name = str(dist).split(" ")[0]

    version, manifest, url = get_data(dist)

    library_hash = get_hash(name, version)

    result = create_package(version, manifest, file_count, name, url, library_hash)
    result["used_files"] = total_used_files

    return result


def create_package(version, manifest, class_count, name, url, sha):
    """
    Add package with appropriate metadata to the library; reason for this is that protobuf cannot
    be pickled.

    :param version: version of package
    :param manifest: all metadata
    :param class_count: count of files
    :param name: name of the library
    :param url: homepage of package
    :param sha: sha1 hash of the name(space)version
    :return: dict of package
    """
    current_time = int(time.time() * 1000)

    return {
        "version": version,
        "manifest": manifest,
        "class_count": class_count,
        "file_path": name,
        "url": url,
        "hash_code": sha,
        "external_ms": current_time,
        "internal_ms": current_time,
    }


def search_modules(package_path, directory_name):
    """
    Searches every module to see if it is associated with the current package being looked at.
    If it is, update used file count.

    :param directory_name: name of top level directory (i.e package or subpackage name)
    :param package_path: Path to the package currently being looked for
    :return: The list of used files for this directory, package_found, name, version, manifest,
        and url
    name, version, manifest, and url are empty strings if the package was already found
    """
    used_files = set()
    sys_modules = get_loaded_modules().values()
    top_level_package = directory_name + "."

    for module in sys_modules:
        if (
            module.__name__
            and (module.__name__ + ".").startswith(top_level_package)
            and hasattr(module, "__file__")
            and module.__file__
        ):
            # First we need to check and see if the top-level package (the name up till the
            # first dot) of __name__ is the same as the current directory_name/module name

            # Then, process __file__ (Path to the loaded file)
            # __file__ is optional. If set, this attribute’s value must be a string. The import
            # system may opt to leave __file__ unset if it has no semantic meaning
            # (e.g. a module loaded from a database).

            # using os.path.realpath because in some environments there is a symlink to the
            # lib/python_version_here/site-packages directory. When checking module.__file__
            # the symlink path is used instead of the resolved path. Resolved path is needed
            # because that is what the package_path will point to.
            module_file = os.path.realpath(module.__file__)

            file_path = module_file if module_file.startswith(package_path) else ""
            if len(file_path) > 0:
                used_files.add(normalize_file_name(file_path))

    return used_files


def find_module_file(module_name, module_location):
    """
    If module is found, returns the module file path under the location specified
    by module_location.

    :param module_name: Name of the module
    :param module_location: Path to the module file
    """
    module_file = None

    try:
        # A module spec contains information used to import a module. Such information,
        # origin/location is needed here to get the location of the module found under
        # module_location. PEP 451 contains additional info on this.
        module_spec = find_spec(module_name)
        if hasattr(module_spec, "origin") and module_spec.origin:
            if module_spec.origin.startswith(module_location):
                module_file = module_spec.origin
        elif hasattr(module_spec, "location") and module_spec.location:
            if module_spec.location.startswith(module_location):
                module_file = module_spec.location
    except Exception as e:
        # Raises ModuleNotFoundError (new in 3.6+) instead of AttributeError if package is in
        # fact not a package (i.e. lacks a __path__ attribute). This is something that
        # shouldn't happen so we log it here
        logger.debug(
            "WARNING - Module %s at location %s is not a package: %s",
            module_name,
            module_location,
            e,
        )

    if module_file and os.path.isfile(module_file):
        return module_file

    return None


def get_all_files(package_path):
    """
    Get all .py and .so files in all directories in package_path.

    :param package_path: Path to the current package being looked or a path to the module file.
        If its a module, we need to get the full path to the module file
    since package_path only includes the first portion of the module name.
    :return: list of files
    """
    total_files = []

    if not os.path.isdir(package_path):
        module_name = os.path.basename(package_path)
        module_location = os.path.dirname(package_path)

        module_file = find_module_file(module_name, module_location)
        if module_file:
            total_files.append(module_file)
        # These next 2 cases can occur if __location__ or __origin__ exists as an attr
        # but they don't point to the location of the module
        # This happened for bottle
        elif os.path.exists(package_path + PY_SUFFIX):
            total_files.append(package_path + PY_SUFFIX)

        return total_files

    for _, _, files in os.walk(package_path):
        for f in files:
            if f.endswith((PY_SUFFIX, SO_SUFFIX)):
                total_files.append(f)

    return total_files
