1 files changed, 37 insertions, 51 deletions
diff --git a/libpathod/language/websockets.py b/libpathod/language/websockets.py
index 3869bd85..44c12f64 100644
--- a/libpathod/language/websockets.py
+++ b/libpathod/language/websockets.py
@@ -68,60 +68,31 @@ class Length(base.Integer):
 class Times(base.Integer):
     preamble = "x"
 
-
-class NestedFrame(base.Token):
-    def __init__(self, value):
-        self.value = value
-        try:
-            self.parsed = WebsocketFrame(
-                Response.expr().parseString(
-                    value.val,
-                    parseAll=True
-                )
-            )
-        except pp.ParseException as v:
-            raise exceptions.ParseException(v.msg, v.line, v.col)
-
-    @classmethod
-    def expr(klass):
-        e = pp.Literal("wf").suppress()
-        e = e + base.TokValueLiteral.expr()
-        return e.setParseAction(lambda x: klass(*x))
-
-    def values(self, settings):
-        return [
-            self.value.get_generator(settings),
-        ]
-
-    def spec(self):
-        return "s%s" % (self.value.spec())
-
-    def freeze(self, settings):
-        f = self.parsed.freeze(settings).spec()
-        return NestedFrame(base.TokValueLiteral(f.encode("string_escape")))
+COMPONENTS = (
+    OpCode,
+    Length,
+    # Bit flags
+    Fin,
+    RSV1,
+    RSV2,
+    RSV3,
+    Mask,
+    actions.PauseAt,
+    actions.DisconnectAt,
+    actions.InjectAt,
+    KeyNone,
+    Key,
+    Times,
+
+    Body,
+    RawBody,
+)
 
 
 class WebsocketFrame(message.Message):
-    comps = (
-        OpCode,
-        Length,
-        # Bit flags
-        Fin,
-        RSV1,
-        RSV2,
-        RSV3,
-        Mask,
-        actions.PauseAt,
-        actions.DisconnectAt,
-        actions.InjectAt,
-        KeyNone,
-        Key,
-        Times,
-
-        Body,
-        RawBody,
-    )
+    components = COMPONENTS
     logattrs = ["body"]
+    unique_name = "body"
 
     @property
     def actions(self):
@@ -177,7 +148,7 @@ class WebsocketFrame(message.Message):
 
     @classmethod
     def expr(klass):
-        parts = [i.expr() for i in klass.comps]
+        parts = [i.expr() for i in klass.components]
         atom = pp.MatchFirst(parts)
         resp = pp.And(
             [
@@ -246,3 +217,18 @@ class WebsocketFrame(message.Message):
 
     def spec(self):
         return ":".join([i.spec() for i in self.tokens])
+
+
+class NestedFrame(base.NestedMessage):
+    preamble = "f"
+    nest_type = WebsocketFrame
+
+
+class WebsocketClientFrame(WebsocketFrame):
+    components = COMPONENTS + (
+        NestedFrame,
+    )
+
+    @property
+    def nested_frame(self):
+        return self.tok(NestedFrame)