From eddef85f5f98988b47d43a7a698689febf7c6c62 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Sat, 5 Aug 2017 02:32:55 +0200
Subject: [PATCH] [sans-io] add options to context

---
 mitmproxy/addons/proxyserver.py   | 14 +++++++++++---
 mitmproxy/proxy2/context.py       |  9 ++++++++-
 mitmproxy/proxy2/server.py        |  7 ++++---
 mitmproxy/proxy2/test/conftest.py |  4 +++-
 mitmproxy/proxy2/test/tutils.py   |  3 ++-
 5 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/mitmproxy/addons/proxyserver.py b/mitmproxy/addons/proxyserver.py
index 7d8ba53e7..501b9e553 100644
--- a/mitmproxy/addons/proxyserver.py
+++ b/mitmproxy/addons/proxyserver.py
@@ -23,10 +23,10 @@ class ProxyConnectionHandler(server.ConnectionHandler):
     event_queue: queue.Queue
     loop: asyncio.AbstractEventLoop
 
-    def __init__(self, event_queue, loop, r, w):
+    def __init__(self, event_queue, loop, r, w, options):
         self.event_queue = event_queue
         self.loop = loop
-        super().__init__(r, w)
+        super().__init__(r, w, options)
 
     async def handle_hook(self, hook: commands.Hook) -> None:
         q = asyncio.Queue()
@@ -41,6 +41,7 @@ class ProxyConnectionHandler(server.ConnectionHandler):
         x.reply = controller.DummyReply()
         self.event_queue.put(("log",  x))
 
+
 class Proxyserver:
     """
     This addon runs the actual proxy server.
@@ -51,6 +52,7 @@ class Proxyserver:
         self.loop = asyncio.get_event_loop()
         self.listen_port = None
         self.event_queue = None
+        self.options = ctx.options
         self._lock = asyncio.Lock()
 
     def running(self):
@@ -73,7 +75,13 @@ class Proxyserver:
             )
 
     async def handle_connection(self, r, w):
-        await ProxyConnectionHandler(self.event_queue, self.loop, r, w).handle_client()
+        await ProxyConnectionHandler(
+            self.event_queue,
+            self.loop,
+            r,
+            w,
+            self.options
+        ).handle_client()
 
     def configure(self, updated):
         if "listen_port" in updated:
diff --git a/mitmproxy/proxy2/context.py b/mitmproxy/proxy2/context.py
index 4075000b3..14620a6e9 100644
--- a/mitmproxy/proxy2/context.py
+++ b/mitmproxy/proxy2/context.py
@@ -1,4 +1,6 @@
-from typing import Optional
+from typing import Optional, List
+
+from mitmproxy.options import Options
 
 
 class Connection:
@@ -37,11 +39,16 @@ class Context:
 
     client: Client
     server: Optional[Server]
+    options: Options
+    layers: List["mitmproxy.proxy2.layer.Layer"]
 
     def __init__(
             self,
             client: Client,
             server: Optional[Server],
+            options: Options,
     ) -> None:
         self.client = client
         self.server = server
+        self.options = options
+        self.layers = []
diff --git a/mitmproxy/proxy2/server.py b/mitmproxy/proxy2/server.py
index 6a17f7bc0..6f7b7ed61 100644
--- a/mitmproxy/proxy2/server.py
+++ b/mitmproxy/proxy2/server.py
@@ -11,6 +11,7 @@ import asyncio
 import socket
 import typing
 
+from mitmproxy import options
 from mitmproxy.proxy2 import events, commands
 from mitmproxy.proxy2.context import Client, Context, Connection
 from mitmproxy.proxy2.layers.modes import ReverseProxy
@@ -24,11 +25,11 @@ class StreamIO(typing.NamedTuple):
 class ConnectionHandler(metaclass=abc.ABCMeta):
     transports: typing.MutableMapping[Connection, StreamIO]
 
-    def __init__(self, reader, writer):
+    def __init__(self, reader, writer, options):
         addr = writer.get_extra_info('peername')
 
         self.client = Client(addr)
-        self.context = Context(self.client, None)
+        self.context = Context(self.client, None, options)
 
         # self.layer = ReverseProxy(self.context, ("localhost", 443))
         self.layer = ReverseProxy(self.context, ("localhost", 8000))
@@ -134,7 +135,7 @@ if __name__ == "__main__":
 
 
     async def handle(reader, writer):
-        await SimpleConnectionHandler(reader, writer).handle_client()
+        await SimpleConnectionHandler(reader, writer, options.Options()).handle_client()
 
 
     coro = asyncio.start_server(handle, '127.0.0.1', 8080, loop=loop)
diff --git a/mitmproxy/proxy2/test/conftest.py b/mitmproxy/proxy2/test/conftest.py
index 4267183fb..8a01ee0e6 100644
--- a/mitmproxy/proxy2/test/conftest.py
+++ b/mitmproxy/proxy2/test/conftest.py
@@ -1,5 +1,6 @@
 import pytest
 
+from mitmproxy import options
 from mitmproxy.proxy2 import context
 
 
@@ -7,5 +8,6 @@ from mitmproxy.proxy2 import context
 def tctx():
     return context.Context(
         context.Client("client"),
-        context.Server("server")
+        context.Server("server"),
+        options.Options()
     )
diff --git a/mitmproxy/proxy2/test/tutils.py b/mitmproxy/proxy2/test/tutils.py
index e7900ffde..72568c84c 100644
--- a/mitmproxy/proxy2/test/tutils.py
+++ b/mitmproxy/proxy2/test/tutils.py
@@ -138,7 +138,8 @@ class playbook:
         # Playbooks are only executed on assert (which signals that the playbook is partially
         # complete), so we need to signal if someone forgets to assert and playbooks aren't
         # evaluated.
-        if not self._errored and len(self.actual) < len(self.expected):
+        is_final_destruct = not hasattr(self, "_errored")
+        if is_final_destruct or (not self._errored and len(self.actual) < len(self.expected)):
             raise RuntimeError("Unfinished playbook!")
 
     def fork(self):