import json
from typing import Optional
from unittest.mock import Mock

from assertpy import assert_that
from messagebus import Message
from messagebusmock import MessageBusSpy

from httpmessageprocessor import ApiRequest
from httpmessageprocessorutil.automapped import AutoMappedMessageProcessor
from httpmessageprocessorutil.sanitizer import Sanitizer, Nh3Sanitizer

XSS_SCRIPT_TAG = '<script>alert("xss")</script>'
XSS_IMG_TAG = '<img src=x onerror=alert(1)>'
XSS_EVENT_ATTR = '"><button onwebkitmouseforcechanged=alert(sessionStorage.getItem("key"))>'

_SANITIZER = Nh3Sanitizer()


def _sanitize(value: str) -> str:
    return _SANITIZER.sanitize(value)


def _get_sanitizer_mock() -> Mock:
    mock = Mock(Sanitizer)
    mock.sanitize.side_effect = _sanitize
    mock.sanitize_request.side_effect = _sanitize
    mock.sanitize_response.side_effect = _sanitize
    return mock

def _get_sanitizer_sanitize_mock_calls(sanitizer_mock) -> list:
    calls = []
    for method in [sanitizer_mock.sanitize_request, sanitizer_mock.sanitize_response]:
        calls.extend([mc[1][0] for mc in method.mock_calls])
    return calls

class ParameterInteractor(object):
    def execute(self, obj):
        return obj


class NoParameterInteractor(object):
    def execute(self):
        pass


def _build_api_request(
    method="POST",
    path_parameters=None,
    query_parameters=None,
    headers=None,
    body=None,
):
    body_str = body or ""
    if not isinstance(body, str):
        body_str = json.dumps(body or {}, default=str)
    return ApiRequest(
        method=method,
        api_path="/test",
        path_parameters=path_parameters or {},
        query_parameters=query_parameters or {},
        headers=headers or {},
        body=body_str.encode("UTF-8"),
    )


class SimpleDto(object):
    def __init__(self, name: str, value: Optional[str] = None):
        self.name = name
        self.value = value


msg_bus = MessageBusSpy()
msg = Message(0)


class HeaderDto(object):
    def __init__(self, host: str, value: str):
        self.host = host
        self.value = value


_HEADER_EXTRA_MAPPING = {
    "host": {"location": "header", "name": "X-Host", "type": str},
    "value": {"location": "header", "name": "X-Value", "type": str},
}


class TestSanitizerOnHeaderParameter:
    def test_xssScriptTagInHeader_isSanitized(self):
        processor = AutoMappedMessageProcessor(
            HeaderDto, ParameterInteractor(), extra_mapping=_HEADER_EXTRA_MAPPING
        )

        result = processor.parse_data(
            _build_api_request(headers={"X-Host": XSS_SCRIPT_TAG, "X-Value": "safe"}),
            event_type=None,
        )

        assert_that(result.host).is_equal_to(_sanitize(XSS_SCRIPT_TAG))
        assert_that(result.value).is_equal_to("safe")

    def test_xssImgTagInHeader_isSanitized(self):
        processor = AutoMappedMessageProcessor(
            HeaderDto, ParameterInteractor(), extra_mapping=_HEADER_EXTRA_MAPPING
        )

        result = processor.parse_data(
            _build_api_request(headers={"X-Host": XSS_IMG_TAG, "X-Value": "safe"}),
            event_type=None,
        )

        assert_that(result.host).is_equal_to(_sanitize(XSS_IMG_TAG))

    def test_xssEventAttributeInHeader_isSanitized(self):
        processor = AutoMappedMessageProcessor(
            HeaderDto, ParameterInteractor(), extra_mapping=_HEADER_EXTRA_MAPPING
        )

        result = processor.parse_data(
            _build_api_request(headers={"X-Host": XSS_EVENT_ATTR, "X-Value": "safe"}),
            event_type=None,
        )

        assert_that(result.host).is_equal_to(_sanitize(XSS_EVENT_ATTR))

    def test_xssInMultipleHeaders_allAreSanitized(self):
        processor = AutoMappedMessageProcessor(
            HeaderDto, ParameterInteractor(), extra_mapping=_HEADER_EXTRA_MAPPING
        )

        result = processor.parse_data(
            _build_api_request(
                headers={"X-Host": XSS_SCRIPT_TAG, "X-Value": XSS_IMG_TAG}
            ),
            event_type=None,
        )

        assert_that(result.host).is_equal_to(_sanitize(XSS_SCRIPT_TAG))
        assert_that(result.value).is_equal_to(_sanitize(XSS_IMG_TAG))

    def test_withoutSanitizer_xssInHeaderIsPreservedAsIs(self):
        processor = AutoMappedMessageProcessor(
            HeaderDto, ParameterInteractor(),
            extra_mapping=_HEADER_EXTRA_MAPPING,
            sanitizer=None,
        )

        result = processor.parse_data(
            _build_api_request(headers={"X-Host": XSS_SCRIPT_TAG, "X-Value": "safe"}),
            event_type=None,
        )

        assert_that(result.host).is_equal_to(XSS_SCRIPT_TAG)

    def test_sanitizerIsCalledForEachHeaderField(self):
        sanitizer = _get_sanitizer_mock()
        processor = AutoMappedMessageProcessor(
            HeaderDto, ParameterInteractor(),
            extra_mapping=_HEADER_EXTRA_MAPPING,
            sanitizer=sanitizer,
        )

        processor.parse_data(
            _build_api_request(
                headers={"X-Host": XSS_SCRIPT_TAG, "X-Value": XSS_IMG_TAG}
            ),
            event_type=None,
        )

        assert_that(_get_sanitizer_sanitize_mock_calls(sanitizer)).contains(
            XSS_SCRIPT_TAG, XSS_IMG_TAG
        )


