from unittest import mock

import pytest

from ddtrace.appsec._iast._iast_request_context import get_iast_reporter
from ddtrace.appsec._iast._taint_tracking import OriginType
from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
from ddtrace.appsec._iast._taint_tracking.aspects import add_aspect
from ddtrace.appsec._iast.constants import VULN_SSRF
from ddtrace.appsec._iast.taint_sinks.ssrf import _iast_report_ssrf
from tests.appsec.iast.iast_utils import _end_iast_context_and_oce
from tests.appsec.iast.iast_utils import _get_iast_data
from tests.appsec.iast.iast_utils import _iast_patched_module
from tests.appsec.iast.iast_utils import _start_iast_context_and_oce
from tests.appsec.iast.iast_utils import get_line_and_hash
from tests.appsec.iast.taint_sinks._taint_sinks_utils import NON_TEXT_TYPES_TEST_DATA


FIXTURES_PATH = "tests/appsec/iast/fixtures/taint_sinks/ssrf.py"


def _get_tainted_url():
    tainted_path = taint_pyobject(
        pyobject="forbidden_dir/",
        source_name="test_ssrf",
        source_value="forbidden_dir/",
        source_origin=OriginType.PARAMETER,
    )
    return add_aspect("http://localhost/", tainted_path), tainted_path


def _check_report(tainted_path, label):
    data = _get_iast_data()

    vulnerability = data["vulnerabilities"][0]
    source = data["sources"][0]
    assert vulnerability["type"] == VULN_SSRF
    assert vulnerability["evidence"]["valueParts"] == [
        {"value": "http://localhost/"},
        {"source": 0, "value": tainted_path},
    ]
    assert "value" not in vulnerability["evidence"].keys()
    assert vulnerability["evidence"].get("pattern") is None
    assert vulnerability["evidence"].get("redacted") is None
    assert source["name"] == "test_ssrf"
    assert source["origin"] == OriginType.PARAMETER
    assert source["value"] == tainted_path

    line, hash_value = get_line_and_hash(label, VULN_SSRF, filename=FIXTURES_PATH)
    assert vulnerability["location"]["path"] == FIXTURES_PATH
    assert vulnerability["location"]["line"] == line
    assert vulnerability["location"]["method"] == label
    assert "class" not in vulnerability["location"]
    assert type(vulnerability["location"]["spanId"]) is int
    assert vulnerability["hash"] == hash_value


def test_ssrf_requests(iast_context_defaults):
    tainted_url, tainted_path = _get_tainted_url()
    mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.ssrf")

    mod.pt_requests_get(tainted_url)

    _check_report(tainted_path, "pt_requests_get")


def test_ssrf_urllib3(iast_context_defaults):
    tainted_url, tainted_path = _get_tainted_url()
    mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.ssrf")

    mod.pt_urllib3_request(tainted_url)

    _check_report(tainted_path, "pt_urllib3_request")


def test_ssrf_urllib3_poolmanager(iast_context_defaults):
    tainted_url, tainted_path = _get_tainted_url()
    mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.ssrf")

    mod.pt_urllib3_poolmanager(tainted_url)

    _check_report(tainted_path, "pt_urllib3_poolmanager")


def test_ssrf_httplib(iast_context_defaults):
    tainted_url, tainted_path = _get_tainted_url()
    mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.ssrf")

    mod.pt_httplib_request(tainted_url)

    _check_report(tainted_path, "pt_httplib_request")


def test_ssrf_webbrowser(iast_context_defaults):
    tainted_url, tainted_path = _get_tainted_url()
    mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.ssrf")
    mod.pt_webbrowser_open(tainted_url)

    _check_report(tainted_path, "pt_webbrowser_open")


def test_urllib_request(iast_context_defaults):
    tainted_url, tainted_path = _get_tainted_url()
    mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.ssrf")

    mod.pt_urllib_request(tainted_url)

    _check_report(tainted_path, "pt_urllib_request")


def _check_no_report_if_deduplicated(num_vuln_expected):
    span_report = get_iast_reporter()
    if num_vuln_expected == 0:
        assert span_report is None
    else:
        assert span_report

        assert len(span_report.vulnerabilities) == num_vuln_expected


