From e7b61995adcb8f6e5391847043d5be3c69bfc3f9 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Tue, 14 Jun 2016 12:34:22 +1200 Subject: [PATCH 1/4] A tiny wrapper to make script invocations nicer --- mitmproxy/script/script.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/mitmproxy/script/script.py b/mitmproxy/script/script.py index 70f74817d..48efcde54 100644 --- a/mitmproxy/script/script.py +++ b/mitmproxy/script/script.py @@ -9,12 +9,21 @@ from __future__ import absolute_import, print_function, division import os import shlex import sys +import contextlib import six from mitmproxy import exceptions +@contextlib.contextmanager +def setargs(args): + oldargs = sys.argv + sys.argv = args + yield + sys.argv = oldargs + + class Script(object): """ @@ -113,7 +122,8 @@ class Script(object): f = self.ns.get(name) if f: try: - return f(self.ctx, *args, **kwargs) + with setargs(self.args): + return f(self.ctx, *args, **kwargs) except Exception: six.reraise( exceptions.ScriptException, From 804efe9d38449d484a37430459d1af45c7fcf0cf Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Mon, 13 Jun 2016 18:08:57 -0700 Subject: [PATCH 2/4] fix contextmanager, remove args argument to start() --- mitmproxy/script/script.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/mitmproxy/script/script.py b/mitmproxy/script/script.py index 48efcde54..9d149f4d8 100644 --- a/mitmproxy/script/script.py +++ b/mitmproxy/script/script.py @@ -6,10 +6,12 @@ by the mitmproxy-specific ScriptContext. # Do not import __future__ here, this would apply transitively to the inline scripts. from __future__ import absolute_import, print_function, division +import inspect import os import shlex import sys import contextlib +import warnings import six @@ -20,8 +22,10 @@ from mitmproxy import exceptions def setargs(args): oldargs = sys.argv sys.argv = args - yield - sys.argv = oldargs + try: + yield + finally: + sys.argv = oldargs class Script(object): @@ -98,7 +102,15 @@ class Script(object): finally: sys.path.pop() sys.path.pop() - return self.run("start", self.args) + + start_fn = self.ns.get("start") + if len(inspect.getargspec(start_fn).args) == 2: + warnings.warn( + "The 'args' argument of the start() script hook is deprecated. " + "Please use sys.argv instead." + ) + return self.run("start", self.args) + return self.run("start") def unload(self): try: From 63c6660895c83967c5d93499bcda3e233f722ded Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Mon, 13 Jun 2016 18:17:09 -0700 Subject: [PATCH 3/4] update examples, tests, docs --- docs/scripting/inlinescripts.rst | 2 +- examples/custom_contentviews.py | 2 +- examples/filt.py | 8 ++++---- examples/flowwriter.py | 8 ++++---- examples/har_extractor.py | 7 ++++--- examples/iframe_injector.py | 7 ++++--- examples/modify_response_body.py | 8 +++++--- examples/proxapp.py | 2 +- examples/sslstrip.py | 2 +- examples/stub.py | 2 +- examples/tcp_message.py | 4 ++-- examples/tls_passthrough.py | 7 ++++--- test/mitmproxy/data/scripts/a.py | 6 ++++-- test/mitmproxy/data/scripts/concurrent_decorator_err.py | 2 +- test/mitmproxy/data/scripts/starterr.py | 4 ++-- 15 files changed, 39 insertions(+), 32 deletions(-) diff --git a/docs/scripting/inlinescripts.rst b/docs/scripting/inlinescripts.rst index d282dfa6d..2065923dc 100644 --- a/docs/scripting/inlinescripts.rst +++ b/docs/scripting/inlinescripts.rst @@ -44,7 +44,7 @@ to store any form of state you require. Script Lifecycle Events ^^^^^^^^^^^^^^^^^^^^^^^ -.. py:function:: start(context, argv) +.. py:function:: start(context) Called once on startup, before any other events. diff --git a/examples/custom_contentviews.py b/examples/custom_contentviews.py index 034f356c0..05ebeb69f 100644 --- a/examples/custom_contentviews.py +++ b/examples/custom_contentviews.py @@ -62,7 +62,7 @@ class ViewPigLatin(contentviews.View): pig_view = ViewPigLatin() -def start(context, argv): +def start(context): context.add_contentview(pig_view) diff --git a/examples/filt.py b/examples/filt.py index f99b675ca..1a423845c 100644 --- a/examples/filt.py +++ b/examples/filt.py @@ -1,13 +1,13 @@ # This scripts demonstrates how to use mitmproxy's filter pattern in inline scripts. # Usage: mitmdump -s "filt.py FILTER" - +import sys from mitmproxy import filt -def start(context, argv): - if len(argv) != 2: +def start(context): + if len(sys.argv) != 2: raise ValueError("Usage: -s 'filt.py FILTER'") - context.filter = filt.parse(argv[1]) + context.filter = filt.parse(sys.argv[1]) def response(context, flow): diff --git a/examples/flowwriter.py b/examples/flowwriter.py index 8fb8cc60f..cb5ccb0dc 100644 --- a/examples/flowwriter.py +++ b/examples/flowwriter.py @@ -4,14 +4,14 @@ import sys from mitmproxy.flow import FlowWriter -def start(context, argv): - if len(argv) != 2: +def start(context): + if len(sys.argv) != 2: raise ValueError('Usage: -s "flowriter.py filename"') - if argv[1] == "-": + if sys.argv[1] == "-": f = sys.stdout else: - f = open(argv[1], "wb") + f = open(sys.argv[1], "wb") context.flow_writer = FlowWriter(f) diff --git a/examples/har_extractor.py b/examples/har_extractor.py index 6806989df..c21f1a8f7 100644 --- a/examples/har_extractor.py +++ b/examples/har_extractor.py @@ -3,6 +3,7 @@ https://github.com/JustusW/harparser to generate a HAR log object. """ import six +import sys from harparser import HAR from datetime import datetime @@ -52,15 +53,15 @@ class _HARLog(HAR.log): return self.__page_list__ -def start(context, argv): +def start(context): """ On start we create a HARLog instance. You will have to adapt this to suit your actual needs of HAR generation. As it will probably be necessary to cluster logs by IPs or reset them from time to time. """ context.dump_file = None - if len(argv) > 1: - context.dump_file = argv[1] + if len(sys.argv) > 1: + context.dump_file = sys.argv[1] else: raise ValueError( 'Usage: -s "har_extractor.py filename" ' diff --git a/examples/iframe_injector.py b/examples/iframe_injector.py index ad844f197..9495da93d 100644 --- a/examples/iframe_injector.py +++ b/examples/iframe_injector.py @@ -1,13 +1,14 @@ # Usage: mitmdump -s "iframe_injector.py url" # (this script works best with --anticache) +import sys from bs4 import BeautifulSoup from mitmproxy.models import decoded -def start(context, argv): - if len(argv) != 2: +def start(context): + if len(sys.argv) != 2: raise ValueError('Usage: -s "iframe_injector.py url"') - context.iframe_url = argv[1] + context.iframe_url = sys.argv[1] def response(context, flow): diff --git a/examples/modify_response_body.py b/examples/modify_response_body.py index d68bcf63e..3034892e0 100644 --- a/examples/modify_response_body.py +++ b/examples/modify_response_body.py @@ -1,14 +1,16 @@ # Usage: mitmdump -s "modify_response_body.py mitmproxy bananas" # (this script works best with --anticache) +import sys + from mitmproxy.models import decoded -def start(context, argv): - if len(argv) != 3: +def start(context): + if len(sys.argv) != 3: raise ValueError('Usage: -s "modify_response_body.py old new"') # You may want to use Python's argparse for more sophisticated argument # parsing. - context.old, context.new = argv[1], argv[2] + context.old, context.new = sys.argv[1], sys.argv[2] def response(context, flow): diff --git a/examples/proxapp.py b/examples/proxapp.py index 4d8e7b582..613d3f8b0 100644 --- a/examples/proxapp.py +++ b/examples/proxapp.py @@ -15,7 +15,7 @@ def hello_world(): # Register the app using the magic domain "proxapp" on port 80. Requests to # this domain and port combination will now be routed to the WSGI app instance. -def start(context, argv): +def start(context): context.app_registry.add(app, "proxapp", 80) # SSL works too, but the magic domain needs to be resolvable from the mitmproxy machine due to mitmproxy's design. diff --git a/examples/sslstrip.py b/examples/sslstrip.py index 1bc899463..8dde8e3ef 100644 --- a/examples/sslstrip.py +++ b/examples/sslstrip.py @@ -3,7 +3,7 @@ import re from six.moves import urllib -def start(context, argv): +def start(context): # set of SSL/TLS capable hosts context.secure_hosts = set() diff --git a/examples/stub.py b/examples/stub.py index 516b71a53..a0f73538e 100644 --- a/examples/stub.py +++ b/examples/stub.py @@ -3,7 +3,7 @@ """ -def start(context, argv): +def start(context): """ Called once on script startup, before any other events. """ diff --git a/examples/tcp_message.py b/examples/tcp_message.py index 2c2106185..78500c19c 100644 --- a/examples/tcp_message.py +++ b/examples/tcp_message.py @@ -1,4 +1,4 @@ -''' +""" tcp_message Inline Script Hook API Demonstration ------------------------------------------------ @@ -7,7 +7,7 @@ tcp_message Inline Script Hook API Demonstration example cmdline invocation: mitmdump -T --host --tcp ".*" -q -s examples/tcp_message.py -''' +""" from netlib import strutils diff --git a/examples/tls_passthrough.py b/examples/tls_passthrough.py index 0c6d450d7..50aab65b0 100644 --- a/examples/tls_passthrough.py +++ b/examples/tls_passthrough.py @@ -24,6 +24,7 @@ from __future__ import (absolute_import, print_function, division) import collections import random +import sys from enum import Enum from mitmproxy.exceptions import TlsProtocolException @@ -110,9 +111,9 @@ class TlsFeedback(TlsLayer): # inline script hooks below. -def start(context, argv): - if len(argv) == 2: - context.tls_strategy = ProbabilisticStrategy(float(argv[1])) +def start(context): + if len(sys.argv) == 2: + context.tls_strategy = ProbabilisticStrategy(float(sys.argv[1])) else: context.tls_strategy = ConservativeStrategy() diff --git a/test/mitmproxy/data/scripts/a.py b/test/mitmproxy/data/scripts/a.py index d4272ac85..33dbaa64c 100644 --- a/test/mitmproxy/data/scripts/a.py +++ b/test/mitmproxy/data/scripts/a.py @@ -1,11 +1,13 @@ +import sys + from a_helper import parser var = 0 -def start(ctx, argv): +def start(ctx): global var - var = parser.parse_args(argv[1:]).var + var = parser.parse_args(sys.argv[1:]).var def here(ctx): diff --git a/test/mitmproxy/data/scripts/concurrent_decorator_err.py b/test/mitmproxy/data/scripts/concurrent_decorator_err.py index 071b8889f..349e5dd64 100644 --- a/test/mitmproxy/data/scripts/concurrent_decorator_err.py +++ b/test/mitmproxy/data/scripts/concurrent_decorator_err.py @@ -2,5 +2,5 @@ from mitmproxy.script import concurrent @concurrent -def start(context, argv): +def start(context): pass diff --git a/test/mitmproxy/data/scripts/starterr.py b/test/mitmproxy/data/scripts/starterr.py index b217bdfee..82d773bd9 100644 --- a/test/mitmproxy/data/scripts/starterr.py +++ b/test/mitmproxy/data/scripts/starterr.py @@ -1,3 +1,3 @@ -def start(ctx, argv): - raise ValueError +def start(ctx): + raise ValueError() From c2b5a13e3f1cfd193184c2fca9df0d531501a8ab Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Tue, 14 Jun 2016 14:39:07 +1200 Subject: [PATCH 4/4] Don't fail when start() is not defined --- mitmproxy/script/script.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mitmproxy/script/script.py b/mitmproxy/script/script.py index 9d149f4d8..9ff79f528 100644 --- a/mitmproxy/script/script.py +++ b/mitmproxy/script/script.py @@ -104,7 +104,7 @@ class Script(object): sys.path.pop() start_fn = self.ns.get("start") - if len(inspect.getargspec(start_fn).args) == 2: + if start_fn and len(inspect.getargspec(start_fn).args) == 2: warnings.warn( "The 'args' argument of the start() script hook is deprecated. " "Please use sys.argv instead."