diff options
| author | Aldo Cortesi <aldo@corte.si> | 2018-05-08 14:53:56 +1200 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2018-05-08 14:53:56 +1200 |
| commit | 8a682d3532ec071f20f4ba0f4700fbcdb7921699 (patch) | |
| tree | cc437a95786c34200340e43e39c145f0c5ea5d25 | |
| parent | 7ec9c5524f64bbceaf52b58696c8c9342a4887dd (diff) | |
| parent | 717fbaa99076545d11f554187759005dce1aa67b (diff) | |
| download | mitmproxy-8a682d3532ec071f20f4ba0f4700fbcdb7921699.tar.gz mitmproxy-8a682d3532ec071f20f4ba0f4700fbcdb7921699.tar.bz2 mitmproxy-8a682d3532ec071f20f4ba0f4700fbcdb7921699.zip | |
Merge pull request #3105 from cortesi/opts
Add deferred options, tweak benchmarks, document done event
| -rw-r--r-- | examples/addons/events.py | 8 | ||||
| -rw-r--r-- | mitmproxy/addonmanager.py | 1 | ||||
| -rw-r--r-- | mitmproxy/optmanager.py | 74 | ||||
| -rw-r--r-- | mitmproxy/tools/main.py | 2 | ||||
| -rw-r--r-- | test/bench/benchmark.py | 10 | ||||
| -rwxr-xr-x | test/bench/run-mitmdump (renamed from test/bench/run) | 0 | ||||
| -rwxr-xr-x | test/bench/run-mitmproxy | 4 | ||||
| -rw-r--r-- | test/mitmproxy/test_optmanager.py | 13 |
8 files changed, 84 insertions, 28 deletions
diff --git a/examples/addons/events.py b/examples/addons/events.py index cd714ba0..f83c8f11 100644 --- a/examples/addons/events.py +++ b/examples/addons/events.py @@ -142,8 +142,12 @@ class Events: def done(self): """ - Called when the addon shuts down, either by being removed from the - mitmproxy instance, or when mitmproxy itself shuts down. + Called when the addon shuts down, either by being removed from + the mitmproxy instance, or when mitmproxy itself shuts down. On + shutdown, this event is called after the event loop is + terminated, guaranteeing that it will be the final event an addon + sees. Note that log handlers are shut down at this point, so + calls to log functions will produce no output. """ def load(self, entry: mitmproxy.addonmanager.Loader): diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py index a9f6551f..64c957ba 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -173,6 +173,7 @@ class AddonManager: self.lookup[name] = a for a in traverse([addon]): self.master.commands.collect_commands(a) + self.master.options.process_deferred() return addon def add(self, *addons): diff --git a/mitmproxy/optmanager.py b/mitmproxy/optmanager.py index bb9e3030..a473e9df 100644 --- a/mitmproxy/optmanager.py +++ b/mitmproxy/optmanager.py @@ -91,9 +91,12 @@ class OptManager: mutation doesn't change the option state inadvertently. """ def __init__(self): - self.__dict__["_options"] = {} - self.__dict__["changed"] = blinker.Signal() - self.__dict__["errored"] = blinker.Signal() + self._deferred: typing.Dict[str, str] = {} + self.changed = blinker.Signal() + self.errored = blinker.Signal() + # Options must be the last attribute here - after that, we raise an + # error for attribute assigment to unknown options. + self._options: typing.Dict[str, typing.Any] = {} def add_option( self, @@ -168,7 +171,14 @@ class OptManager: raise AttributeError("No such option: %s" % attr) def __setattr__(self, attr, value): - self.update(**{attr: value}) + # This is slightly tricky. We allow attributes to be set on the instance + # until we have an _options attribute. After that, assignment is sent to + # the update function, and will raise an error for unknown options. + opts = self.__dict__.get("_options") + if not opts: + super().__setattr__(attr, value) + else: + self.update(**{attr: value}) def keys(self): return set(self._options.keys()) @@ -272,20 +282,48 @@ class OptManager: options=options ) - def set(self, *spec): + def set(self, *spec, defer=False): + """ + Takes a list of set specification in standard form (option=value). + Options that are known are updated immediately. If defer is true, + options that are not known are deferred, and will be set once they + are added. + """ vals = {} + unknown = {} for i in spec: - vals.update(self._setspec(i)) + parts = i.split("=", maxsplit=1) + if len(parts) == 1: + optname, optval = parts[0], None + else: + optname, optval = parts[0], parts[1] + if optname in self._options: + vals[optname] = self.parse_setval(self._options[optname], optval) + else: + unknown[optname] = optval + if defer: + self._deferred.update(unknown) + elif unknown: + raise exceptions.OptionsError("Unknown options: %s" % ", ".join(unknown.keys())) self.update(**vals) - def parse_setval(self, optname: str, optstr: typing.Optional[str]) -> typing.Any: + def process_deferred(self): + """ + Processes options that were deferred in previous calls to set, and + have since been added. + """ + update = {} + for optname, optval in self._deferred.items(): + if optname in self._options: + update[optname] = self.parse_setval(self._options[optname], optval) + self.update(**update) + for k in update.keys(): + del self._deferred[k] + + def parse_setval(self, o: _Option, optstr: typing.Optional[str]) -> typing.Any: """ Convert a string to a value appropriate for the option type. """ - if optname not in self._options: - raise exceptions.OptionsError("No such option %s" % optname) - o = self._options[optname] - if o.typespec in (str, typing.Optional[str]): return optstr elif o.typespec in (int, typing.Optional[int]): @@ -295,7 +333,7 @@ class OptManager: except ValueError: raise exceptions.OptionsError("Not an integer: %s" % optstr) elif o.typespec == int: - raise exceptions.OptionsError("Option is required: %s" % optname) + raise exceptions.OptionsError("Option is required: %s" % o.name) else: return None elif o.typespec == bool: @@ -313,19 +351,9 @@ class OptManager: if not optstr: return [] else: - return getattr(self, optname) + [optstr] + return getattr(self, o.name) + [optstr] raise NotImplementedError("Unsupported option type: %s", o.typespec) - def _setspec(self, spec): - d = {} - parts = spec.split("=", maxsplit=1) - if len(parts) == 1: - optname, optval = parts[0], None - else: - optname, optval = parts[0], parts[1] - d[optname] = self.parse_setval(optname, optval) - return d - def make_parser(self, parser, optname, metavar=None, short=None): """ Auto-Create a command-line parser entry for a named option. If the diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index cf370f29..0b271b91 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -110,7 +110,7 @@ def run( if args.commands: master.commands.dump() sys.exit(0) - opts.set(*args.setoptions) + opts.set(*args.setoptions, defer=True) if extra: opts.update(**extra(args)) diff --git a/test/bench/benchmark.py b/test/bench/benchmark.py index 8d208088..84ec6005 100644 --- a/test/bench/benchmark.py +++ b/test/bench/benchmark.py @@ -11,6 +11,15 @@ class Benchmark: self.pr = cProfile.Profile() self.started = False + self.resps = 0 + self.reqs = 0 + + def request(self, f): + self.reqs += 1 + + def response(self, f): + self.resps += 1 + async def procs(self): ctx.log.error("starting benchmark") backend = await asyncio.create_subprocess_exec("devd", "-q", "-p", "10001", ".") @@ -23,6 +32,7 @@ class Benchmark: ) stdout, _ = await traf.communicate() open(ctx.options.benchmark_save_path + ".bench", mode="wb").write(stdout) + ctx.log.error("Proxy saw %s requests, %s responses" % (self.reqs, self.resps)) ctx.log.error(stdout.decode("ascii")) backend.kill() ctx.master.shutdown() diff --git a/test/bench/run b/test/bench/run-mitmdump index 9925b578..9925b578 100755 --- a/test/bench/run +++ b/test/bench/run-mitmdump diff --git a/test/bench/run-mitmproxy b/test/bench/run-mitmproxy new file mode 100755 index 00000000..4f1f626e --- /dev/null +++ b/test/bench/run-mitmproxy @@ -0,0 +1,4 @@ +#!/bin/sh + +mkdir -p results +mitmproxy -p0 -q --set benchmark_save_path=./results/mitmproxy -s ./benchmark.py
\ No newline at end of file diff --git a/test/mitmproxy/test_optmanager.py b/test/mitmproxy/test_optmanager.py index 1c49c0b8..4356434b 100644 --- a/test/mitmproxy/test_optmanager.py +++ b/test/mitmproxy/test_optmanager.py @@ -70,7 +70,7 @@ def test_defaults(): def test_required_int(): o = TO() with pytest.raises(exceptions.OptionsError): - o.parse_setval("required_int", None) + o.parse_setval(o._options["required_int"], None) def test_deepcopy(): @@ -426,4 +426,13 @@ def test_set(): assert opts.seqstr == [] with pytest.raises(exceptions.OptionsError): - opts.set("nonexistent=wobble") + opts.set("deferred=wobble") + + opts.set("deferred=wobble", defer=True) + assert "deferred" in opts._deferred + opts.process_deferred() + assert "deferred" in opts._deferred + opts.add_option("deferred", str, "default", "help") + opts.process_deferred() + assert "deferred" not in opts._deferred + assert opts.deferred == "wobble" |