import struct
import sys

from hypothesis import settings
from hypothesis.stateful import initialize, precondition, rule
import hypothesis.strategies as st

from llvm_one_input import C
import sanitizers


class GrammarTester(sanitizers.RuleBasedStateMachine):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        with sanitizers.leaky_region():
            assert C.LLVMFuzzerTestOneInput(b"", 0) == 0

        self.last_len = 0
        self.buf = b""

    def teardown(self):
        print("CHECK")
        sys.stdout.flush()
        assert C.LLVMFuzzerTestOneInput(self.buf, len(self.buf)) == 0

    @initialize(
        tls_ver=st.sampled_from([1, 2, 3]),
        payload=st.binary(),
        padding=st.binary(min_size=16),
    )
    def add_heartbeat(self, tls_ver, payload, padding):
        hb_payload = bytes([0x01]) + struct.pack(">H", len(payload)) + payload + padding
        self.buf += bytes([0x18, 0x03, tls_ver])
        self.buf += struct.pack(">H", len(hb_payload))
        self.buf += hb_payload

    @precondition(lambda self: self.buf)
    @rule(data=st.data(), value=st.integers(min_value=0, max_value=255))
    def replace_byte(self, data, value):
        index = data.draw(st.integers(min_value=0, max_value=len(self.buf) - 1))
        prefix = self.buf[0:index]
        suffix = self.buf[index + 1 :]
        self.buf = prefix + bytes([value]) + suffix

    @precondition(lambda self: self.buf)
    @rule(data=st.data())
    def strip_suffix(self, data):
        count = data.draw(st.integers(min_value=1, max_value=len(self.buf)))
        self.buf = self.buf[0:-count]

    @rule(suffix=st.binary(min_size=1))
    def add_suffix(self, suffix):
        self.buf += suffix


TestWithGrammar = GrammarTester.TestCase
TestWithGrammar.settings = settings(max_examples=10000, deadline=None)