A demo of the peels library. Parse a websocket frame in < 50 lines of code!

example.py

from io import BytesIO
from peels import Peel, SharedConsumption, Segment, ConstantConsumption, Condition, ValueCondition


class PayloadLengthContConsumer(ConstantConsumption):

    def length(self, ctx, _):
        length = ctx["payload_len"]
        if length == 126:
            return 2
        else:
            return 8


class PayloadLengthContCondition(Condition):

    def __call__(self, ctx, _):
        return ctx[self.target] > 126


class PayloadConsumer(ConstantConsumption):

    def length(self, ctx, _):
        if "payload_len_cont" in ctx:
            return ctx["payload_len_cont"]
        return ctx["payload_len"]


header_flags = SharedConsumption(1, {"fin": 128, "rsv1": 64, "rsv2": 32, "rsv3": 16, "opcode": 0xF})
mask_payload = SharedConsumption(1, {"has_mask": 128, "payload_len": 127})

websocket_frame = Peel(
    Segment("fin", header_flags, "B"),
    Segment("rsv1", header_flags, "B"),
    Segment("rsv2", header_flags, "B"),
    Segment("rsv3", header_flags, "B"),
    Segment("opcode", header_flags, "-"),
    Segment("has_mask", mask_payload, "B"),
    Segment("payload_len", mask_payload, "-"),
    Segment("payload_len_cont", PayloadLengthContConsumer(None), "I", PayloadLengthContCondition("payload_len")),
    Segment("mask_key", 2, "Y", ValueCondition("has_mask")),
    Segment("payload", PayloadConsumer(None), "Y")
)

stream = BytesIO(b'\x81\x0cHello World!')
stream.seek(0)
result = websocket_frame.parse_stream(stream)
print(result)
# OrderedDict([('fin', True), ('rsv1', False), ('rsv2', False), ('rsv3', False), ('opcode', 1), ('has_mask', False), ('payload_len', 12), ('payload', b'Hello World!')])