# -*- 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,
)

from contrast.agent import scope
from contrast.api import Library
from contrast.utils.patch_utils import get_loaded_modules
from contrast_vendor 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_ts_message_func):
        self.settings = settings
        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

        from contrast.reporting import teamserver_messages

        msg = teamserver_messages.LibraryUsage(reportable_dists)

        logger.debug("Sending list of modules loaded to Contrast")

        self.send_ts_message_func(msg)


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.
    Then, assuming library was loaded, appends that metadata to the self.libraries.

    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):
    """
    Generate a json-serializable representation of a package for teamserver

    :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
