add websocket stub

This commit is contained in:
Maximilian Hils 2017-07-12 14:54:36 +02:00
parent 9757403008
commit 2e3704faf1
7 changed files with 68 additions and 12 deletions

View File

@ -83,5 +83,7 @@ class Proxyserver:
self.loop.call_soon_threadsafe(lambda: asyncio.ensure_future(self.start())) self.loop.call_soon_threadsafe(lambda: asyncio.ensure_future(self.start()))
def request(self, flow): def request(self, flow):
print("Changing port...") pass
ctx.options.listen_port += 1 # test live options changes.
# print("Changing port...")
# ctx.options.listen_port += 1

View File

@ -8,6 +8,7 @@ The counterpart to commands are events.
""" """
import typing import typing
from mitmproxy import log
from mitmproxy.proxy.protocol2.context import Connection from mitmproxy.proxy.protocol2.context import Connection
@ -79,4 +80,9 @@ class Hook(Command):
self.data = data self.data = data
class Log(Hook):
def __init__(self, *args, level="info"):
super().__init__("log", log.LogEntry(repr(args), level))
TCommandGenerator = typing.Generator[Command, typing.Any, None] TCommandGenerator = typing.Generator[Command, typing.Any, None]

View File

@ -1,8 +1,10 @@
import typing import typing
from warnings import warn from warnings import warn
import sys
import h11 import h11
from mitmproxy.proxy.protocol2 import events, commands from mitmproxy.proxy.protocol2 import events, commands, websocket
from mitmproxy.proxy.protocol2.context import ClientServerContext from mitmproxy.proxy.protocol2.context import ClientServerContext
from mitmproxy.proxy.protocol2.layer import Layer from mitmproxy.proxy.protocol2.layer import Layer
from mitmproxy.proxy.protocol2.utils import expect from mitmproxy.proxy.protocol2.utils import expect
@ -28,7 +30,7 @@ class HTTPLayer(Layer):
def log_event(orig): def log_event(orig):
def next_event(): def next_event():
e = orig() e = orig()
print(e) print(e, file=sys.__stdout__)
return e return e
return next_event return next_event
@ -56,7 +58,7 @@ class HTTPLayer(Layer):
if event is h11.NEED_DATA: if event is h11.NEED_DATA:
return return
elif isinstance(event, h11.Request): elif isinstance(event, h11.Request):
yield commands.Hook("requestheaders", event) yield commands.Log("requestheaders", event)
if self.client_conn.client_is_waiting_for_100_continue: if self.client_conn.client_is_waiting_for_100_continue:
raise NotImplementedError() raise NotImplementedError()
@ -76,7 +78,7 @@ class HTTPLayer(Layer):
self.flow_events[0].append(event) self.flow_events[0].append(event)
elif isinstance(event, h11.EndOfMessage): elif isinstance(event, h11.EndOfMessage):
self.flow_events[0].append(event) self.flow_events[0].append(event)
yield commands.Hook("request", self.flow_events) yield commands.Log("request", self.flow_events)
yield from self._send_request() yield from self._send_request()
return return
else: else:
@ -84,7 +86,12 @@ class HTTPLayer(Layer):
def _send_request(self): def _send_request(self):
if not self.context.server.connected: if not self.context.server.connected:
yield commands.OpenConnection(self.context.server) err = yield commands.OpenConnection(self.context.server)
if err:
yield commands.Log("error", err)
yield commands.CloseConnection(self.context.client)
self._handle_event = self.done
return
for e in self.flow_events[0]: for e in self.flow_events[0]:
bytes_to_send = self.server_conn.send(e) bytes_to_send = self.server_conn.send(e)
yield commands.SendData(self.context.server, bytes_to_send) yield commands.SendData(self.context.server, bytes_to_send)
@ -95,11 +102,18 @@ class HTTPLayer(Layer):
if event is h11.NEED_DATA: if event is h11.NEED_DATA:
return return
elif isinstance(event, h11.Response): elif isinstance(event, h11.Response):
yield commands.Hook("responseheaders", event) yield commands.Log("responseheaders", event)
self.flow_events[1].append(event) self.flow_events[1].append(event)
self.state = self.read_response_body self.state = self.read_response_body
yield from self.read_response_body() # there may already be further events. yield from self.read_response_body() # there may already be further events.
elif isinstance(event, h11.InformationalResponse):
if event.status_code == 101:
# FIXME: check if this is actually WebSocket
child_layer = websocket.WebsocketLayer(self.context, None)
yield from child_layer.handle_event(events.Start())
self._handle_event = child_layer.handle_event
return
else: else:
raise TypeError(f"Unexpected event: {event}") raise TypeError(f"Unexpected event: {event}")
@ -112,7 +126,7 @@ class HTTPLayer(Layer):
self.flow_events[1].append(event) self.flow_events[1].append(event)
elif isinstance(event, h11.EndOfMessage): elif isinstance(event, h11.EndOfMessage):
self.flow_events[1].append(event) self.flow_events[1].append(event)
yield commands.Hook("response", self.flow_events) yield commands.Log("response", self.flow_events)
yield from self._send_response() yield from self._send_response()
return return
else: else:
@ -128,3 +142,7 @@ class HTTPLayer(Layer):
self.flow_events = [[], []] self.flow_events = [[], []]
self.client_conn.start_next_cycle() self.client_conn.start_next_cycle()
self.server_conn.start_next_cycle() self.server_conn.start_next_cycle()
@expect(events.DataReceived, events.ConnectionClosed)
def done(self, _):
yield from ()

