mitmproxy/libpathod/rparse.py

import operator, string, random, mmap, os, time
import contrib.pyparsing as pp
from netlib import http_status
import utils

BLOCKSIZE = 1024
TRUNCATE = 1024

class ParseException(Exception):
    def __init__(self, msg, s, col):
        Exception.__init__(self)
        self.msg = msg
        self.s = s
        self.col = col

    def marked(self):
        return "%s\n%s"%(self.s, " "*(self.col-1) + "^")

    def __str__(self):
        return self.msg


class ServerError(Exception): pass


def ready_actions(length, lst):
    ret = []
    for i in lst:
        itms = list(i)
        if i[0] == "r":
            itms[0] = random.randrange(length)
        if i[0] == "a":
            itms[0] = length+1
        ret.append(tuple(itms))
    ret.sort()
    return ret


def write_values(fp, vals, actions, sofar=0, skip=0, blocksize=BLOCKSIZE):
    """
        vals: A list of values, which may be strings or Value objects.
        actions: A list of (offset, action, arg) tuples. Action may be "pause" or "disconnect".

        Return True if connection should disconnect.
    """
    while vals:
        part = vals.pop()
        for i in range(skip, len(part), blocksize):
            d = part[i:i+blocksize]
            if actions and actions[-1][0] < (sofar + len(d)):
                p = actions.pop()
                offset = p[0]-sofar
                vals.append(part)
                if p[1] == "pause":
                    fp.write(d[:offset])
                    time.sleep(p[2])
                    return write_values(
                        fp, vals, actions,
                        sofar=sofar+offset,
                        skip=i+offset,
                        blocksize=blocksize
                    )
                elif p[1] == "disconnect":
                    fp.write(d[:offset])
                    return True
            fp.write(d)
            sofar += len(d)
        skip = 0


DATATYPES = dict(
    ascii_letters = string.ascii_letters,
    ascii_lowercase = string.ascii_lowercase,
    ascii_uppercase = string.ascii_uppercase,
    digits = string.digits,
    hexdigits = string.hexdigits,
    letters = string.letters,
    lowercase = string.lowercase,
    octdigits = string.octdigits,
    printable = string.printable,
    punctuation = string.punctuation,
    uppercase = string.uppercase,
    whitespace = string.whitespace,
    ascii = string.printable,
    bytes = "".join(chr(i) for i in range(256))
)


v_integer = pp.Regex(r"[+-]?\d+")\
    .setName("integer")\
    .setParseAction(lambda toks: int(toks[0]))


v_literal = pp.MatchFirst(
    [
        pp.QuotedString("\"", escChar="\\", unquoteResults=True),
        pp.QuotedString("'", escChar="\\", unquoteResults=True),
    ]
)

v_naked_literal = pp.MatchFirst(
    [
        v_literal,
        pp.Word("".join(i for i in pp.printables if i not in ",:"))
    ]
)


class LiteralGenerator:
    def __init__(self, s):
        self.s = s

    def __eq__(self, other):
        return self[:] == other

    def __len__(self):
        return len(self.s)

    def __getitem__(self, x):
        return self.s.__getitem__(x)

    def __getslice__(self, a, b):
        return self.s.__getslice__(a, b)


class RandomGenerator:
    def __init__(self, chars, length):
        self.chars = chars
        self.length = length

    def __len__(self):
        return self.length

    def __getitem__(self, x):
        return random.choice(self.chars)

    def __getslice__(self, a, b):
        b = min(b, self.length)
        return "".join(random.choice(self.chars) for x in range(a, b))


class FileGenerator:
    def __init__(self, path):
        self.path = path
        self.fp = file(path, "r")
        self.map = mmap.mmap(self.fp.fileno(), 0, prot=mmap.PROT_READ)

    def __len__(self):
        return len(self.map)

    def __getitem__(self, x):
        return self.map.__getitem__(x)

    def __getslice__(self, a, b):
        return self.map.__getslice__(a, b)


class ValueLiteral:
    def __init__(self, val):
        self.val = val

    def get_generator(self, settings):
        return LiteralGenerator(self.val)

    @classmethod
    def expr(klass):
        e = v_literal.copy()
        return e.setParseAction(lambda x: klass(*x))

    def __str__(self):
        return self.val


class ValueGenerate:
    UNITS = dict(
        b = 1024**0,
        k = 1024**1,
        m = 1024**2,
        g = 1024**3,
        t = 1024**4,
    )
    def __init__(self, usize, unit, datatype):
        if not unit:
            unit = "b"
        self.usize, self.unit, self.datatype = usize, unit, datatype

    def bytes(self):
        return self.usize * self.UNITS[self.unit]

    def get_generator(self, settings):
        return RandomGenerator(DATATYPES[self.datatype], self.bytes())

    @classmethod
    def expr(klass):
        e = pp.Literal("@").suppress() + v_integer

        u = reduce(operator.or_, [pp.Literal(i) for i in klass.UNITS.keys()])
        e = e + pp.Optional(u, default=None)

        s = pp.Literal(",").suppress()
        s += reduce(operator.or_, [pp.Literal(i) for i in DATATYPES.keys()])
        e += pp.Optional(s, default="bytes")
        return e.setParseAction(lambda x: klass(*x))

    def __str__(self):
        return "@%s%s,%s"%(self.usize, self.unit, self.datatype)


