diff options
author | Maximilian Hils <git@maximilianhils.com> | 2017-04-26 13:40:51 +0200 |
---|---|---|
committer | GitHub <noreply@github.com> | 2017-04-26 13:40:51 +0200 |
commit | d5ea08db6244e4957f9134a4313a90d3efb8dde9 (patch) | |
tree | 14069d73669933f6f5341ab885e6ad7091c5dd67 /mitmproxy | |
parent | 0a8e54edeab4c1a250333874129dfb2360d27af3 (diff) | |
parent | ca2827886a97de88f4ab2937e71588fc9320ba4d (diff) | |
download | mitmproxy-d5ea08db6244e4957f9134a4313a90d3efb8dde9.tar.gz mitmproxy-d5ea08db6244e4957f9134a4313a90d3efb8dde9.tar.bz2 mitmproxy-d5ea08db6244e4957f9134a4313a90d3efb8dde9.zip |
Merge pull request #2258 from mhils/readfile
Integrate readstdin into readfile
Diffstat (limited to 'mitmproxy')
-rw-r--r-- | mitmproxy/addons/__init__.py | 2 | ||||
-rw-r--r-- | mitmproxy/addons/readfile.py | 52 | ||||
-rw-r--r-- | mitmproxy/addons/readstdin.py | 26 | ||||
-rw-r--r-- | mitmproxy/tools/console/master.py | 10 | ||||
-rw-r--r-- | mitmproxy/tools/dump.py | 4 | ||||
-rw-r--r-- | mitmproxy/tools/web/master.py | 2 | ||||
-rw-r--r-- | mitmproxy/utils/typecheck.py | 2 |
7 files changed, 49 insertions, 49 deletions
diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index b4367d78..7a45106c 100644 --- a/mitmproxy/addons/__init__.py +++ b/mitmproxy/addons/__init__.py @@ -8,7 +8,6 @@ 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 @@ -38,6 +37,5 @@ 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 index 949da15d..05b6c309 100644 --- a/mitmproxy/addons/readfile.py +++ b/mitmproxy/addons/readfile.py @@ -1,38 +1,56 @@ import os.path +import sys +import typing from mitmproxy import ctx -from mitmproxy import io from mitmproxy import exceptions +from mitmproxy import io class ReadFile: """ An addon that handles reading from file on startup. """ - def load_flows_file(self, path: str) -> int: - path = os.path.expanduser(path) + + def load_flows(self, fo: typing.IO[bytes]) -> int: cnt = 0 + freader = io.FlowReader(fo) 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: + for flow in freader.stream(): + ctx.master.load_flow(flow) + cnt += 1 + except (IOError, exceptions.FlowReadException) as e: if cnt: - ctx.log.warn( - "Flow file corrupted - loaded %i flows." % cnt, - ) + ctx.log.warn("Flow file corrupted - loaded %i flows." % cnt) else: ctx.log.error("Flow file corrupted.") - raise exceptions.FlowReadException(v) + raise exceptions.FlowReadException(str(e)) from e + else: + return cnt + + def load_flows_from_path(self, path: str) -> int: + path = os.path.expanduser(path) + try: + with open(path, "rb") as f: + return self.load_flows(f) + except IOError as e: + ctx.log.error("Cannot load flows: {}".format(e)) + raise exceptions.FlowReadException(str(e)) from e def running(self): if ctx.options.rfile: try: - self.load_flows_file(ctx.options.rfile) - except exceptions.FlowReadException as v: - raise exceptions.OptionsError(v) + self.load_flows_from_path(ctx.options.rfile) + except exceptions.FlowReadException as e: + raise exceptions.OptionsError(e) from e finally: ctx.master.addons.trigger("processing_complete") + + +class ReadFileStdin(ReadFile): + """Support the special case of "-" for reading from stdin""" + def load_flows_from_path(self, path: str) -> int: + if path == "-": + return self.load_flows(sys.stdin.buffer) + else: + return super().load_flows_from_path(path) diff --git a/mitmproxy/addons/readstdin.py b/mitmproxy/addons/readstdin.py deleted file mode 100644 index 93a99f01..00000000 --- a/mitmproxy/addons/readstdin.py +++ /dev/null @@ -1,26 +0,0 @@ -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 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) - ctx.master.addons.trigger("processing_complete") diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py index c1d584ac..e7a2c6ae 100644 --- a/mitmproxy/tools/console/master.py +++ b/mitmproxy/tools/console/master.py @@ -17,8 +17,9 @@ from mitmproxy import exceptions from mitmproxy import master from mitmproxy import io from mitmproxy import log -from mitmproxy.addons import view from mitmproxy.addons import intercept +from mitmproxy.addons import readfile +from mitmproxy.addons import view from mitmproxy.tools.console import flowlist from mitmproxy.tools.console import flowview from mitmproxy.tools.console import grideditor @@ -91,7 +92,12 @@ class ConsoleMaster(master.Master): signals.sig_add_log.connect(self.sig_add_log) self.addons.add(Logger()) self.addons.add(*addons.default_addons()) - self.addons.add(intercept.Intercept(), self.view, UnsupportedLog()) + self.addons.add( + intercept.Intercept(), + self.view, + UnsupportedLog(), + readfile.ReadFile(), + ) def sigint_handler(*args, **kwargs): self.prompt_for_exit() diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py index 6329f6b7..4d0ccf4b 100644 --- a/mitmproxy/tools/dump.py +++ b/mitmproxy/tools/dump.py @@ -1,7 +1,7 @@ from mitmproxy import addons from mitmproxy import options from mitmproxy import master -from mitmproxy.addons import dumper, termlog, termstatus, readstdin, keepserving +from mitmproxy.addons import dumper, termlog, termstatus, keepserving, readfile class ErrorCheck: @@ -30,7 +30,7 @@ class DumpMaster(master.Master): if with_dumper: self.addons.add(dumper.Dumper()) self.addons.add( - readstdin.ReadStdin(), keepserving.KeepServing(), + readfile.ReadFileStdin(), self.errorcheck ) diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py index 0db5a09f..c09fe0a2 100644 --- a/mitmproxy/tools/web/master.py +++ b/mitmproxy/tools/web/master.py @@ -7,6 +7,7 @@ from mitmproxy import log from mitmproxy import master from mitmproxy.addons import eventstore from mitmproxy.addons import intercept +from mitmproxy.addons import readfile from mitmproxy.addons import termlog from mitmproxy.addons import view from mitmproxy.addons import termstatus @@ -32,6 +33,7 @@ class WebMaster(master.Master): self.addons.add(*addons.default_addons()) self.addons.add( intercept.Intercept(), + readfile.ReadFile(), self.view, self.events, ) diff --git a/mitmproxy/utils/typecheck.py b/mitmproxy/utils/typecheck.py index 5df4ea4b..628ea642 100644 --- a/mitmproxy/utils/typecheck.py +++ b/mitmproxy/utils/typecheck.py @@ -68,5 +68,7 @@ def check_type(name: str, value: typing.Any, typeinfo: typing.Any) -> None: return else: raise e + elif typename.startswith("typing.Any"): + return elif not isinstance(value, typeinfo): raise e |