diff options
| -rw-r--r-- | libmproxy/flow.py | 31 | ||||
| -rw-r--r-- | libmproxy/proxy.py | 12 | ||||
| -rw-r--r-- | test/test_flow.py | 17 |
3 files changed, 52 insertions, 8 deletions
diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 9596d416..6031a009 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -2,9 +2,9 @@ This module provides more sophisticated flow tracking. These match requests with their responses, and provide filtering and interception facilities. """ -import subprocess, base64, sys +import subprocess, base64, sys, json from contrib import bson -import proxy, threading +import proxy, threading, netstring class RunException(Exception): def __init__(self, msg, returncode, errout): @@ -302,3 +302,30 @@ class State: rt = ReplayThread(f, masterq) rt.start() #end nocover + + + +class FlowWriter: + def __init__(self, fo): + self.fo = fo + self.ns = netstring.FileEncoder(fo) + + def add(self, flow): + d = flow.get_state() + s = json.dumps(d) + self.ns.write(s) + + +class FlowReader: + def __init__(self, fo): + self.fo = fo + self.ns = netstring.decode_file(fo) + + def stream(self): + """ + Yields Flow objects from the dump. + """ + for i in self.ns: + data = json.loads(i) + yield Flow.from_state(data) + diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py index 0941155f..ba47ce7b 100644 --- a/libmproxy/proxy.py +++ b/libmproxy/proxy.py @@ -5,7 +5,7 @@ Development started from Neil Schemenauer's munchy.py """ -import sys, os, time, string, socket, urlparse, re, select, copy +import sys, os, time, string, socket, urlparse, re, select, copy, base64 import SocketServer, ssl import utils, controller @@ -147,7 +147,7 @@ class Request(controller.Msg): method = self.method, path = self.path, headers = self.headers.get_state(), - content = self.content, + content = base64.encodestring(self.content), timestamp = self.timestamp, ) @@ -161,7 +161,7 @@ class Request(controller.Msg): state["method"], state["path"], utils.Headers.from_state(state["headers"]), - state["content"], + base64.decodestring(state["content"]), state["timestamp"] ) @@ -242,7 +242,7 @@ class Response(controller.Msg): msg = self.msg, headers = self.headers.get_state(), timestamp = self.timestamp, - content = self.content + content = base64.encodestring(self.content) ) @classmethod @@ -252,7 +252,7 @@ class Response(controller.Msg): state["code"], state["msg"], utils.Headers.from_state(state["headers"]), - state["content"], + base64.decodestring(state["content"]), state["timestamp"], ) @@ -307,7 +307,7 @@ class ClientConnection(controller.Msg): controller.Msg.__init__(self) def get_state(self): - return self.address + return list(self.address) @classmethod def from_state(klass, state): diff --git a/test/test_flow.py b/test/test_flow.py index b71ce6af..4dfa059e 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -1,3 +1,4 @@ +from cStringIO import StringIO from libmproxy import console, proxy, filt, flow import utils import libpry @@ -249,8 +250,24 @@ class uState(libpry.AutoTree): c.accept_all() +class uSerialize(libpry.AutoTree): + def test_roundtrip(self): + sio = StringIO() + f = utils.tflow() + w = flow.FlowWriter(sio) + w.add(f) + + sio.seek(0) + r = flow.FlowReader(sio) + l = list(r.stream()) + assert len(l) == 1 + assert l[0] == f + + + tests = [ uFlow(), uState(), + uSerialize() ] |