class ValueFile:
    def __init__(self, path):
        self.path = path

    @classmethod
    def expr(klass):
        e = pp.Literal("<").suppress()
        e = e + v_naked_literal
        return e.setParseAction(lambda x: klass(*x))

    def get_generator(self, settings):
        sd = settings.get("staticdir")
        if not sd:
            raise ServerError("No static directory specified.")
        path = os.path.join(sd, self.path)
        if not os.path.exists(path):
            raise ServerError("Static file does not exist: %s"%path)
        return FileGenerator(path)

    def __str__(self):
        return "<%s"%(self.path)


Value = pp.MatchFirst(
    [
        ValueGenerate.expr(),
        ValueFile.expr(),
        ValueLiteral.expr()
    ]
)


class ShortcutContentType:
    def __init__(self, value):
        self.value = value

    def accept(self, settings, r):
        r.headers.append(
            (
                LiteralGenerator("Content-Type"),
                self.value.get_generator(settings)
            )
        )

    @classmethod
    def expr(klass):
        e = pp.Literal("c").suppress()
        e = e + Value
        return e.setParseAction(lambda x: klass(*x))



class ShortcutLocation:
    def __init__(self, value):
        self.value = value

    def accept(self, settings, r):
        r.headers.append(
            (
                LiteralGenerator("Location"),
                self.value.get_generator(settings)
            )
        )

    @classmethod
    def expr(klass):
        e = pp.Literal("l").suppress()
        e = e + Value
        return e.setParseAction(lambda x: klass(*x))


class Body:
    def __init__(self, value):
        self.value = value

    def accept(self, settings, r):
        r.body = self.value.get_generator(settings)

    @classmethod
    def expr(klass):
        e = pp.Literal("b").suppress()
        e = e + Value
        return e.setParseAction(lambda x: klass(*x))


class Path:
    def __init__(self, value):
        if isinstance(value, basestring):
            value = ValueLiteral(value)
        self.value = value

    def accept(self, settings, r):
        r.path = self.value.get_generator(settings)

    @classmethod
    def expr(klass):
        e = v_naked_literal.copy()
        return e.setParseAction(lambda x: klass(*x))



class Method:
    methods = [
        "get",
        "head",
        "post",
        "put",
        "delete",
        "options",
        "trace",
        "connect",
    ]
    def __init__(self, value):
        # If it's a string, we were passed one of the methods, so we upper-case
        # it to be canonical. The user can specify a different case by using a
        # string value literal.
        if isinstance(value, basestring):
            value = ValueLiteral(value.upper())
        self.value = value

    def accept(self, settings, r):
        r.method = self.value.get_generator(settings)

    @classmethod
    def expr(klass):
        parts = [pp.CaselessLiteral(i) for i in klass.methods]
        m = pp.MatchFirst(parts)
        spec = m | Value.copy()
        spec = spec.setParseAction(lambda x: klass(*x))
        return spec


class PauseAt:
    def __init__(self, seconds, offset):
        self.seconds, self.offset = seconds, offset

    @classmethod
    def expr(klass):
        e = pp.Literal("p").suppress()
        e += pp.MatchFirst(
                    [
                        v_integer,
                        pp.Literal("f")
                    ]
                )
        e += pp.Literal(",").suppress()
        e += pp.MatchFirst(
                    [
                        v_integer,
                        pp.Literal("r"),
                        pp.Literal("a"),
                    ]
                )
        return e.setParseAction(lambda x: klass(*x))

    def accept(self, settings, r):
        r.actions.append((self.offset, "pause", self.seconds))


class DisconnectAt:
    def __init__(self, value):
        self.value = value

    def accept(self, settings, r):
        r.actions.append((self.value, "disconnect"))

    @classmethod
    def expr(klass):
        e = pp.Literal("d").suppress()
        e = e + pp.MatchFirst(
                    [
                        v_integer,
                        pp.Literal("r")
                    ]
                )
        return e.setParseAction(lambda x: klass(*x))


class Header:
    def __init__(self, key, value):
        self.key, self.value = key, value

    def accept(self, settings, r):
        r.headers.append(
            (
                self.key.get_generator(settings),
                self.value.get_generator(settings)
            )
        )

    @classmethod
    def expr(klass):
        e = pp.Literal("h").suppress()
        e += Value
        e += pp.Literal("=").suppress()
        e += Value
        return e.setParseAction(lambda x: klass(*x))


