diff options
| -rw-r--r-- | libmproxy/dump.py | 3 | ||||
| -rw-r--r-- | libmproxy/flow.py | 13 | ||||
| -rw-r--r-- | test/test_flow.py | 10 |
3 files changed, 21 insertions, 5 deletions
diff --git a/libmproxy/dump.py b/libmproxy/dump.py index 2eb0035e..97232ac7 100644 --- a/libmproxy/dump.py +++ b/libmproxy/dump.py @@ -93,7 +93,7 @@ class DumpMaster(flow.FlowMaster): try: f = file(path, "r") flows = list(flow.FlowReader(f).stream()) - except IOError, v: + except (IOError, flow.FlowReadError), v: raise DumpError(v.strerror) return flows @@ -182,7 +182,6 @@ class DumpMaster(flow.FlowMaster): return f - # begin nocover def run(self): try: diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 60a6ebd2..d555c926 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -560,6 +560,10 @@ class FlowWriter: s = json.dumps(d) self.ns.write(s) +class FlowReadError(Exception): + @property + def strerror(self): + return self.args[0] class FlowReader: def __init__(self, fo): @@ -570,7 +574,10 @@ class FlowReader: """ Yields Flow objects from the dump. """ - for i in self.ns: - data = json.loads(i) - yield Flow.from_state(data) + try: + for i in self.ns: + data = json.loads(i) + yield Flow.from_state(data) + except netstring.DecoderError: + raise FlowReadError("Invalid data format.") diff --git a/test/test_flow.py b/test/test_flow.py index 78892b5d..d2cb85dc 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -356,6 +356,16 @@ class uSerialize(libpry.AutoTree): assert len(l) == 1 assert l[0] == f + def test_error(self): + sio = StringIO() + sio.write("bogus") + sio.seek(0) + r = flow.FlowReader(sio) + libpry.raises(flow.FlowReadError, list, r.stream()) + + f = flow.FlowReadError("foo") + assert f.strerror == "foo" + class uFlowMaster(libpry.AutoTree): def test_all(self): |