[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 event_queue: queue.Queue
loop: asyncio.AbstractEventLoop 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.event_queue = event_queue
self.loop = loop self.loop = loop
super().__init__(r, w) super().__init__(r, w, options)
async def handle_hook(self, hook: commands.Hook) -> None: async def handle_hook(self, hook: commands.Hook) -> None:
q = asyncio.Queue() q = asyncio.Queue()
@ -41,6 +41,7 @@ class ProxyConnectionHandler(server.ConnectionHandler):
x.reply = controller.DummyReply() x.reply = controller.DummyReply()
self.event_queue.put(("log", x)) self.event_queue.put(("log", x))
class Proxyserver: class Proxyserver:
""" """
This addon runs the actual proxy server. This addon runs the actual proxy server.
@ -51,6 +52,7 @@ class Proxyserver:
self.loop = asyncio.get_event_loop() self.loop = asyncio.get_event_loop()
self.listen_port = None self.listen_port = None
self.event_queue = None self.event_queue = None
self.options = ctx.options
self._lock = asyncio.Lock() self._lock = asyncio.Lock()
def running(self): def running(self):
@ -73,7 +75,13 @@ class Proxyserver:
) )
async def handle_connection(self, r, w): 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): def configure(self, updated):
if "listen_port" in 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: class Connection:
@ -37,11 +39,16 @@ class Context:
client: Client client: Client
server: Optional[Server] server: Optional[Server]
options: Options
layers: List["mitmproxy.proxy2.layer.Layer"]
def __init__( def __init__(
self, self,
client: Client, client: Client,
server: Optional[Server], server: Optional[Server],
options: Options,
) -> None: ) -> None:
self.client = client self.client = client
self.server = server self.server = server
self.options = options
self.layers = []

View File

@ -11,6 +11,7 @@ import asyncio
import socket import socket
import typing import typing
from mitmproxy import options
from mitmproxy.proxy2 import events, commands from mitmproxy.proxy2 import events, commands
from mitmproxy.proxy2.context import Client, Context, Connection from mitmproxy.proxy2.context import Client, Context, Connection
from mitmproxy.proxy2.layers.modes import ReverseProxy from mitmproxy.proxy2.layers.modes import ReverseProxy
@ -24,11 +25,11 @@ class StreamIO(typing.NamedTuple):
class ConnectionHandler(metaclass=abc.ABCMeta): class ConnectionHandler(metaclass=abc.ABCMeta):
transports: typing.MutableMapping[Connection, StreamIO] transports: typing.MutableMapping[Connection, StreamIO]
def __init__(self, reader, writer): def __init__(self, reader, writer, options):
addr = writer.get_extra_info('peername') addr = writer.get_extra_info('peername')
self.client = Client(addr) 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", 443))
self.layer = ReverseProxy(self.context, ("localhost", 8000)) self.layer = ReverseProxy(self.context, ("localhost", 8000))
@ -134,7 +135,7 @@ if __name__ == "__main__":
async def handle(reader, writer): 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) coro = asyncio.start_server(handle, '127.0.0.1', 8080, loop=loop)

View File

@ -1,5 +1,6 @@
import pytest import pytest
from mitmproxy import options
from mitmproxy.proxy2 import context from mitmproxy.proxy2 import context
@ -7,5 +8,6 @@ from mitmproxy.proxy2 import context
def tctx(): def tctx():
return context.Context( return context.Context(
context.Client("client"), 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 # 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 # complete), so we need to signal if someone forgets to assert and playbooks aren't
# evaluated. # 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!") raise RuntimeError("Unfinished playbook!")
def fork(self): def fork(self):