<commit_before>"""
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
from mitmproxy.proxy.protocol2.context import Connection
from mitmproxy.proxy.protocol2.context import Context, Client
from mitmproxy.proxy.protocol2.reverse_proxy import ReverseProxy


class TCPServer:
    def __init__(self, address):
        self.connections = {}  # type: MutableMapping[Connection, socket.socket]

        self.sel = selectors.DefaultSelector()
        sock = socket.socket()
        sock.bind(address)
        sock.listen(100)
        sock.setblocking(False)
        self.sel.register(sock, selectors.EVENT_READ, self.accept)

    def accept(self, server_sock, mask):
        sock, addr = server_sock.accept()  # Should be ready
        sock.setblocking(False)

        client = Client(addr)
        context = Context(client)
        layer = ReverseProxy(context, ("example.com", 80))

        self.receive_event(layer, events.Start())

        callback = functools.partial(self.read, layer, client)

        self.sel.register(
            sock,
            selectors.EVENT_READ,
            callback
        )

        self.connections[client] = sock

    def read(self, layer, conn, sock, mask):
        data = sock.recv(4096)  # Should be ready
        if data:
            self.receive_event(layer, events.ReceiveData(conn, data))
        else:
            # TODO: Needs proper teardown.
            self.sel.unregister(sock)
            sock.close()
            self.receive_event(layer, events.CloseConnection(conn))

    def receive_event(self, layer, event):
        print(">>", event)
        layer_events = layer.handle_event(event)
        for event in layer_events:
            print("<<", event)
            if isinstance(event, events.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, events.SendData):
                # FIXME: This may fail.
                self.connections[event.connection].sendall(event.data)
            else:
                raise NotImplementedError("Unexpected event: {}".format(event))

    def run(self):
        while True:
            events = self.sel.select()
            for key, mask in events:
                callback = key.data
                callback(key.fileobj, mask)


if __name__ == '__main__':
    s = TCPServer(('', 8080))
    s.run()<commit_msg>minor changes to sync server<commit_after>"""
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
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, events.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.ReceiveData(conn, data))
        else:
            # TODO: Needs proper teardown.
            self.sel.unregister(sock)
            sock.close()
            self.server_event(layer, events.CloseConnection(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, events.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, events.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()

"""
1) full async
uh?
2) notifier

3) thread per connection
"""
{'fraction_non_alphanumeric': 0.09571532316630356, 'fraction_numerical': 0.006535947712418301, 'mean_word_length': 2.2313467855466915, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 8, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2537884', 'n_tokens_mistral': 1909, 'n_tokens_neox': 1805, 'n_words': 448}