[sans-io] add options to context

This commit is contained in:
Maximilian Hils 2017-08-05 02:32:55 +02:00
parent e7b67e96a2
commit eddef85f5f
5 changed files with 28 additions and 9 deletions

View File

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

View File

@ -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 = []

View File

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

View File

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

View File

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