def test_ssrf_requests_deduplication(iast_context_deduplication_enabled):
    mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.ssrf")

    for num_vuln_expected in [1, 0, 0]:
        _start_iast_context_and_oce()
        tainted_url, tainted_path = _get_tainted_url()
        for _ in range(0, 5):
            mod.pt_requests_get(tainted_url)

        _check_no_report_if_deduplicated(num_vuln_expected)
        _end_iast_context_and_oce()


def test_ssrf_urllib3_deduplication(iast_context_deduplication_enabled):
    mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.ssrf")
    for num_vuln_expected in [1, 0, 0]:
        _start_iast_context_and_oce()
        tainted_url, tainted_path = _get_tainted_url()
        for _ in range(0, 5):
            mod.pt_urllib3_request(tainted_url)

        _check_no_report_if_deduplicated(num_vuln_expected)
        _end_iast_context_and_oce()


def test_ssrf_httplib_deduplication(iast_context_deduplication_enabled):
    mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.ssrf")

    for num_vuln_expected in [1, 0, 0]:
        _start_iast_context_and_oce()
        tainted_url, tainted_path = _get_tainted_url()
        for _ in range(0, 5):
            mod.pt_httplib_request(tainted_url)

        _check_no_report_if_deduplicated(num_vuln_expected)
        _end_iast_context_and_oce()


def test_ssrf_webbrowser_deduplication(iast_context_deduplication_enabled):
    mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.ssrf")

    for num_vuln_expected in [1, 0, 0]:
        _start_iast_context_and_oce()
        tainted_url, tainted_path = _get_tainted_url()
        for _ in range(0, 5):
            try:
                # label test_ssrf_webbrowser_deduplication
                mod.pt_webbrowser_open(tainted_url)
            except ConnectionError:
                pass

        _check_no_report_if_deduplicated(num_vuln_expected)
        _end_iast_context_and_oce()


def test_ssrf_urllib_deduplication(iast_context_deduplication_enabled):
    mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.ssrf")

    for num_vuln_expected in [1, 0, 0]:
        _start_iast_context_and_oce()
        tainted_url, tainted_path = _get_tainted_url()
        for _ in range(0, 5):
            mod.pt_urllib_request(tainted_url)

        _check_no_report_if_deduplicated(num_vuln_expected)
        _end_iast_context_and_oce()


@pytest.mark.parametrize("non_text_obj,obj_type", NON_TEXT_TYPES_TEST_DATA)
def test_ssrf_non_text_types_no_vulnerability(non_text_obj, obj_type, iast_context_defaults):
    """Test that non-text types don't trigger SSRF vulnerabilities."""
    # Taint the non-text object
    tainted_obj = taint_pyobject(
        non_text_obj,
        source_name="test_source",
        source_value=str(non_text_obj),
        source_origin=OriginType.PARAMETER,
    )

    # Mock a function that would be in the SSRF function mapping
    def mock_request_func(url):
        return f"Response for {url}"

    # Mock get_argument_value to return our tainted non-text object
    with mock.patch("ddtrace.appsec._iast.taint_sinks.ssrf.get_argument_value", return_value=tainted_obj):
        # Add our mock function to the SSRF function mapping temporarily
        from ddtrace.appsec._iast.taint_sinks.ssrf import _FUNC_TO_URL_ARGUMENT

        original_mapping = _FUNC_TO_URL_ARGUMENT.copy()
        _FUNC_TO_URL_ARGUMENT["my_testing_cool_module"] = (0, "url")

        try:
            # Call the SSRF reporting function
            _iast_report_ssrf("my_cool_function", "my_testing_cool_module", mock_request_func, tainted_obj)
        finally:
            # Restore original mapping
            _FUNC_TO_URL_ARGUMENT.clear()
            _FUNC_TO_URL_ARGUMENT.update(original_mapping)

    # Assert no vulnerability was reported
    span_report = get_iast_reporter()
    assert span_report is None, f"Vulnerability reported for {obj_type}: {non_text_obj}"