View File

@ -8,6 +8,7 @@ from abc import ABCMeta, abstractmethod
from mitmproxy.proxy.protocol2 import commands, events from mitmproxy.proxy.protocol2 import commands, events
from mitmproxy.proxy.protocol2.context import Context from mitmproxy.proxy.protocol2.context import Context
from mitmproxy.proxy.protocol2.events import Event from mitmproxy.proxy.protocol2.events import Event
from mitmproxy.proxy.protocol2.utils import expect
class Paused(typing.NamedTuple): class Paused(typing.NamedTuple):

View File

@ -12,8 +12,8 @@ class ReverseProxy(Layer):
server = Server(server_addr) server = Server(server_addr)
self.child_context = ClientServerContext(context.client, server) self.child_context = ClientServerContext(context.client, server)
# self.child_layer = TLSLayer(self.child_context, True, True) # self.child_layer = TLSLayer(self.child_context, True, True)
self.child_layer = TCPLayer(self.child_context, False) # self.child_layer = TCPLayer(self.child_context, False)
# self.child_layer = HTTPLayer(self.child_context) self.child_layer = HTTPLayer(self.child_context)
def _handle_event(self, event: Event) -> TCommandGenerator: def _handle_event(self, event: Event) -> TCommandGenerator:
yield from self.child_layer.handle_event(event) yield from self.child_layer.handle_event(event)

View File

@ -32,7 +32,7 @@ class ConnectionHandler(metaclass=abc.ABCMeta):
self.context = Context(self.client) self.context = Context(self.client)
# self.layer = ReverseProxy(self.context, ("localhost", 443)) # self.layer = ReverseProxy(self.context, ("localhost", 443))
self.layer = ReverseProxy(self.context, ("localhost", 80)) self.layer = ReverseProxy(self.context, ("localhost", 8000))
self.transports = { self.transports = {
self.client: StreamIO(reader, writer) self.client: StreamIO(reader, writer)

View File

@ -0,0 +1,29 @@
from mitmproxy import websocket, http
from mitmproxy.proxy.protocol2 import events, commands
from mitmproxy.proxy.protocol2.context import ClientServerContext
from mitmproxy.proxy.protocol2.layer import Layer
from mitmproxy.proxy.protocol2.utils import expect
class WebsocketLayer(Layer):
"""
Simple TCP layer that just relays messages right now.
"""
context: ClientServerContext = None
flow: websocket.WebSocketFlow
def __init__(self, context: ClientServerContext, handshake_flow: http.HTTPFlow):
super().__init__(context)
self.flow = websocket.WebSocketFlow(context.client, context.server, handshake_flow)
assert context.server.connected
@expect(events.Start)
def start(self, _) -> commands.TCommandGenerator:
yield from ()
self._handle_event = self.relay_messages
@expect(events.DataReceived, events.ConnectionClosed)
def relay_messages(self, event: events.Event) -> commands.TCommandGenerator:
raise NotImplementedError()
_handle_event = start