class TestSanitizerOnParseData:
    def test_xssInBody_isSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor()
        )

        result = processor.parse_data(
            _build_api_request(body={"name": XSS_SCRIPT_TAG, "value": "safe"}),
            event_type=None,
        )

        assert_that(result.name).is_equal_to(_sanitize(XSS_SCRIPT_TAG))
        assert_that(result.value).is_equal_to("safe")

    def test_xssImgTagInBody_isSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor()
        )

        result = processor.parse_data(
            _build_api_request(body={"name": XSS_IMG_TAG}),
            event_type=None,
        )

        assert_that(result.name).is_equal_to(_sanitize(XSS_IMG_TAG))

    def test_xssInQueryParameter_isSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor()
        )

        result = processor.parse_data(
            _build_api_request(
                method="GET",
                query_parameters={"name": XSS_EVENT_ATTR},
            ),
            event_type=None,
        )

        assert_that(result.name).is_equal_to(_sanitize(XSS_EVENT_ATTR))

    def test_xssImgTagInQueryParameter_isSanitized(self):
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor())

        result = processor.parse_data(
            _build_api_request(
                method="GET",
                query_parameters={"name": XSS_IMG_TAG},
            ),
            event_type=None,
        )

        assert_that(result.name).is_equal_to(_sanitize(XSS_IMG_TAG))

    def test_xssInMultipleQueryParameters_allAreSanitized(self):
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor())

        result = processor.parse_data(
            _build_api_request(
                method="GET",
                query_parameters={"name": XSS_SCRIPT_TAG, "value": XSS_IMG_TAG},
            ),
            event_type=None,
        )

        assert_that(result.name).is_equal_to(_sanitize(XSS_SCRIPT_TAG))
        assert_that(result.value).is_equal_to(_sanitize(XSS_IMG_TAG))

    def test_withoutSanitizer_xssInQueryParameterIsPreservedAsIs(self):
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor(), sanitizer=None)

        result = processor.parse_data(
            _build_api_request(
                method="GET",
                query_parameters={"name": XSS_SCRIPT_TAG},
            ),
            event_type=None,
        )

        assert_that(result.name).is_equal_to(XSS_SCRIPT_TAG)

    def test_sanitizerIsCalledForQueryParameterFields(self):
        sanitizer = _get_sanitizer_mock()
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor(), sanitizer=sanitizer)

        processor.parse_data(
            _build_api_request(
                method="GET",
                query_parameters={"name": XSS_SCRIPT_TAG, "value": XSS_IMG_TAG},
            ),
            event_type=None,
        )

        assert_that(_get_sanitizer_sanitize_mock_calls(sanitizer)).contains(
            XSS_SCRIPT_TAG, XSS_IMG_TAG
        )

    def test_xssInPathParameter_isSanitized(self):
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor())

        result = processor.parse_data(
            _build_api_request(
                method="GET",
                path_parameters={"name": XSS_SCRIPT_TAG},
            ),
            event_type=None,
        )

        assert_that(result.name).is_equal_to(_sanitize(XSS_SCRIPT_TAG))

    def test_xssImgTagInPathParameter_isSanitized(self):
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor())

        result = processor.parse_data(
            _build_api_request(
                method="GET",
                path_parameters={"name": XSS_IMG_TAG},
            ),
            event_type=None,
        )

        assert_that(result.name).is_equal_to(_sanitize(XSS_IMG_TAG))

    def test_xssInMultiplePathParameters_allAreSanitized(self):
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor())

        result = processor.parse_data(
            _build_api_request(
                method="GET",
                path_parameters={"name": XSS_SCRIPT_TAG, "value": XSS_EVENT_ATTR},
            ),
            event_type=None,
        )

        assert_that(result.name).is_equal_to(_sanitize(XSS_SCRIPT_TAG))
        assert_that(result.value).is_equal_to(_sanitize(XSS_EVENT_ATTR))

    def test_withoutSanitizer_xssInPathParameterIsPreservedAsIs(self):
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor(), sanitizer=None)

        result = processor.parse_data(
            _build_api_request(
                method="GET",
                path_parameters={"name": XSS_SCRIPT_TAG},
            ),
            event_type=None,
        )

        assert_that(result.name).is_equal_to(XSS_SCRIPT_TAG)

    def test_sanitizerIsCalledForPathParameterFields(self):
        sanitizer = _get_sanitizer_mock()
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor(), sanitizer=sanitizer)

        processor.parse_data(
            _build_api_request(
                method="GET",
                path_parameters={"name": XSS_SCRIPT_TAG, "value": XSS_IMG_TAG},
            ),
            event_type=None,
        )

        assert_that(_get_sanitizer_sanitize_mock_calls(sanitizer)).contains(
            XSS_SCRIPT_TAG, XSS_IMG_TAG
        )

    def test_multipleXssStringsInBody_allAreSanitized(self):
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor())

        result = processor.parse_data(
            _build_api_request(
                body={"name": XSS_SCRIPT_TAG, "value": XSS_IMG_TAG}
            ),
            event_type=None,
        )

        assert_that(result.name).is_equal_to(_sanitize(XSS_SCRIPT_TAG))
        assert_that(result.value).is_equal_to(_sanitize(XSS_IMG_TAG))

    def test_withoutSanitizer_xssInputIsPreservedAsIs(self):
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor(), sanitizer=None)

        result = processor.parse_data(
            _build_api_request(body={"name": XSS_SCRIPT_TAG}),
            event_type=None,
        )

        assert_that(result.name).is_equal_to(XSS_SCRIPT_TAG)

    def test_sanitizerIsCalledForEachStringField(self):
        sanitizer = Mock(Sanitizer)
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor(), sanitizer=sanitizer)

        processor.parse_data(
            _build_api_request(
                body={"name": XSS_SCRIPT_TAG, "value": XSS_IMG_TAG}
            ),
            event_type=None,
        )

        assert_that(_get_sanitizer_sanitize_mock_calls(sanitizer)).contains(XSS_SCRIPT_TAG, XSS_IMG_TAG)


