[sans-io] improve folder structure

mitmproxy/addons/proxyserver.py

@@ -3,9 +3,9 @@ import queue
 import threading
 
 from mitmproxy import ctx, controller, log
-from mitmproxy.proxy.protocol2 import commands
-from mitmproxy.proxy.protocol2 import events
-from mitmproxy.proxy.protocol2.server import server_async
+from mitmproxy.proxy2 import commands
+from mitmproxy.proxy2 import events
+from mitmproxy.proxy2 import server
 
 
 class AsyncReply(controller.Reply):
@@ -19,7 +19,7 @@ class AsyncReply(controller.Reply):
         self.submit(self.q.get_nowait())
 
 
-class ProxyConnectionHandler(server_async.ConnectionHandler):
+class ProxyConnectionHandler(server.ConnectionHandler):
     event_queue: queue.Queue
     loop: asyncio.AbstractEventLoop
 

mitmproxy/proxy/protocol2/server/server.py

@@ -1,109 +0,0 @@
-# This is outdated, only the async version is kept up to date.
-"""
-Minimal server implementation based on https://docs.python.org/3/library/selectors.html#examples.
-May be worth to replace this with something asyncio-based to overcome the issues outlined by the
-FIXMEs below.
-"""
-import functools
-import selectors
-import socket
-from typing import MutableMapping
-
-from mitmproxy.proxy.protocol2 import events, commands
-from mitmproxy.proxy.protocol2.context import Connection
-from mitmproxy.proxy.protocol2.context import Context, Client
-from mitmproxy.proxy.protocol2.events import Event
-from mitmproxy.proxy.protocol2.layer import Layer
-from mitmproxy.proxy.protocol2.reverse_proxy import ReverseProxy
-
-
-class ConnectionHandler:
-    def __init__(self, sel: selectors.BaseSelector, sock: socket.socket, addr: tuple) -> None:
-        self.connections = {}  # type: MutableMapping[Connection, socket.socket]
-        self.sel = sel
-        self.sock = sock
-        self.addr = addr
-
-        sock.setblocking(False)
-
-        client = Client(addr)
-        context = Context(client)
-
-        layer = ReverseProxy(context, ("example.com", 80))
-
-        # self.server_event(layer, commands.OpenConnection(client))
-
-        callback = functools.partial(self.read, layer, client)
-
-        self.sel.register(
-            sock,
-            selectors.EVENT_READ,
-            callback
-        )
-
-        self.connections[client] = sock
-
-    def read(self, layer: Layer, conn: Connection, sock: socket.socket, mask: int):
-        data = sock.recv(4096)
-        if data:
-            self.server_event(layer, events.DataReceived(conn, data))
-        else:
-            # TODO: Needs proper teardown.
-            self.sel.unregister(sock)
-            sock.close()
-            self.server_event(layer, events.ConnectionClosed(conn))
-
-    def server_event(self, layer: Layer, event: Event):
-        print(">>", event)
-        layer_events = layer.handle_event(event)
-        for event in layer_events:
-            print("<<", event)
-            if isinstance(event, commands.OpenConnection):
-                # FIXME: This is blocking!
-                sock = socket.create_connection(event.connection.address)
-                sock.setblocking(False)
-                event.connection.connected = True
-                layer_events.send(42)
-                callback = functools.partial(self.read, layer, event.connection)
-                self.sel.register(
-                    sock,
-                    selectors.EVENT_READ,
-                    callback
-                )
-                self.connections[event.connection] = sock
-            elif isinstance(event, commands.SendData):
-                # FIXME: This may fail.
-                self.connections[event.connection].sendall(event.data)
-            else:
-                raise NotImplementedError("Unexpected event: {}".format(event))
-
-
-class TCPServer:
-    def __init__(self, address):
-        self.connections = set()
-        self.server_sock = socket.socket()
-        self.server_sock.bind(address)
-        self.server_sock.listen(100)
-        self.server_sock.setblocking(False)
-
-        self.sel = selectors.DefaultSelector()
-        self.sel.register(self.server_sock, selectors.EVENT_READ, self.accept)
-
-    def accept(self, server_sock, mask):
-        sock, addr = server_sock.accept()  # Should be ready
-        connection = ConnectionHandler(self.sel, sock, addr)
-        self.connections.add(connection)
-
-    def run(self):
-        while True:
-            self.tick(None)
-
-    def tick(self, timeout):
-        for key, mask in self.sel.select(timeout):
-            callback = key.data
-            callback(key.fileobj, mask)
-
-
-if __name__ == '__main__':
-    s = TCPServer(('', 8080))
-    s.run()

