mitmproxy/pathod/language/message.py

import abc
import typing  # noqa

import pyparsing as pp

from mitmproxy.utils import strutils
from . import actions, exceptions, base

LOG_TRUNCATE = 1024


class Message:
    __metaclass__ = abc.ABCMeta
    logattrs: typing.List[str] = []

    def __init__(self, tokens):
        track = set()
        for i in tokens:
            if i.unique_name:
                if i.unique_name in track:
                    raise exceptions.ParseException(
                        "Message has multiple %s clauses, "
                        "but should only have one." % i.unique_name,
                        0, 0
                    )
                else:
                    track.add(i.unique_name)
        self.tokens = tokens

    def strike_token(self, name):
        toks = [i for i in self.tokens if i.unique_name != name]
        return self.__class__(toks)

    def toks(self, klass):
        """
            Fetch all tokens that are instances of klass
        """
        return [i for i in self.tokens if isinstance(i, klass)]

    def tok(self, klass):
        """
            Fetch first token that is an instance of klass
        """
        l = self.toks(klass)
        if l:
            return l[0]

    def length(self, settings):
        """
            Calculate the length of the base message without any applied
            actions.
        """
        return sum(len(x) for x in self.values(settings))

    def preview_safe(self):
        """
            Return a copy of this message that is safe for previews.
        """
        tokens = [i for i in self.tokens if not isinstance(i, actions.PauseAt)]
        return self.__class__(tokens)

    def maximum_length(self, settings):
        """
            Calculate the maximum length of the base message with all applied
            actions.
        """
        l = self.length(settings)
        for i in self.actions:
            if isinstance(i, actions.InjectAt):
                l += len(i.value.get_generator(settings))
        return l

    @classmethod
    def expr(cls):  # pragma: no cover
        pass

    def log(self, settings):
        """
            A dictionary that should be logged if this message is served.
        """
        ret = {}
        for i in self.logattrs:
            v = getattr(self, i)
            # Careful not to log any VALUE specs without sanitizing them first.
            # We truncate at 1k.
            if hasattr(v, "values"):
                v = [x[:LOG_TRUNCATE] for x in v.values(settings)]
                v = strutils.bytes_to_escaped_str(b"".join(v))
            elif hasattr(v, "__len__"):
                v = v[:LOG_TRUNCATE]
                v = strutils.bytes_to_escaped_str(v)
            ret[i] = v
        ret["spec"] = self.spec()
        return ret

    def freeze(self, settings):
        r = self.resolve(settings)
        return self.__class__([i.freeze(settings) for i in r.tokens])

    def __repr__(self):
        return self.spec()


class NestedMessage(base.Token):
    """
        A nested message, as an escaped string with a preamble.
    """
    preamble = ""
    nest_type: typing.Optional[typing.Type[Message]] = None

    def __init__(self, value):
        super().__init__()
        self.value = value
        try:
            self.parsed = self.nest_type(
                self.nest_type.expr().parseString(
                    value.val.decode(),
                    parseAll=True
                )
            )
        except pp.ParseException as v:
            raise exceptions.ParseException(v.msg, v.line, v.col)

    @classmethod
    def expr(cls):
        e = pp.Literal(cls.preamble).suppress()
        e = e + base.TokValueLiteral.expr()
        return e.setParseAction(lambda x: cls(*x))

    def values(self, settings):
        return [
            self.value.get_generator(settings),
        ]

    def spec(self):
        return f"{self.preamble}{self.value.spec()}"

    def freeze(self, settings):
        f = self.parsed.freeze(settings).spec()
        return self.__class__(
            base.TokValueLiteral(
                strutils.bytes_to_escaped_str(f.encode(), escape_single_quotes=True)
            )
        )