from inspect import currentframe
from inspect import getframeinfo
import traceback

import pytest

import ddtrace.appsec._common_module_patches as cmp


def test_lfi_normal_exception():
    """
    Ensure the top frame is the one where the exception is raised in the customer code
    """
    exception_repr = """Traceback (most recent call last):
  File "{}", line {}, in test_lfi_normal_exception
    with open("/unknown/do_not_exist_test.txt", "w"):
"""
    try:
        cmp.patch_common_modules()
        with pytest.raises(Exception) as e:
            with open("/unknown/do_not_exist_test.txt", "w"):
                pass
        assert e.type is FileNotFoundError
        # ensure the last frame is from the file where open was called
        assert e.traceback[-1].path.as_posix() == __file__
        # Does not work as we can't remove futur frames at raising point
        # assert len(e.traceback) == 1
        line_number = getframeinfo(currentframe()).lineno
        try:
            with open("/unknown/do_not_exist_test.txt", "w"):
                pass
        except Exception as e:
            assert e.__class__.__name__ == "FileNotFoundError"
            assert e.__traceback__.tb_frame.f_code.co_filename == __file__
            assert traceback.format_exc(limit=1).startswith(exception_repr.format(__file__, line_number + 2))
    finally:
        cmp.unpatch_common_modules()


def get_result(v: str) -> str:
    return "A" + v


def test_wrapping_twice():
    """
    Ensure wrapping mechanism is not wrapping twice the same function in a row, but can still wrap different functions
    """

    def wrapper1(original, instance, args, kargs):
        return original("B" + args[0])

    def wrapper2(original, instance, args, kargs):
        return original("C" + args[0])

    assert get_result("1") == "A1"
    cmp.try_wrap_function_wrapper(__name__, "get_result", wrapper1)
    assert get_result("1") == "AB1"
    # different wrap is applied
    cmp.try_wrap_function_wrapper(__name__, "get_result", wrapper2)
    assert get_result("1") == "ABC1"
    cmp.try_wrap_function_wrapper(__name__, "get_result", wrapper2)
    # second wrap with same wrapper is ignored
    assert get_result("1") == "ABC1"
    cmp.try_unwrap(__name__, "get_result")
    assert get_result("1") == "A1"
    cmp.try_wrap_function_wrapper(__name__, "get_result", wrapper1)
    assert get_result("1") == "AB1"
    # second wrap with same wrapper is ignored
    cmp.try_wrap_function_wrapper(__name__, "get_result", wrapper1)
    assert get_result("1") == "AB1"
    cmp.try_unwrap(__name__, "get_result")
    assert get_result("1") == "A1"
