websockets: code specification

1 files changed, 22 insertions, 1 deletions

diff --git a/libpathod/language/websockets.py b/libpathod/language/websockets.py
index ae755f30..ddffdab9 100644
--- a/libpathod/language/websockets.py
+++ b/libpathod/language/websockets.py
@@ -17,6 +17,19 @@ class WF(base.CaselessLiteral):
     TOK = "wf"
 
 
+class Code(base.IntField):
+    names = {
+        "continue": netlib.websockets.OPCODE.CONTINUE,
+        "text": netlib.websockets.OPCODE.TEXT,
+        "binary": netlib.websockets.OPCODE.BINARY,
+        "close": netlib.websockets.OPCODE.CLOSE,
+        "ping": netlib.websockets.OPCODE.PING,
+        "pong": netlib.websockets.OPCODE.PONG,
+    }
+    max = 15
+    preamble = "c"
+
+
 class Body(base.Value):
     preamble = "b"
 
@@ -24,6 +37,7 @@ class Body(base.Value):
 class WebsocketFrame(message.Message):
     comps = (
         Body,
+        Code,
         actions.PauseAt,
         actions.DisconnectAt,
         actions.InjectAt
@@ -37,6 +51,10 @@ class WebsocketFrame(message.Message):
     def body(self):
         return self.tok(Body)
 
+    @property
+    def code(self):
+        return self.tok(Code)
+
     @classmethod
     def expr(klass):
         parts = [i.expr() for i in klass.comps]
@@ -59,10 +77,13 @@ class WebsocketFrame(message.Message):
         else:
             bodygen = None
             length = 0
-        frame = netlib.websockets.FrameHeader(
+        frameparts = dict(
             mask = True,
             payload_length = length
         )
+        if self.code:
+            frameparts["opcode"] = self.code.value
+        frame = netlib.websockets.FrameHeader(**frameparts)
         vals = [frame.to_bytes()]
         if self.body:
             masker = netlib.websockets.Masker(frame.masking_key)