Merge pull request #2265 from cortesi/addons

Addons and addon testing

3 files changed, 38 insertions, 31 deletions

diff --git a/examples/complex/har_dump.py b/examples/complex/har_dump.py
index b6b9ffae..21bcc341 100644
--- a/examples/complex/har_dump.py
+++ b/examples/complex/har_dump.py
@@ -4,7 +4,6 @@ This inline script can be used to dump flows as HAR files.
 
 
 import json
-import sys
 import base64
 import zlib
 import os
@@ -15,6 +14,7 @@ from datetime import timezone
 import mitmproxy
 
 from mitmproxy import version
+from mitmproxy import ctx
 from mitmproxy.utils import strutils
 from mitmproxy.net.http import cookies
 
@@ -26,16 +26,12 @@ SERVERS_SEEN = set()
 
 
 def load(l):
-    """
-        Called once on script startup before any other events.
-    """
-    if len(sys.argv) != 2:
-        raise ValueError(
-            'Usage: -s "har_dump.py filename" '
-            '(- will output to stdout, filenames ending with .zhar '
-            'will result in compressed har)'
-        )
+    l.add_option(
+        "hardump", str, "", "HAR dump path.",
+    )
+
 
+def configure(updated):
     HAR.update({
         "log": {
             "version": "1.2",
@@ -156,21 +152,20 @@ def done():
     """
         Called once on script shutdown, after any other events.
     """
-    dump_file = sys.argv[1]
+    if ctx.options.hardump:
+        json_dump = json.dumps(HAR, indent=2)  # type: str
 
-    json_dump = json.dumps(HAR, indent=2)  # type: str
-
-    if dump_file == '-':
-        mitmproxy.ctx.log(json_dump)
-    else:
-        raw = json_dump.encode()  # type: bytes
-        if dump_file.endswith('.zhar'):
-            raw = zlib.compress(raw, 9)
+        if ctx.options.hardump == '-':
+            mitmproxy.ctx.log(json_dump)
+        else:
+            raw = json_dump.encode()  # type: bytes
+            if ctx.options.hardump.endswith('.zhar'):
+                raw = zlib.compress(raw, 9)
 
-        with open(os.path.expanduser(dump_file), "wb") as f:
-            f.write(raw)
+            with open(os.path.expanduser(ctx.options.hardump), "wb") as f:
+                f.write(raw)
 
-        mitmproxy.ctx.log("HAR dump finished (wrote %s bytes to file)" % len(json_dump))
+            mitmproxy.ctx.log("HAR dump finished (wrote %s bytes to file)" % len(json_dump))
 
 
 def format_cookies(cookie_list):
diff --git a/examples/complex/tls_passthrough.py b/examples/complex/tls_passthrough.py
index 72c0244b..9bb27d25 100644
--- a/examples/complex/tls_passthrough.py
+++ b/examples/complex/tls_passthrough.py
@@ -23,10 +23,10 @@ Authors: Maximilian Hils, Matthew Tuusberg
 import collections
 import random
 
-import sys
 from enum import Enum
 
 import mitmproxy
+from mitmproxy import ctx
 from mitmproxy.exceptions import TlsProtocolException
 from mitmproxy.proxy.protocol import TlsLayer, RawTCPLayer
 
@@ -113,9 +113,15 @@ tls_strategy = None
 
 
 def load(l):
+    l.add_option(
+        "tlsstrat", int, 0, "TLS passthrough strategy (0-100)",
+    )
+
+
+def configure(updated):
     global tls_strategy
-    if len(sys.argv) == 2:
-        tls_strategy = ProbabilisticStrategy(float(sys.argv[1]))
+    if ctx.options.tlsstrat > 0:
+        tls_strategy = ProbabilisticStrategy(float(ctx.options.tlsstrat) / 100.0)
     else:
         tls_strategy = ConservativeStrategy()
 
diff --git a/examples/simple/filter_flows.py b/examples/simple/filter_flows.py
index 896fa54a..fd49425a 100644
--- a/examples/simple/filter_flows.py
+++ b/examples/simple/filter_flows.py
@@ -1,15 +1,21 @@
 """
 This scripts demonstrates how to use mitmproxy's filter pattern in scripts.
-Usage:
-    mitmdump -s "flowfilter.py FILTER"
 """
-import sys
 from mitmproxy import flowfilter
+from mitmproxy import ctx
 
 
 class Filter:
-    def __init__(self, spec):
-        self.filter = flowfilter.parse(spec)
+    def __init__(self):
+        self.filter = None
+
+    def configure(self, updated):
+        self.filter = flowfilter.parse(ctx.options.flowfilter)
+
+    def load(self, l):
+        l.add_option(
+            "flowfilter", str, "", "Check that flow matches filter."
+        )
 
     def response(self, flow):
         if flowfilter.match(self.filter, flow):
@@ -17,4 +23,4 @@ class Filter:
             print(flow)
 
 
-addons = [Filter(sys.argv[1])]
+addons = [Filter()]