class Code:
    def __init__(self, code, msg=None):
        self.code, self.msg = code, msg
        if msg is None:
            self.msg = ValueLiteral(http_status.RESPONSES.get(self.code, "Unknown code"))

    def accept(self, settings, r):
        r.code = self.code
        r.msg = self.msg.get_generator(settings)

    @classmethod
    def expr(klass):
        e = v_integer
        e = e + pp.Optional(
            Value
        )
        return e.setParseAction(lambda x: klass(*x))



class Message:
    version = "HTTP/1.1"
    def length(self):
        l = sum(len(x) for x in self.preamble())
        l += 2
        for i in self.headers:
            l += len(i[0]) + len(i[1])
            l += 4
        l += 2
        l += len(self.body)
        return l

    def serve(self, fp):
        started = time.time()
        if self.body and not utils.get_header("Content-Length", self.headers):
            self.headers.append(
                (
                    LiteralGenerator("Content-Length"),
                    LiteralGenerator(str(len(self.body))),
                )
            )

        hdrs = []
        for k, v in self.headers:
            hdrs.extend([
                k,
                ": ",
                v,
                "\r\n",
            ])
        vals = self.preamble()
        vals.append("\r\n")
        vals.extend(hdrs)
        vals.append("\r\n")
        if self.body:
            vals.append(self.body)
        vals.reverse()
        actions = ready_actions(self.length(), self.actions)
        actions.reverse()
        disconnect = write_values(fp, vals, actions[:])
        duration = time.time() - started
        ret = dict(
            disconnect = disconnect,
            started = started,
            duration = duration,
            actions = actions,
        )
        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, "__len__"):
                v = v[:TRUNCATE]
            ret[i] = v
        return ret


class Response(Message):
    comps = (
        Body,
        Header,
        PauseAt,
        DisconnectAt,
        ShortcutContentType,
        ShortcutLocation,
    )
    logattrs = ["code", "version"]
    def __init__(self):
        self.headers = []
        self.actions = []
        self.code = 200
        self.msg = LiteralGenerator(http_status.RESPONSES[self.code])
        self.body = LiteralGenerator("")

    def preamble(self):
        return [self.version, " ", str(self.code), " ", self.msg]

    @classmethod
    def expr(klass):
        parts = [i.expr() for i in klass.comps]
        atom = pp.MatchFirst(parts)
        resp = pp.And(
            [
                Code.expr(),
                pp.ZeroOrMore(pp.Literal(":").suppress() + atom)
            ]
        )
        return resp

    def __str__(self):
        parts = [
            "%s %s"%(self.code, self.msg[:])
        ]
        return "\n".join(parts)


class Request(Message):
    comps = (
        Body,
        Header,
        PauseAt,
        DisconnectAt,
        ShortcutContentType,
    )
    logattrs = ["method", "path"]
    def __init__(self):
        self.method = None
        self.path = None
        self.body = LiteralGenerator("")
        self.headers = []
        self.actions = []

    def preamble(self):
        return [self.method, " ", self.path, " ", self.version]

    @classmethod
    def expr(klass):
        parts = [i.expr() for i in klass.comps]
        atom = pp.MatchFirst(parts)
        resp = pp.And(
            [
                Method.expr(),
                pp.Literal(":").suppress(),
                Path.expr(),
                pp.ZeroOrMore(pp.Literal(":").suppress() + atom)
            ]
        )
        return resp

    def __str__(self):
        parts = [
            "%s %s"%(self.method[:], self.path[:])
        ]
        return "\n".join(parts)


class CraftedRequest(Request):
    def __init__(self, settings, spec, tokens):
        Request.__init__(self)
        self.spec, self.tokens = spec, tokens
        for i in tokens:
            i.accept(settings, self)

    def serve(self, fp):
        d = Request.serve(self, fp)
        d["spec"] = self.spec
        return d


class CraftedResponse(Response):
    def __init__(self, settings, spec, tokens):
        Response.__init__(self)
        self.spec, self.tokens = spec, tokens
        for i in tokens:
            i.accept(settings, self)

    def serve(self, fp):
        d = Response.serve(self, fp)
        d["spec"] = self.spec
        return d


class InternalResponse(Response):
    def __init__(self, code, body):
        Response.__init__(self)
        self.code = code
        self.msg = LiteralGenerator(http_status.RESPONSES.get(code, "Unknown error"))
        self.body = LiteralGenerator(body)
        self.headers = [
            (
                LiteralGenerator("Content-Type"),
                LiteralGenerator("text/plain")
            ),
            (
                LiteralGenerator("Content-Length"),
                LiteralGenerator(str(len(self.body)))
            )
        ]

    def serve(self, fp):
        d = Response.serve(self, fp)
        d["internal"] = True
        return d


def parse_response(settings, s):
    try:
        return CraftedResponse(settings, s, Response.expr().parseString(s, parseAll=True))
    except pp.ParseException, v:
        raise ParseException(v.msg, v.line, v.col)


def parse_request(settings, s):
    try:
        return CraftedRequest(settings, s, Request.expr().parseString(s, parseAll=True))
    except pp.ParseException, v:
        raise ParseException(v.msg, v.line, v.col)