mitmproxy/proxy2/commands.py

@@ -9,7 +9,7 @@ The counterpart to commands are events.
 import typing
 
 from mitmproxy import log
-from mitmproxy.proxy.protocol2.context import Connection
+from mitmproxy.proxy2.context import Connection
 
 
 class Command:

mitmproxy/proxy2/events.py

@@ -5,8 +5,8 @@ The counterpart to events are commands.
 """
 import typing
 
-from mitmproxy.proxy.protocol2 import commands
-from mitmproxy.proxy.protocol2.context import Connection
+from mitmproxy.proxy2 import commands
+from mitmproxy.proxy2.context import Connection
 
 
 class Event:

mitmproxy/proxy2/layer.py

@@ -5,10 +5,8 @@ import collections
 import typing
 from abc import ABCMeta, abstractmethod
 
-from mitmproxy.proxy.protocol2 import commands, events
-from mitmproxy.proxy.protocol2.context import Context
-from mitmproxy.proxy.protocol2.events import Event
-from mitmproxy.proxy.protocol2.utils import expect
+from mitmproxy.proxy2 import commands, events
+from mitmproxy.proxy2.context import Context
 
 
 class Paused(typing.NamedTuple):
@@ -32,12 +30,12 @@ class Layer(metaclass=ABCMeta):
         pass  # print(*args)
 
     @abstractmethod
-    def _handle_event(self, event: Event) -> commands.TCommandGenerator:
+    def _handle_event(self, event: events.Event) -> commands.TCommandGenerator:
         """Handle a proxy server event"""
         if False:
             yield None
 
-    def handle_event(self, event: Event) -> commands.TCommandGenerator:
+    def handle_event(self, event: events. Event) -> commands.TCommandGenerator:
         if self._paused:
             # did we just receive the reply we were waiting for?
             pause_finished = (

mitmproxy/proxy2/layers/http.py

@@ -1,16 +1,16 @@
-import typing
-from warnings import warn
-from unittest import mock
-
 import sys
+import typing
+from unittest import mock
+from warnings import warn
 
 import h11
 from mitmproxy.net import http
-from mitmproxy.proxy.protocol2 import events, commands, websocket
-from mitmproxy.proxy.protocol2.context import ClientServerContext
-from mitmproxy.proxy.protocol2.layer import Layer
-from mitmproxy.proxy.protocol2.utils import expect
 from mitmproxy.net import websockets
+from mitmproxy.proxy2 import events, commands
+from mitmproxy.proxy2.context import ClientServerContext
+from mitmproxy.proxy2.layer import Layer
+from mitmproxy.proxy2.layers import websocket
+from mitmproxy.proxy2.utils import expect
 
 
 class HTTPLayer(Layer):
@@ -25,7 +25,7 @@ class HTTPLayer(Layer):
     def __init__(self, context: ClientServerContext):
         super().__init__(context)
         self.state = self.read_request_headers
-        self.flow = mock.Mock("mitmproxy.http.HTTPFlow")
+        self.flow = mock.Mock()
         self.client_conn = h11.Connection(h11.SERVER)
         self.server_conn = h11.Connection(h11.CLIENT)
 

mitmproxy/proxy2/layers/modes.py

@@ -1,9 +1,8 @@
-from mitmproxy.proxy.protocol2.commands import TCommandGenerator
-from mitmproxy.proxy.protocol2.context import ClientServerContext, Context, Server
-from mitmproxy.proxy.protocol2.events import Event
-from mitmproxy.proxy.protocol2.http import HTTPLayer
-from mitmproxy.proxy.protocol2.layer import Layer
-from mitmproxy.proxy.protocol2.tcp import TCPLayer
+from mitmproxy.proxy2 import events
+from mitmproxy.proxy2.commands import TCommandGenerator
+from mitmproxy.proxy2.context import Context, Server, ClientServerContext
+from mitmproxy.proxy2.layer import Layer
+from mitmproxy.proxy2.layers.http import HTTPLayer
 
 
 class ReverseProxy(Layer):
@@ -15,5 +14,5 @@ class ReverseProxy(Layer):
         # self.child_layer = TCPLayer(self.child_context, False)
         self.child_layer = HTTPLayer(self.child_context)
 
-    def _handle_event(self, event: Event) -> TCommandGenerator:
+    def _handle_event(self, event: events.Event) -> TCommandGenerator:
         yield from self.child_layer.handle_event(event)

mitmproxy/proxy2/layers/tcp.py

@@ -1,10 +1,8 @@
-import typing
-
 from mitmproxy import tcp, flow
-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
+from mitmproxy.proxy2 import commands, events
+from mitmproxy.proxy2.context import ClientServerContext
+from mitmproxy.proxy2.layer import Layer
+from mitmproxy.proxy2.utils import expect
 
 
 class TCPLayer(Layer):

mitmproxy/proxy2/layers/tls.py

@@ -10,11 +10,11 @@ from OpenSSL import SSL
 
 from mitmproxy.certs import CertStore
 from mitmproxy.proxy.protocol.tls import DEFAULT_CLIENT_CIPHERS
-from mitmproxy.proxy.protocol2 import events, commands
-from mitmproxy.proxy.protocol2.context import ClientServerContext, Connection
-from mitmproxy.proxy.protocol2.layer import Layer
-from mitmproxy.proxy.protocol2.tcp import TCPLayer
-from mitmproxy.proxy.protocol2.utils import expect
+from mitmproxy.proxy2 import events, commands
+from mitmproxy.proxy2.context import ClientServerContext, Connection
+from mitmproxy.proxy2.layer import Layer
+from mitmproxy.proxy2.layers.tcp import TCPLayer
+from mitmproxy.proxy2.utils import expect
 
 
 class TLSLayer(Layer):

mitmproxy/proxy2/layers/websocket.py

@@ -1,8 +1,8 @@
 from mitmproxy import websocket, http, flow
-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
+from mitmproxy.proxy2 import events, commands
+from mitmproxy.proxy2.context import ClientServerContext
+from mitmproxy.proxy2.layer import Layer
+from mitmproxy.proxy2.utils import expect
 from wsproto import connection as wsconn
 from wsproto import events as wsevents
 

mitmproxy/proxy2/server.py

@@ -11,10 +11,9 @@ import asyncio
 import socket
 import typing
 
-from mitmproxy.proxy.protocol2 import events, commands
-from mitmproxy.proxy.protocol2.context import Client, Context
-from mitmproxy.proxy.protocol2.context import Connection
-from mitmproxy.proxy.protocol2.reverse_proxy import ReverseProxy
+from mitmproxy.proxy2 import events, commands
+from mitmproxy.proxy2.context import Client, Context, Connection
+from mitmproxy.proxy2.layers.modes import ReverseProxy
 
 
 class StreamIO(typing.NamedTuple):

mitmproxy/proxy2/test/conftest.py

@@ -1,6 +1,6 @@
 import pytest
 
-from mitmproxy.proxy.protocol2 import context
+from mitmproxy.proxy2 import context
 
 
 @pytest.fixture

mitmproxy/proxy2/test/test_tcp.py

@@ -1,7 +1,6 @@
-from . import tutils
-from .. import commands
-from .. import events
-from .. import tcp
+from mitmproxy.proxy2 import commands, events
+from mitmproxy.proxy2.layers import tcp
+from mitmproxy.proxy2.test import tutils
 
 
 def test_open_connection(tctx):

mitmproxy/proxy2/test/test_tutils.py

@@ -2,9 +2,9 @@ import typing
 
 import pytest
 
-from mitmproxy.proxy.protocol2 import events, commands
-from mitmproxy.proxy.protocol2.layer import Layer
-from mitmproxy.proxy.protocol2.test import tutils
+from mitmproxy.proxy2 import events, commands
+from mitmproxy.proxy2.layer import Layer
+from mitmproxy.proxy2.test import tutils
 
 
 class TEvent(events.Event):

mitmproxy/proxy2/test/test_websocket.py

@@ -2,10 +2,9 @@ from unittest import mock
 
 import pytest
 
-from . import tutils
-from .. import commands
-from .. import events
-from .. import websocket
+from mitmproxy.proxy2 import commands, events
+from mitmproxy.proxy2.layers import websocket
+from mitmproxy.proxy2.test import tutils
 from mitmproxy.test import tflow
 
 

mitmproxy/proxy2/test/tutils.py

@@ -3,9 +3,9 @@ import difflib
 import itertools
 import typing
 
-from mitmproxy.proxy.protocol2 import commands
-from mitmproxy.proxy.protocol2 import events
-from mitmproxy.proxy.protocol2 import layer
+from mitmproxy.proxy2 import commands
+from mitmproxy.proxy2 import events
+from mitmproxy.proxy2 import layer
 
 TPlaybookEntry = typing.Union[commands.Command, events.Event]
 TPlaybook = typing.List[TPlaybookEntry]

mitmproxy/proxy2/utils.py

@@ -4,7 +4,7 @@ Utility decorators that help build state machines
 import functools
 from typing import Optional
 
-from mitmproxy.proxy.protocol2 import events
+from mitmproxy.proxy2 import events
 
 
 # This is not used at the moment.