Merge pull request #1254 from mitmproxy/scriptargs

Inline Scripts: use sys.argv instead of args argument.

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.
 

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)
 
 

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):
+def start(context):
-    if len(argv) != 2:
+    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):

examples/flowwriter.py

@@ -4,14 +4,14 @@ import sys
 from mitmproxy.flow import FlowWriter
 
 
-def start(context, argv):
+def start(context):
-    if len(argv) != 2:
+    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)
 
 

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:
+    if len(sys.argv) > 1:
-        context.dump_file = argv[1]
+        context.dump_file = sys.argv[1]
     else:
         raise ValueError(
             'Usage: -s "har_extractor.py filename" '

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):
+def start(context):
-    if len(argv) != 2:
+    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):

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):
+def start(context):
-    if len(argv) != 3:
+    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):

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.

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()
 

examples/stub.py

@@ -3,7 +3,7 @@
 """
 
 
-def start(context, argv):
+def start(context):
     """
         Called once on script startup, before any other events.
     """

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
 
 

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):
+def start(context):
-    if len(argv) == 2:
+    if len(sys.argv) == 2:
-        context.tls_strategy = ProbabilisticStrategy(float(argv[1]))
+        context.tls_strategy = ProbabilisticStrategy(float(sys.argv[1]))
     else:
         context.tls_strategy = ConservativeStrategy()
 

mitmproxy/script/script.py

@@ -6,15 +6,28 @@ 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
 
 from mitmproxy import exceptions
 
 
+@contextlib.contextmanager
+def setargs(args):
+    oldargs = sys.argv
+    sys.argv = args
+    try:
+        yield
+    finally:
+        sys.argv = oldargs
+
+
 class Script(object):
 
     """
@@ -89,7 +102,15 @@ class Script(object):
         finally:
             sys.path.pop()
             sys.path.pop()
+
+        start_fn = self.ns.get("start")
+        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."
+            )
             return self.run("start", self.args)
+        return self.run("start")
 
     def unload(self):
         try:
@@ -113,6 +134,7 @@ class Script(object):
         f = self.ns.get(name)
         if f:
             try:
+                with setargs(self.args):
                     return f(self.ctx, *args, **kwargs)
             except Exception:
                 six.reraise(

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):

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

test/mitmproxy/data/scripts/starterr.py

@@ -1,3 +1,3 @@
 
-def start(ctx, argv):
+def start(ctx):
-    raise ValueError
+    raise ValueError()