aboutsummaryrefslogtreecommitdiffstats
path: root/libmproxy/dump.py
blob: d43da44bcad103edbde3f3384d34faf71ef5bd22 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
import sys, os, traceback
import flow, filt, utils

class DumpError(Exception): pass


class Options(object):
    __slots__ = [
        "verbosity",
        "wfile",
        "request_script",
        "response_script",
        "replay",
    ]
    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            setattr(self, k, v)
        for i in self.__slots__:
            if not hasattr(self, i):
                setattr(self, i, None)


class DumpMaster(flow.FlowMaster):
    def __init__(self, server, options, filtstr, outfile=sys.stdout):
        flow.FlowMaster.__init__(self, server, flow.State())
        self.outfile = outfile
        self.o = options

        if filtstr:
            self.filt = filt.parse(filtstr)
        else:
            self.filt = None

        if options.wfile:
            path = os.path.expanduser(options.wfile)
            try:
                f = file(path, "wb")
                self.fwriter = flow.FlowWriter(f)
            except IOError, v:
                raise DumpError(v.strerror)

        if options.replay:
            path = os.path.expanduser(options.replay)
            try:
                f = file(path, "r")
                flows = list(flow.FlowReader(f).stream())
            except IOError, v:
                raise DumpError(v.strerror)
            self.start_playback(flows)

    def _runscript(self, f, script):
        try:
            ret = f.run_script(script)
            if self.o.verbosity > 0:
                print >> self.outfile, ret
        except flow.RunException, e:
            if e.errout:
                eout = "Script output:\n" + self.indent(4, e.errout) + "\n"
            else:
                eout = ""
            raise DumpError(
                    "%s: %s\n%s"%(script, e.args[0], eout)
                )

    def handle_request(self, r):
        f = flow.FlowMaster.handle_request(self, r)
        if self.o.request_script:
            self._runscript(f, self.o.request_script)
        if not self.playback(f):
            r.ack()

    def indent(self, n, t):
        l = str(t).strip().split("\n")
        return "\n".join(" "*n + i for i in l)

    def handle_response(self, msg):
        f = flow.FlowMaster.handle_response(self, msg)
        if f:
            if self.o.response_script:
                self._runscript(f, self.o.response_script)
            msg.ack()
            if self.filt and not f.match(self.filt):
                    return
            sz = utils.pretty_size(len(f.response.content))
            if self.o.verbosity == 1:
                print >> self.outfile, f.request.client_conn.address[0],
                print >> self.outfile, f.request.short()
                print >> self.outfile, "  <<",
                print >> self.outfile, f.response.short(), sz
            elif self.o.verbosity == 2:
                print >> self.outfile, f.request.client_conn.address[0],
                print >> self.outfile, f.request.short()
                print >> self.outfile, self.indent(4, f.request.headers)
                print >> self.outfile
                print >> self.outfile, " <<", f.response.short(), sz
                print >> self.outfile, self.indent(4, f.response.headers)
                print >> self.outfile, "\n"
            elif self.o.verbosity == 3:
                print >> self.outfile, f.request.client_conn.address[0],
                print >> self.outfile, f.request.short()
                print >> self.outfile, self.indent(4, f.request.headers)
                if utils.isBin(f.request.content):
                    print >> self.outfile, self.indent(4, utils.hexdump(f.request.content))
                elif f.request.content:
                    print >> self.outfile, self.indent(4, f.request.content)
                print >> self.outfile
                print >> self.outfile, " <<", f.response.short(), sz
                print >> self.outfile, self.indent(4, f.response.headers)
                if utils.isBin(f.response.content):
                    print >> self.outfile, self.indent(4, utils.hexdump(f.response.content))
                elif f.response.content:
                    print >> self.outfile, self.indent(4, f.response.content)
                print >> self.outfile, "\n"
            self.state.delete_flow(f)
            if self.o.wfile:
                self.fwriter.add(f)

# begin nocover
    def run(self):
        try:
            return flow.FlowMaster.run(self)
        except BaseException, v:
            self.shutdown()
            raise