import json
from typing import Optional

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, SanitizeMode

XSS_SCRIPT_TAG = '<script>alert("xss")</script>'
XSS_IMG_TAG = '<img src=x onerror=alert(1)>'

_SANITIZER = Nh3Sanitizer()


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


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


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


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


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

msg_bus = MessageBusSpy()
msg = Message(0)


class TestSanitizeModeRequestOnly:
    def test_requestBody_isSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor(),
            sanitize_mode=SanitizeMode.REQUEST_ONLY,
        )

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

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

    def test_queryParameter_isSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor(),
            sanitize_mode=SanitizeMode.REQUEST_ONLY,
        )

        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(_sanitize(XSS_SCRIPT_TAG))

    def test_pathParameter_isSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor(),
            sanitize_mode=SanitizeMode.REQUEST_ONLY,
        )

        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_header_isSanitized(self):
        processor = AutoMappedMessageProcessor(
            HeaderDto, ParameterInteractor(),
            extra_mapping=_HEADER_EXTRA_MAPPING,
            sanitize_mode=SanitizeMode.REQUEST_ONLY,
        )

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

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

    def test_response_isNotSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor(),
            sanitize_mode=SanitizeMode.REQUEST_ONLY,
        )

        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)


class TestSanitizeModeResponseOnly:
    def test_requestBody_isNotSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor(),
            sanitize_mode=SanitizeMode.RESPONSE_ONLY,
        )

        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_queryParameter_isNotSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor(),
            sanitize_mode=SanitizeMode.RESPONSE_ONLY,
        )

        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(XSS_IMG_TAG)

    def test_response_isSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor(),
            sanitize_mode=SanitizeMode.RESPONSE_ONLY,
        )

        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))


class TestSanitizeModeNone:
    def test_requestBody_isNotSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor(),
            sanitize_mode=SanitizeMode.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_response_isNotSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor(),
            sanitize_mode=SanitizeMode.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)


class TestSanitizeModeDefault:
    def test_request_isSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor(),
        )

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

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

    def test_response_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))


class TestSanitizeModeEdgeCases:
    def test_sanitizerNone_withRequestOnlyMode_nothingIsSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor(),
            sanitizer=None,
            sanitize_mode=SanitizeMode.REQUEST_ONLY,
        )

        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_sanitizerNone_withResponseOnlyMode_nothingIsSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor(),
            sanitizer=None,
            sanitize_mode=SanitizeMode.RESPONSE_ONLY,
        )

        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_noneModeWithCustomSanitizer_nothingIsSanitized(self):
        processor = AutoMappedMessageProcessor(
            SimpleDto, ParameterInteractor(),
            sanitizer=Nh3Sanitizer(),
            sanitize_mode=SanitizeMode.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)

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

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


class _UpperCaseSanitizer(Sanitizer):
    def _types_to_sanitize(self):
        return [str]

    def _sanitize(self, value):
        if isinstance(value, str):
            return value.upper()
        return value


class TestSetDefaultSanitizer:
    def _restore_default(self):
        AutoMappedMessageProcessor.set_default_sanitizer(Nh3Sanitizer())

    def test_setDefaultSanitizer_newInstancesUseIt(self):
        try:
            AutoMappedMessageProcessor.set_default_sanitizer(_UpperCaseSanitizer())

            processor = AutoMappedMessageProcessor(
                SimpleDto, ParameterInteractor(),
            )

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

            assert_that(result.name).is_equal_to("HELLO")
        finally:
            self._restore_default()

    def test_setDefaultSanitizerToNone_disablesSanitizationByDefault(self):
        try:
            AutoMappedMessageProcessor.set_default_sanitizer(None)

            processor = AutoMappedMessageProcessor(
                SimpleDto, ParameterInteractor(),
            )

            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)
        finally:
            self._restore_default()

    def test_setDefaultSanitizerToNone_responseIsAlsoNotSanitized(self):
        try:
            AutoMappedMessageProcessor.set_default_sanitizer(None)

            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(XSS_SCRIPT_TAG)
        finally:
            self._restore_default()

    def test_explicitSanitizer_overridesClassDefault(self):
        try:
            AutoMappedMessageProcessor.set_default_sanitizer(None)

            processor = AutoMappedMessageProcessor(
                SimpleDto, ParameterInteractor(),
                sanitizer=Nh3Sanitizer(),
            )

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

            assert_that(result.name).is_equal_to(_sanitize(XSS_SCRIPT_TAG))
        finally:
            self._restore_default()

    def test_explicitSanitizerNone_overridesClassDefault(self):
        try:
            AutoMappedMessageProcessor.set_default_sanitizer(_UpperCaseSanitizer())

            processor = AutoMappedMessageProcessor(
                SimpleDto, ParameterInteractor(),
                sanitizer=None,
            )

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

            assert_that(result.name).is_equal_to("hello")
        finally:
            self._restore_default()

    def test_setDefaultSanitizer_doesNotAffectExistingInstances(self):
        try:
            processor = AutoMappedMessageProcessor(
                SimpleDto, ParameterInteractor(),
            )

            AutoMappedMessageProcessor.set_default_sanitizer(None)

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

            assert_that(result.name).is_equal_to(_sanitize(XSS_SCRIPT_TAG))
        finally:
            self._restore_default()

    def test_setDefaultSanitizer_worksWithSanitizeMode(self):
        try:
            AutoMappedMessageProcessor.set_default_sanitizer(_UpperCaseSanitizer())

            processor = AutoMappedMessageProcessor(
                SimpleDto, ParameterInteractor(),
                sanitize_mode=SanitizeMode.REQUEST_ONLY,
            )

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

            assert_that(result.name).is_equal_to("HELLO")

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

            body = json.loads(msg_bus.reply_message_call.args[1].body)
            assert_that(body["name"]).is_equal_to("hello")
        finally:
            self._restore_default()
