aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
Diffstat
-rw-r--r--mitmproxy/addons/__init__.py2
-rw-r--r--mitmproxy/addons/readfile.py50
-rw-r--r--mitmproxy/addons/readstdin.py34
-rw-r--r--mitmproxy/master.py30
-rw-r--r--mitmproxy/tools/console/master.py20
-rw-r--r--mitmproxy/tools/dump.py17
-rw-r--r--mitmproxy/tools/web/app.py3
-rw-r--r--mitmproxy/tools/web/master.py9
-rw-r--r--test/mitmproxy/addons/test_readfile.py62
-rw-r--r--test/mitmproxy/addons/test_readstdin.py59
-rw-r--r--test/mitmproxy/test_flow.py53
-rw-r--r--test/mitmproxy/tools/console/test_master.py6
-rw-r--r--test/mitmproxy/tools/test_dump.py13
13 files changed, 235 insertions, 123 deletions
diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py
index 7a45106c..b4367d78 100644
--- a/mitmproxy/addons/__init__.py
+++ b/mitmproxy/addons/__init__.py
@@ -8,6 +8,7 @@ from mitmproxy.addons import disable_h2c
from mitmproxy.addons import onboarding
from mitmproxy.addons import proxyauth
from mitmproxy.addons import replace
+from mitmproxy.addons import readfile
from mitmproxy.addons import script
from mitmproxy.addons import serverplayback
from mitmproxy.addons import setheaders
@@ -37,5 +38,6 @@ def default_addons():
stickycookie.StickyCookie(),
streambodies.StreamBodies(),
streamfile.StreamFile(),
+ readfile.ReadFile(),
upstream_auth.UpstreamAuth(),
]
diff --git a/mitmproxy/addons/readfile.py b/mitmproxy/addons/readfile.py
new file mode 100644
index 00000000..a4b92444
--- /dev/null
+++ b/mitmproxy/addons/readfile.py
@@ -0,0 +1,50 @@
+import os.path
+
+from mitmproxy import ctx
+from mitmproxy import io
+from mitmproxy import exceptions
+
+
+class ReadFile:
+ """
+ An addon that handles reading from file on startup.
+ """
+ def __init__(self):
+ self.path = None
+ self.keepserving = False
+
+ def load_flows_file(self, path: str) -> int:
+ path = os.path.expanduser(path)
+ cnt = 0
+ try:
+ with open(path, "rb") as f:
+ freader = io.FlowReader(f)
+ for i in freader.stream():
+ cnt += 1
+ ctx.master.load_flow(i)
+ return cnt
+ except (IOError, exceptions.FlowReadException) as v:
+ if cnt:
+ ctx.log.warn(
+ "Flow file corrupted - loaded %i flows." % cnt,
+ )
+ else:
+ ctx.log.error("Flow file corrupted.")
+ raise exceptions.FlowReadException(v)
+
+ def configure(self, options, updated):
+ if "keepserving" in updated:
+ self.keepserving = options.keepserving
+ if "rfile" in updated and options.rfile:
+ self.path = options.rfile
+
+ def running(self):
+ if self.path:
+ try:
+ self.load_flows_file(self.path)
+ except exceptions.FlowReadException as v:
+ raise exceptions.OptionsError(v)
+ finally:
+ self.path = None
+ if not self.keepserving:
+ ctx.master.shutdown()
diff --git a/mitmproxy/addons/readstdin.py b/mitmproxy/addons/readstdin.py
new file mode 100644
index 00000000..e45d25b8
--- /dev/null
+++ b/mitmproxy/addons/readstdin.py
@@ -0,0 +1,34 @@
+from mitmproxy import ctx
+from mitmproxy import io
+from mitmproxy import exceptions
+import sys
+
+
+class ReadStdin:
+ """
+ An addon that reads from stdin if we're not attached to (someting like)
+ a tty.
+ """
+ def __init__(self):
+ self.keepserving = False
+
+ def configure(self, options, updated):
+ if "keepserving" in updated:
+ self.keepserving = options.keepserving
+
+ def running(self, stdin = sys.stdin):
+ if not stdin.isatty():
+ ctx.log.info("Reading from stdin")
+ try:
+ stdin.buffer.read(0)
+ except Exception as e:
+ ctx.log.warn("Cannot read from stdin: {}".format(e))
+ return
+ freader = io.FlowReader(stdin.buffer)
+ try:
+ for i in freader.stream():
+ ctx.master.load_flow(i)
+ except exceptions.FlowReadException as e:
+ ctx.log.error("Error reading from stdin: %s" % e)
+ if not self.keepserving:
+ ctx.master.shutdown()
diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index 79747a97..69359de6 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -1,8 +1,6 @@
-import os
import threading
import contextlib
import queue
-import sys
from mitmproxy import addonmanager
from mitmproxy import options
@@ -12,7 +10,6 @@ from mitmproxy import exceptions
from mitmproxy import connections
from mitmproxy import http
from mitmproxy import log
-from mitmproxy import io
from mitmproxy.proxy.protocol import http_replay
from mitmproxy.types import basethread
import mitmproxy.net.http
@@ -160,33 +157,6 @@ class Master:
for e, o in eventsequence.iterate(f):
getattr(self, e)(o)
- def load_flows(self, fr: io.FlowReader) -> int:
- """
- Load flows from a FlowReader object.
- """
- cnt = 0
- for i in fr.stream():
- cnt += 1
- self.load_flow(i)
- return cnt
-
- def load_flows_file(self, path: str) -> int:
- path = os.path.expanduser(path)
- try:
- if path == "-":
- try:
- sys.stdin.buffer.read(0)
- except Exception as e:
- raise IOError("Cannot read from stdin: {}".format(e))
- freader = io.FlowReader(sys.stdin.buffer)
- return self.load_flows(freader)
- else:
- with open(path, "rb") as f:
- freader = io.FlowReader(f)
- return self.load_flows(freader)
- except IOError as v:
- raise exceptions.FlowReadException(v.strerror)
-
def replay_request(
self,
f: http.HTTPFlow,
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index e75105cf..b6339817 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -256,19 +256,6 @@ class ConsoleMaster(master.Master):
)
self.ab = statusbar.ActionBar()
- if self.options.rfile:
- ret = self.load_flows_path(self.options.rfile)
- if ret and self.view.store_count():
- signals.add_log(
- "File truncated or corrupted. "
- "Loaded as many flows as possible.",
- "error"
- )
- elif ret and not self.view.store_count():
- self.shutdown()
- print("Could not load file: {}".format(ret), file=sys.stderr)
- sys.exit(1)
-
self.loop.set_alarm_in(0.01, self.ticker)
self.loop.set_alarm_in(
@@ -289,7 +276,10 @@ class ConsoleMaster(master.Master):
print("Shutting down...", file=sys.stderr)
finally:
sys.stderr.flush()
- self.shutdown()
+ super().shutdown()
+
+ def shutdown(self):
+ raise urwid.ExitMainLoop
def view_help(self, helpctx):
signals.push_view_state.send(
@@ -402,7 +392,7 @@ class ConsoleMaster(master.Master):
def quit(self, a):
if a != "n":
- raise urwid.ExitMainLoop
+ self.shutdown()
def clear_events(self):
self.logbuffer[:] = []
diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py
index 4bfe2dc4..be83fb1d 100644
--- a/mitmproxy/tools/dump.py
+++ b/mitmproxy/tools/dump.py
@@ -1,9 +1,8 @@
from mitmproxy import controller
-from mitmproxy import exceptions
from mitmproxy import addons
from mitmproxy import options
from mitmproxy import master
-from mitmproxy.addons import dumper, termlog, termstatus
+from mitmproxy.addons import dumper, termlog, termstatus, readstdin
class DumpMaster(master.Master):
@@ -22,21 +21,9 @@ class DumpMaster(master.Master):
self.addons.add(*addons.default_addons())
if with_dumper:
self.addons.add(dumper.Dumper())
-
- if options.rfile:
- try:
- self.load_flows_file(options.rfile)
- except exceptions.FlowReadException as v:
- self.add_log("Flow file corrupted.", "error")
- raise exceptions.OptionsError(v)
+ self.addons.add(readstdin.ReadStdin())
@controller.handler
def log(self, e):
if e.level == "error":
self.has_errored = True
-
- def run(self): # pragma: no cover
- if self.options.rfile and not self.options.keepserving:
- self.addons.done()
- return
- super().run()
diff --git a/mitmproxy/tools/web/app.py b/mitmproxy/tools/web/app.py
index eddaa3e1..002513b9 100644
--- a/mitmproxy/tools/web/app.py
+++ b/mitmproxy/tools/web/app.py
@@ -230,7 +230,8 @@ class DumpFlows(RequestHandler):
def post(self):
self.view.clear()
bio = BytesIO(self.filecontents)
- self.master.load_flows(io.FlowReader(bio))
+ for i in io.FlowReader(bio).stream():
+ self.master.load_flow(i)
bio.close()
diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py
index 8c7f579d..e28bd002 100644
--- a/mitmproxy/tools/web/master.py
+++ b/mitmproxy/tools/web/master.py
@@ -3,7 +3,6 @@ import webbrowser
import tornado.httpserver
import tornado.ioloop
from mitmproxy import addons
-from mitmproxy import exceptions
from mitmproxy import log
from mitmproxy import master
from mitmproxy.addons import eventstore
@@ -42,14 +41,6 @@ class WebMaster(master.Master):
)
# This line is just for type hinting
self.options = self.options # type: Options
- if options.rfile:
- try:
- self.load_flows_file(options.rfile)
- except exceptions.FlowReadException as v:
- self.add_log(
- "Could not read flow file: %s" % v,
- "error"
- )
def _sig_view_add(self, view, flow):
app.ClientConnection.broadcast(
diff --git a/test/mitmproxy/addons/test_readfile.py b/test/mitmproxy/addons/test_readfile.py
new file mode 100644
index 00000000..c0cf97ae
--- /dev/null
+++ b/test/mitmproxy/addons/test_readfile.py
@@ -0,0 +1,62 @@
+from mitmproxy.addons import readfile
+from mitmproxy.test import taddons
+from mitmproxy.test import tflow
+from mitmproxy import io
+from mitmproxy import exceptions
+from unittest import mock
+
+import pytest
+
+
+def write_data(path, corrupt=False):
+ with open(path, "wb") as tf:
+ w = io.FlowWriter(tf)
+ for i in range(3):
+ f = tflow.tflow(resp=True)
+ w.add(f)
+ for i in range(3):
+ f = tflow.tflow(err=True)
+ w.add(f)
+ f = tflow.ttcpflow()
+ w.add(f)
+ f = tflow.ttcpflow(err=True)
+ w.add(f)
+ if corrupt:
+ tf.write(b"flibble")
+
+
+@mock.patch('mitmproxy.master.Master.load_flow')
+def test_configure(mck, tmpdir):
+
+ rf = readfile.ReadFile()
+ with taddons.context() as tctx:
+ tf = str(tmpdir.join("tfile"))
+ write_data(tf)
+ tctx.configure(rf, rfile=str(tf), keepserving=False)
+ assert not mck.called
+ rf.running()
+ assert mck.called
+
+ write_data(tf, corrupt=True)
+ tctx.configure(rf, rfile=str(tf), keepserving=False)
+ with pytest.raises(exceptions.OptionsError):
+ rf.running()
+
+
+@mock.patch('mitmproxy.master.Master.load_flow')
+def test_corruption(mck, tmpdir):
+
+ rf = readfile.ReadFile()
+ with taddons.context() as tctx:
+ with pytest.raises(exceptions.FlowReadException):
+ rf.load_flows_file("nonexistent")
+ assert not mck.called
+ assert len(tctx.master.event_log) == 1
+
+ tfc = str(tmpdir.join("tfile"))
+ write_data(tfc, corrupt=True)
+
+ with pytest.raises(exceptions.FlowReadException):
+ rf.load_flows_file(tfc)
+ assert mck.called
+ assert len(tctx.master.event_log) == 2
diff --git a/test/mitmproxy/addons/test_readstdin.py b/test/mitmproxy/addons/test_readstdin.py
new file mode 100644
index 00000000..bbef81fc
--- /dev/null
+++ b/test/mitmproxy/addons/test_readstdin.py
@@ -0,0 +1,59 @@
+
+import io
+from mitmproxy.addons import readstdin
+from mitmproxy.test import taddons
+from mitmproxy.test import tflow
+import mitmproxy.io
+from unittest import mock
+
+
+def gen_data(corrupt=False):
+ tf = io.BytesIO()
+ w = mitmproxy.io.FlowWriter(tf)
+ for i in range(3):
+ f = tflow.tflow(resp=True)
+ w.add(f)
+ for i in range(3):
+ f = tflow.tflow(err=True)
+ w.add(f)
+ f = tflow.ttcpflow()
+ w.add(f)
+ f = tflow.ttcpflow(err=True)
+ w.add(f)
+ if corrupt:
+ tf.write(b"flibble")
+ tf.seek(0)
+ return tf
+
+
+def test_configure(tmpdir):
+ rf = readstdin.ReadStdin()
+ with taddons.context() as tctx:
+ tctx.configure(rf, keepserving=False)
+
+
+class mStdin:
+ def __init__(self, d):
+ self.buffer = d
+
+ def isatty(self):
+ return False
+
+
+@mock.patch('mitmproxy.master.Master.load_flow')
+def test_read(m, tmpdir):
+ rf = readstdin.ReadStdin()
+ with taddons.context() as tctx:
+ assert not m.called
+ rf.running(stdin=mStdin(gen_data()))
+ assert m.called
+
+ rf.running(stdin=mStdin(None))
+ assert tctx.master.event_log
+ tctx.master.clear()
+
+ m.reset_mock()
+ assert not m.called
+ rf.running(stdin=mStdin(gen_data(corrupt=True)))
+ assert m.called
+ assert tctx.master.event_log
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index f4d32cbb..1fb33bb2 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -3,12 +3,13 @@ import pytest
from mitmproxy.test import tflow
import mitmproxy.io
-from mitmproxy import flowfilter, options
+from mitmproxy import flowfilter
+from mitmproxy import options
+from mitmproxy.proxy import config
from mitmproxy.contrib import tnetstring
from mitmproxy.exceptions import FlowReadException
from mitmproxy import flow
from mitmproxy import http
-from mitmproxy.proxy import ProxyConfig
from mitmproxy.proxy.server import DummyServer
from mitmproxy import master
from . import tservers
@@ -16,23 +17,6 @@ from . import tservers
class TestSerialize:
- def _treader(self):
- sio = io.BytesIO()
- w = mitmproxy.io.FlowWriter(sio)
- for i in range(3):
- f = tflow.tflow(resp=True)
- w.add(f)
- for i in range(3):
- f = tflow.tflow(err=True)
- w.add(f)
- f = tflow.ttcpflow()
- w.add(f)
- f = tflow.ttcpflow(err=True)
- w.add(f)
-
- sio.seek(0)
- return mitmproxy.io.FlowReader(sio)
-
def test_roundtrip(self):
sio = io.BytesIO()
f = tflow.tflow()
@@ -51,26 +35,6 @@ class TestSerialize:
assert f2.request == f.request
assert f2.marked
- def test_load_flows(self):
- r = self._treader()
- s = tservers.TestState()
- fm = master.Master(None, DummyServer())
- fm.addons.add(s)
- fm.load_flows(r)
- assert len(s.flows) == 6
-
- def test_load_flows_reverse(self):
- r = self._treader()
- s = tservers.TestState()
- opts = options.Options(
- mode="reverse:https://use-this-domain"
- )
- conf = ProxyConfig(opts)
- fm = master.Master(opts, DummyServer(conf))
- fm.addons.add(s)
- fm.load_flows(r)
- assert s.flows[0].request.host == "use-this-domain"
-
def test_filter(self):
sio = io.BytesIO()
flt = flowfilter.parse("~c 200")
@@ -122,6 +86,17 @@ class TestSerialize:
class TestFlowMaster:
+ def test_load_flow_reverse(self):
+ s = tservers.TestState()
+ opts = options.Options(
+ mode="reverse:https://use-this-domain"
+ )
+ conf = config.ProxyConfig(opts)
+ fm = master.Master(opts, DummyServer(conf))
+ fm.addons.add(s)
+ f = tflow.tflow(resp=True)
+ fm.load_flow(f)
+ assert s.flows[0].request.host == "use-this-domain"
def test_replay(self):
fm = master.Master(None, DummyServer())
diff --git a/test/mitmproxy/tools/console/test_master.py b/test/mitmproxy/tools/console/test_master.py
index 6c716ad1..45908450 100644
--- a/test/mitmproxy/tools/console/test_master.py
+++ b/test/mitmproxy/tools/console/test_master.py
@@ -5,6 +5,7 @@ from mitmproxy import proxy
from mitmproxy import options
from mitmproxy.tools.console import common
from ... import tservers
+import urwid
def test_format_keyvals():
@@ -35,7 +36,10 @@ class TestMaster(tservers.MasterTest):
def test_basic(self):
m = self.mkmaster()
for i in (1, 2, 3):
- self.dummy_cycle(m, 1, b"")
+ try:
+ self.dummy_cycle(m, 1, b"")
+ except urwid.ExitMainLoop:
+ pass
assert len(m.view) == i
def test_run_script_once(self):
diff --git a/test/mitmproxy/tools/test_dump.py b/test/mitmproxy/tools/test_dump.py
index a15bf583..624bf08f 100644
--- a/test/mitmproxy/tools/test_dump.py
+++ b/test/mitmproxy/tools/test_dump.py
@@ -2,7 +2,6 @@ import pytest
from unittest import mock
from mitmproxy import proxy
-from mitmproxy import exceptions
from mitmproxy import log
from mitmproxy import controller
from mitmproxy import options
@@ -17,18 +16,6 @@ class TestDumpMaster(tservers.MasterTest):
m = dump.DumpMaster(o, proxy.DummyServer(), with_termlog=False, with_dumper=False)
return m
- def test_read(self, tmpdir):
- p = str(tmpdir.join("read"))
- self.flowfile(p)
- self.dummy_cycle(
- self.mkmaster(None, rfile=p),
- 1, b"",
- )
- with pytest.raises(exceptions.OptionsError):
- self.mkmaster(None, rfile="/nonexistent")
- with pytest.raises(exceptions.OptionsError):
- self.mkmaster(None, rfile="test_dump.py")
-
def test_has_error(self):
m = self.mkmaster(None)
ent = log.LogEntry("foo", "error")
generated by cgit v1.2.3 (git 2.25.1) at 2024-07-26 14:20:50 +0000