class TestSanitizerOnFormatResponse:
    def test_xssInResponseObject_isSanitized(self):
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor())

        processor.format_response(
            msg_bus, msg, None, None, SimpleDto(name=XSS_SCRIPT_TAG)
        )

        body = json.loads(msg_bus.reply_message_call.args[1].body)
        assert_that(body["name"]).is_equal_to(_sanitize(XSS_SCRIPT_TAG))

    def test_xssImgTagInResponseObject_isSanitized(self):
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor())

        processor.format_response(
            msg_bus, msg, None, None, SimpleDto(name=XSS_IMG_TAG, value=XSS_EVENT_ATTR)
        )

        body = json.loads(msg_bus.reply_message_call.args[1].body)
        assert_that(body["name"]).is_equal_to(_sanitize(XSS_IMG_TAG))
        assert_that(body["value"]).is_equal_to(_sanitize(XSS_EVENT_ATTR))

    def test_withoutSanitizer_xssInResponseIsPreservedAsIs(self):
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor(), sanitizer=None)

        processor.format_response(
            msg_bus, msg, None, None, SimpleDto(name=XSS_SCRIPT_TAG)
        )

        body = json.loads(msg_bus.reply_message_call.args[1].body)
        assert_that(body["name"]).is_equal_to(XSS_SCRIPT_TAG)

    def test_sanitizerIsCalledForEachStringFieldInResponse(self):
        sanitizer = _get_sanitizer_mock()
        processor = AutoMappedMessageProcessor(SimpleDto, ParameterInteractor(), sanitizer=sanitizer)

        processor.format_response(
            msg_bus, msg, None, None, SimpleDto(name=XSS_SCRIPT_TAG, value=XSS_IMG_TAG)
        )

        assert_that(_get_sanitizer_sanitize_mock_calls(sanitizer)).contains(XSS_SCRIPT_TAG, XSS_IMG_TAG)
