<commit_before>import asyncio
import logging
import time
import uuid

from collections import defaultdict
import websockets

from . import utils
from .protocol import UDP, TCP, WS
from .contact import Contact

LOG = logging.getLogger(__name__)


class Peer:
    def __init__(self, peer_addr, app, *, loop=None):
        self.peer_addr = peer_addr
        self.registered = {}
        self.subscriber = defaultdict(dict)
        self._app = app
        self._protocol = None
        self._loop = loop
        self._dialogs = {}
        self._connected_future = asyncio.Future(loop=loop)
        self._disconnected_future = asyncio.Future(loop=loop)

    @property
    def dialogs(self):
        return self._dialogs

    async def close(self):
        if self._protocol is not None:
            LOG.debug('Closing connection for %s', self)
            self._protocol.transport.close()
            await self._disconnected_future

    def send_message(self, msg):
        self._protocol.send_message(msg, addr=self.peer_addr)

    def _create_dialog(self, method, from_details, to_details, contact_details=None, password=<PASSWORD>, call_id=None,
                       headers=None, payload=None, cseq=0, inbound=False):

        if not call_id:
            call_id = str(uuid.uuid4())

        if not contact_details:
            host, port = self.local_addr

            # No way to get the public local addr in UDP. Allow an override or select the From host
            # Maybe with https://bugs.python.org/issue31203
            if self._app.defaults['override_contact_host']:
                host = self._app.defaults['override_contact_host']
            elif host == '0.0.0.0' or host.startswith('127.'):
                host = from_details['uri']['host']

            contact_details = Contact(
                {
                    'uri': 'sip:{username}@{host_and_port};transport={protocol}'.format(
                        username=from_details['uri']['user'],
                        host_and_port=utils.format_host_and_port(host, port),
                        protocol=type(self._protocol).__name__.lower()
                    )
                }
            )

        dialog = self._app.dialog_factory(
            method=method,
            app=self._app,
            from_details=from_details,
            to_details=to_details,
            contact_details=contact_details,
            call_id=call_id,
            peer=self,
            password=password,
            headers=headers,
            payload=payload,
            cseq=cseq,
            inbound=inbound,
        )
        LOG.debug('Creating: %s', dialog)
        self._dialogs[call_id] = dialog
        return dialog

    async def subscribe(self, from_details, to_details, contact_details=None, password=<PASSWORD>, call_id=None, cseq=0,
                        expires=3600):
        dialog = self._create_dialog(method="SUBSCRIBE",
                                     from_details=from_details,
                                     to_details=to_details,
                                     contact_details=contact_details,
                                     password=password,
                                     call_id=call_id, cseq=cseq)
        try:
            response = await dialog.start(expires=expires)
            dialog.status_code = response.status_code
            dialog.status_message = response.status_message
            return dialog
        except asyncio.CancelledError:
            dialog.cancel()
            raise

    async def register(self, from_details, to_details, contact_details=None, password=<PASSWORD>, call_id=None, cseq=0,
                       expires=3600):
        dialog = self._create_dialog(method="REGISTER",
                                     from_details=from_details,
                                     to_details=to_details,
                                     contact_details=contact_details,
                                     password=password,
                                     call_id=call_id, cseq=cseq)
        try:
            response = await dialog.start(expires=expires)
            dialog.status_code = response.status_code
            dialog.status_message = response.status_message
            return dialog
        except asyncio.CancelledError:
            dialog.cancel()
            raise

    async def invite(self, from_details, to_details, contact_details=None, password=<PASSWORD>, call_id=None, headers=None,
                     payload=None, cseq=0):

        if not call_id:
            call_id = str(uuid.uuid4())

        if not contact_details:
            host, port = self.local_addr

            # No way to get the public local addr in UDP. Allow an override or select the From host
            # Maybe with https://bugs.python.org/issue31203
            if self._app.defaults['override_contact_host']:
                host = self._app.defaults['override_contact_host']
            elif host == '0.0.0.0' or host.startswith('127.'):
                host = from_details['uri']['host']

            contact_details = Contact(
                {
                    'uri': 'sip:{username}@{host_and_port};transport={protocol}'.format(
                        username=from_details['uri']['user'],
                        host_and_port=utils.format_host_and_port(host, port),
                        protocol=type(self._protocol).__name__.lower()
                    )
                }
            )

        from .dialog import InviteDialog
        dialog = InviteDialog(
            app=self._app,
            from_details=from_details,
            to_details=to_details,
            call_id=call_id,
            peer=self,
            contact_details=contact_details,
            headers=headers,
            payload=payload,
            password=<PASSWORD>,
            cseq=cseq
        )
        self._dialogs[call_id] = dialog
        await dialog.start()
        return dialog

    async def proxy_request(self, dialog, msg, timeout=5):
        if msg.method == 'ACK':
            self.send_message(msg)
            return

        proxy_dialog = self._dialogs.get(dialog.call_id)
        if not proxy_dialog:
            proxy_dialog = self._create_dialog(
                method=msg.method,
                from_details=dialog.from_details,
                to_details=dialog.to_details,
                call_id=dialog.call_id,
            )
        elif msg.cseq in proxy_dialog.transactions[msg.method]:
            proxy_dialog.transactions[msg.method][msg.cseq].retransmit()
            return

        if isinstance(msg.headers['Via'], str):
            msg.headers['Via'] = [msg.headers['Via']]

        host, port = self.local_addr
        if self._app.defaults['override_contact_host']:
            host = self._app.defaults['override_contact_host']

        msg.headers['Via'].insert(0, 'SIP/2.0/%(protocol)s {host}:{port};branch={branch}'.format(
                host=host,
                port=port,
                branch=utils.gen_branch(10)
            )
        )

        async for response in proxy_dialog.start_proxy_transaction(msg, timeout=timeout):
            yield response

        proxy_dialog._maybe_close(msg)

    def _bookkeeping(self, msg, call_id):
        if msg.method not in ('REGISTER', 'SUBSCRIBE'):
            return

        expires = int(msg.headers.get('Expires', 0))

        if msg.method == 'REGISTER' and expires:
            self.registered[msg.contact_details['uri']['user']] = {
                'expires': time.time() + expires,
                'dialog': call_id
            }
        elif msg.method == 'SUBSCRIBE' and expires:
            self.subscriber[msg.contact_details['uri']['user']][msg.to_details['uri']['user']] = {
                'expires': time.time() + expires,
                'dialog': call_id
            }
        if msg.method == 'REGISTER' and not expires:
            try:
                del self.registered[msg.contact_details['uri']['user']]
            except KeyError:
                pass
        elif msg.method == 'SUBSCRIBE' and not expires:
            try:
                del self.subscriber[msg.contact_details['uri']['user']][msg.to_details['uri']['user']]
            except KeyError:
                pass

    def proxy_response(self, msg):
        msg.headers['Via'].pop(0)
        return self.send_message(msg)

    def _close_dialog(self, call_id):
        try:
            del self._dialogs[call_id]
        except KeyError:
            pass

        to_del = list()
        for user, value in self.registered.items():
            if value['dialog'] == call_id:
                to_del.append(user)
        for user in to_del:
            del self.registered[user]

        to_del = list()
        for user, subscriptions in self.subscriber.items():
            for subscribe, values in subscriptions.items():
                if values['dialog'] == call_id:
                    to_del.append((user, subscribe))

        for v in to_del:
            del self.subscriber[v[0]][v[1]]

    @property
    def protocol(self):
        return type(self._protocol)

    @property
    def connected(self):
        return self._connected_future

    def _connected(self, protocol):
        if not self._connected_future.done():
            self._connected_future.set_result(protocol)

        if self._protocol:
            assert self._protocol == protocol
        else:
            self._protocol = protocol

    def _disconnected(self):
        LOG.debug('Lost connection for %s', self)

        for dialog in self._dialogs.values():
            dialog._close()
        self._dialogs = {}

        self.registered = {}
        self.subscriber = defaultdict(dict)

        self._protocol = None
        self._disconnected_future.set_result(None)

    @property
    def local_addr(self):
        if self._protocol:
            return self._protocol.transport.get_extra_info('sockname')
        else:
            return None, None

    @property
    def contacts(self):
        for contact in self.registered:
            yield contact
        for contact in self.subscriber:
            yield contact

    def __repr__(self):
        return '<{0} {1[0]}:{1[1]} {2}, local_addr={3[0]}:{3[1]}>'.format(
            self.__class__.__name__, self.peer_addr, self.protocol.__name__, self.local_addr)


class BaseConnector:
    def __init__(self, app, *, loop=None):
        self._app = app
        self._loop = loop or asyncio.get_event_loop()

        self._protocols = {}
        self._peers = {}
        self._servers = {}

    async def create_server(self, local_addr, sock):
        return await self._create_server(local_addr, sock)

    async def create_peer(self, peer_addr, local_addr=None):
        try:
            if not local_addr:
                peer = [peer for key, peer in self._peers.items() if key[0] == peer_addr][0]
            else:
                peer = self._peers[(peer_addr, local_addr)]
        except (KeyError, IndexError):
            peer = self._create_peer(peer_addr)
            await self._connect_peer(peer, local_addr)
            LOG.debug('Creating: %s', peer)
            return peer
        else:
            await peer.connected
            return peer

    def _create_peer(self, peer_addr):
        peer = Peer(peer_addr, self._app, loop=self._loop)
        self._peers[(peer_addr, None)] = peer
        return peer

    async def _connect_peer(self, peer, local_addr):
        peer._connected(await self._create_connection(peer_addr=peer.peer_addr, local_addr=local_addr))
        if (peer.peer_addr, None) in self._peers:
            del self._peers[(peer.peer_addr, None)]
        if (peer.peer_addr, peer.local_addr) not in self._peers:
            self._peers[(peer.peer_addr, peer.local_addr)] = peer

    async def get_peer(self, protocol, peer_addr):
        return await self._dispatch(protocol, peer_addr)

    def connection_lost(self, protocol):
        for key, peer in list(self._peers.items()):
            if peer._protocol == protocol:
                peer._disconnected()
                self._peers.pop(key)

        for key, proto in list(self._protocols.items()):
            if proto == protocol:
                self._protocols.pop(key)

    async def close(self):
        for peer in list(self._peers.values()):
            await peer.close()

        for server in self._servers.values():
            server.close()
            await server.wait_closed()
        self._servers = {}

    async def _create_server(self, local_addr, sock):
        raise NotImplementedError()

    async def _create_connection(self, peer_addr, local_addr):
        raise NotImplementedError()

    async def _dispatch(self, protocol, peer_addr):
        raise NotImplementedError()


class TCPConnector(BaseConnector):
    async def _create_server(self, local_addr, sock):
        server = await self._loop.create_server(
            lambda: TCP(app=self._app, loop=self._loop),
            host=local_addr[0],
            port=local_addr[1],
            sock=sock)
        self._servers[local_addr] = server
        return server

    async def _create_connection(self, peer_addr, local_addr):
        try:
            return self._protocols[(peer_addr, local_addr)]
        except KeyError:
            transport, proto = await self._loop.create_connection(
                lambda: TCP(app=self._app, loop=self._loop),
                host=peer_addr[0],
                port=peer_addr[1],
                local_addr=local_addr)
            local_addr = transport.get_extra_info('sockname')
            self._protocols[(peer_addr, local_addr)] = proto
            return proto

    async def _dispatch(self, protocol, _):
        peer_addr = protocol.transport.get_extra_info('peername')
        local_addr = protocol.transport.get_extra_info('sockname')
        if (peer_addr, local_addr) not in self._protocols:
            self._protocols[(peer_addr, local_addr)] = protocol
        return await self.create_peer(peer_addr, local_addr)


class UDPServer:
    """
    Shim to present a unified server interface.
    """
    def __init__(self, transport):
        self.transport = transport

    def close(self):
        self.transport.close()

    async def wait_closed(self):
        pass


class UDPConnector(BaseConnector):
    async def _create_connection(self, peer_addr, local_addr):
        try:
            return self._protocols[(peer_addr, local_addr)]
        except KeyError:
            transport, proto = await self._loop.create_datagram_endpoint(
                lambda: UDP(app=self._app, loop=self._loop),
                local_addr=local_addr,
                remote_addr=peer_addr
            )
            local_addr = transport.get_extra_info('sockname')
            self._protocols[(peer_addr, local_addr)] = proto
            return proto

    async def _create_server(self, local_addr=None, sock=None):
        if sock and local_addr:
            raise ValueError('local_addr and sock are mutually exclusive')
        elif not sock and not local_addr:
            raise ValueError('One of local_addr, sock is mandatory')

        try:
            return self._servers[local_addr]
        except KeyError:
            transport, proto = await self._loop.create_datagram_endpoint(
                lambda: UDP(app=self._app, loop=self._loop),
                sock=sock,
                local_addr=local_addr
            )
            proto_addr = proto.transport.get_extra_info('sockname')
            if sock:
                assert sock.getsockname() == proto_addr
            else:
                assert local_addr == proto_addr

            server = UDPServer(transport)
            self._servers[local_addr] = server
            return server

    async def _dispatch(self, protocol, peer_addr):
        local_addr = protocol.transport.get_extra_info('sockname')
        if (peer_addr, local_addr) not in self._protocols:
            self._protocols[(peer_addr, local_addr)] = protocol
        return await self.create_peer(peer_addr, local_addr)


class WSConnector(BaseConnector):
    async def _create_connection(self, peer_addr, local_addr):
        try:
            return self._protocols[(peer_addr, local_addr)]
        except KeyError:
            websocket = await websockets.connect(peer_addr, subprotocols=['sip'])
            local_addr = (utils.gen_str(12) + '.invalid', None)
            proto = WS(app=self._app, loop=self._loop,
                       local_addr=local_addr,
                       peer_addr=peer_addr,
                       websocket=websocket)
            self._protocols[(peer_addr, local_addr)] = proto
            return proto

    async def _create_server(self, local_addr, sock):
        async def hello(websocket, path):
            proto = WS(app=self._app, loop=self._loop,
                       local_addr=local_addr,
                       peer_addr=websocket.remote_address,
                       websocket=websocket)
            await proto.websocket_pump

        try:
            return self._servers[local_addr]
        except KeyError:
            server = await websockets.serve(hello, local_addr[0], local_addr[1],
                                            subprotocols=['sip'])
            self._servers[local_addr] = server
            return server

    async def _dispatch(self, protocol, peer_addr):
        local_addr = protocol.local_addr
        if (peer_addr, local_addr) not in self._protocols:
            self._protocols[(peer_addr, local_addr)] = protocol
        return await self.create_peer(peer_addr, local_addr)
<commit_msg>Add Dialog.request for starting arbitrary dialogs<commit_after>import asyncio
import logging
import time
import uuid

from collections import defaultdict
import websockets

from . import utils
from .protocol import UDP, TCP, WS
from .contact import Contact

LOG = logging.getLogger(__name__)


class Peer:
    def __init__(self, peer_addr, app, *, loop=None):
        self.peer_addr = peer_addr
        self.registered = {}
        self.subscriber = defaultdict(dict)
        self._app = app
        self._protocol = None
        self._loop = loop
        self._dialogs = {}
        self._connected_future = asyncio.Future(loop=loop)
        self._disconnected_future = asyncio.Future(loop=loop)

    @property
    def dialogs(self):
        return self._dialogs

    async def close(self):
        if self._protocol is not None:
            LOG.debug('Closing connection for %s', self)
            self._protocol.transport.close()
            await self._disconnected_future

    def send_message(self, msg):
        self._protocol.send_message(msg, addr=self.peer_addr)

    def _create_dialog(self, method, from_details, to_details, contact_details=None, password=<PASSWORD>, call_id=None,
                       headers=None, payload=None, cseq=0, inbound=False):

        if not call_id:
            call_id = str(uuid.uuid4())

        if not contact_details:
            host, port = self.local_addr

            # No way to get the public local addr in UDP. Allow an override or select the From host
            # Maybe with https://bugs.python.org/issue31203
            if self._app.defaults['override_contact_host']:
                host = self._app.defaults['override_contact_host']
            elif host == '0.0.0.0' or host.startswith('127.'):
                host = from_details['uri']['host']

            contact_details = Contact(
                {
                    'uri': 'sip:{username}@{host_and_port};transport={protocol}'.format(
                        username=from_details['uri']['user'],
                        host_and_port=utils.format_host_and_port(host, port),
                        protocol=type(self._protocol).__name__.lower()
                    )
                }
            )

        dialog = self._app.dialog_factory(
            method=method,
            app=self._app,
            from_details=from_details,
            to_details=to_details,
            contact_details=contact_details,
            call_id=call_id,
            peer=self,
            password=password,
            headers=headers,
            payload=payload,
            cseq=cseq,
            inbound=inbound,
        )
        LOG.debug('Creating: %s', dialog)
        self._dialogs[call_id] = dialog
        return dialog

    async def request(self, method, from_details, to_details, contact_details=None, password=<PASSWORD>, call_id=None,
                      headers=None, payload=None, cseq=0):
        dialog = self._create_dialog(method=method,
                                     from_details=from_details,
                                     to_details=to_details,
                                     contact_details=contact_details,
                                     headers=headers,
                                     payload=payload,
                                     password=password,
                                     call_id=call_id, cseq=cseq)
        try:
            response = await dialog.start()
            dialog.status_code = response.status_code
            dialog.status_message = response.status_message
            return dialog
        except asyncio.CancelledError:
            dialog.cancel()
            raise

    async def subscribe(self, from_details, to_details, contact_details=None, password=<PASSWORD>, call_id=None, cseq=0,
                        expires=3600):
        dialog = self._create_dialog(method="SUBSCRIBE",
                                     from_details=from_details,
                                     to_details=to_details,
                                     contact_details=contact_details,
                                     password=password,
                                     call_id=call_id, cseq=cseq)
        try:
            response = await dialog.start(expires=expires)
            dialog.status_code = response.status_code
            dialog.status_message = response.status_message
            return dialog
        except asyncio.CancelledError:
            dialog.cancel()
            raise

    async def register(self, from_details, to_details, contact_details=None, password=<PASSWORD>, call_id=None, cseq=0,
                       expires=3600):
        dialog = self._create_dialog(method="REGISTER",
                                     from_details=from_details,
                                     to_details=to_details,
                                     contact_details=contact_details,
                                     password=password,
                                     call_id=call_id, cseq=cseq)
        try:
            response = await dialog.start(expires=expires)
            dialog.status_code = response.status_code
            dialog.status_message = response.status_message
            return dialog
        except asyncio.CancelledError:
            dialog.cancel()
            raise

    async def invite(self, from_details, to_details, contact_details=None, password=<PASSWORD>, call_id=None, headers=None,
                     payload=None, cseq=0):

        if not call_id:
            call_id = str(uuid.uuid4())

        if not contact_details:
            host, port = self.local_addr

            # No way to get the public local addr in UDP. Allow an override or select the From host
            # Maybe with https://bugs.python.org/issue31203
            if self._app.defaults['override_contact_host']:
                host = self._app.defaults['override_contact_host']
            elif host == '0.0.0.0' or host.startswith('127.'):
                host = from_details['uri']['host']

            contact_details = Contact(
                {
                    'uri': 'sip:{username}@{host_and_port};transport={protocol}'.format(
                        username=from_details['uri']['user'],
                        host_and_port=utils.format_host_and_port(host, port),
                        protocol=type(self._protocol).__name__.lower()
                    )
                }
            )

        from .dialog import InviteDialog
        dialog = InviteDialog(
            app=self._app,
            from_details=from_details,
            to_details=to_details,
            call_id=call_id,
            peer=self,
            contact_details=contact_details,
            headers=headers,
            payload=payload,
            password=<PASSWORD>,
            cseq=cseq
        )
        self._dialogs[call_id] = dialog
        await dialog.start()
        return dialog

    async def proxy_request(self, dialog, msg, timeout=5):
        if msg.method == 'ACK':
            self.send_message(msg)
            return

        proxy_dialog = self._dialogs.get(dialog.call_id)
        if not proxy_dialog:
            proxy_dialog = self._create_dialog(
                method=msg.method,
                from_details=dialog.from_details,
                to_details=dialog.to_details,
                call_id=dialog.call_id,
            )
        elif msg.cseq in proxy_dialog.transactions[msg.method]:
            proxy_dialog.transactions[msg.method][msg.cseq].retransmit()
            return

        if isinstance(msg.headers['Via'], str):
            msg.headers['Via'] = [msg.headers['Via']]

        host, port = self.local_addr
        if self._app.defaults['override_contact_host']:
            host = self._app.defaults['override_contact_host']

        msg.headers['Via'].insert(0, 'SIP/2.0/%(protocol)s {host}:{port};branch={branch}'.format(
                host=host,
                port=port,
                branch=utils.gen_branch(10)
            )
        )

        async for response in proxy_dialog.start_proxy_transaction(msg, timeout=timeout):
            yield response

        proxy_dialog._maybe_close(msg)

    def _bookkeeping(self, msg, call_id):
        if msg.method not in ('REGISTER', 'SUBSCRIBE'):
            return

        expires = int(msg.headers.get('Expires', 0))

        if msg.method == 'REGISTER' and expires:
            self.registered[msg.contact_details['uri']['user']] = {
                'expires': time.time() + expires,
                'dialog': call_id
            }
        elif msg.method == 'SUBSCRIBE' and expires:
            self.subscriber[msg.contact_details['uri']['user']][msg.to_details['uri']['user']] = {
                'expires': time.time() + expires,
                'dialog': call_id
            }
        if msg.method == 'REGISTER' and not expires:
            try:
                del self.registered[msg.contact_details['uri']['user']]
            except KeyError:
                pass
        elif msg.method == 'SUBSCRIBE' and not expires:
            try:
                del self.subscriber[msg.contact_details['uri']['user']][msg.to_details['uri']['user']]
            except KeyError:
                pass

    def proxy_response(self, msg):
        msg.headers['Via'].pop(0)
        return self.send_message(msg)

    def _close_dialog(self, call_id):
        try:
            del self._dialogs[call_id]
        except KeyError:
            pass

        to_del = list()
        for user, value in self.registered.items():
            if value['dialog'] == call_id:
                to_del.append(user)
        for user in to_del:
            del self.registered[user]

        to_del = list()
        for user, subscriptions in self.subscriber.items():
            for subscribe, values in subscriptions.items():
                if values['dialog'] == call_id:
                    to_del.append((user, subscribe))

        for v in to_del:
            del self.subscriber[v[0]][v[1]]

    @property
    def protocol(self):
        return type(self._protocol)

    @property
    def connected(self):
        return self._connected_future

    def _connected(self, protocol):
        if not self._connected_future.done():
            self._connected_future.set_result(protocol)

        if self._protocol:
            assert self._protocol == protocol
        else:
            self._protocol = protocol

    def _disconnected(self):
        LOG.debug('Lost connection for %s', self)

        for dialog in self._dialogs.values():
            dialog._close()
        self._dialogs = {}

        self.registered = {}
        self.subscriber = defaultdict(dict)

        self._protocol = None
        self._disconnected_future.set_result(None)

    @property
    def local_addr(self):
        if self._protocol:
            return self._protocol.transport.get_extra_info('sockname')
        else:
            return None, None

    @property
    def contacts(self):
        for contact in self.registered:
            yield contact
        for contact in self.subscriber:
            yield contact

    def __repr__(self):
        return '<{0} {1[0]}:{1[1]} {2}, local_addr={3[0]}:{3[1]}>'.format(
            self.__class__.__name__, self.peer_addr, self.protocol.__name__, self.local_addr)


class BaseConnector:
    def __init__(self, app, *, loop=None):
        self._app = app
        self._loop = loop or asyncio.get_event_loop()

        self._protocols = {}
        self._peers = {}
        self._servers = {}

    async def create_server(self, local_addr, sock):
        return await self._create_server(local_addr, sock)

    async def create_peer(self, peer_addr, local_addr=None):
        try:
            if not local_addr:
                peer = [peer for key, peer in self._peers.items() if key[0] == peer_addr][0]
            else:
                peer = self._peers[(peer_addr, local_addr)]
        except (KeyError, IndexError):
            peer = self._create_peer(peer_addr)
            await self._connect_peer(peer, local_addr)
            LOG.debug('Creating: %s', peer)
            return peer
        else:
            await peer.connected
            return peer

    def _create_peer(self, peer_addr):
        peer = Peer(peer_addr, self._app, loop=self._loop)
        self._peers[(peer_addr, None)] = peer
        return peer

    async def _connect_peer(self, peer, local_addr):
        peer._connected(await self._create_connection(peer_addr=peer.peer_addr, local_addr=local_addr))
        if (peer.peer_addr, None) in self._peers:
            del self._peers[(peer.peer_addr, None)]
        if (peer.peer_addr, peer.local_addr) not in self._peers:
            self._peers[(peer.peer_addr, peer.local_addr)] = peer

    async def get_peer(self, protocol, peer_addr):
        return await self._dispatch(protocol, peer_addr)

    def connection_lost(self, protocol):
        for key, peer in list(self._peers.items()):
            if peer._protocol == protocol:
                peer._disconnected()
                self._peers.pop(key)

        for key, proto in list(self._protocols.items()):
            if proto == protocol:
                self._protocols.pop(key)

    async def close(self):
        for peer in list(self._peers.values()):
            await peer.close()

        for server in self._servers.values():
            server.close()
            await server.wait_closed()
        self._servers = {}

    async def _create_server(self, local_addr, sock):
        raise NotImplementedError()

    async def _create_connection(self, peer_addr, local_addr):
        raise NotImplementedError()

    async def _dispatch(self, protocol, peer_addr):
        raise NotImplementedError()


class TCPConnector(BaseConnector):
    async def _create_server(self, local_addr, sock):
        server = await self._loop.create_server(
            lambda: TCP(app=self._app, loop=self._loop),
            host=local_addr[0],
            port=local_addr[1],
            sock=sock)
        self._servers[local_addr] = server
        return server

    async def _create_connection(self, peer_addr, local_addr):
        try:
            return self._protocols[(peer_addr, local_addr)]
        except KeyError:
            transport, proto = await self._loop.create_connection(
                lambda: TCP(app=self._app, loop=self._loop),
                host=peer_addr[0],
                port=peer_addr[1],
                local_addr=local_addr)
            local_addr = transport.get_extra_info('sockname')
            self._protocols[(peer_addr, local_addr)] = proto
            return proto

    async def _dispatch(self, protocol, _):
        peer_addr = protocol.transport.get_extra_info('peername')
        local_addr = protocol.transport.get_extra_info('sockname')
        if (peer_addr, local_addr) not in self._protocols:
            self._protocols[(peer_addr, local_addr)] = protocol
        return await self.create_peer(peer_addr, local_addr)


class UDPServer:
    """
    Shim to present a unified server interface.
    """
    def __init__(self, transport):
        self.transport = transport

    def close(self):
        self.transport.close()

    async def wait_closed(self):
        pass


class UDPConnector(BaseConnector):
    async def _create_connection(self, peer_addr, local_addr):
        try:
            return self._protocols[(peer_addr, local_addr)]
        except KeyError:
            transport, proto = await self._loop.create_datagram_endpoint(
                lambda: UDP(app=self._app, loop=self._loop),
                local_addr=local_addr,
                remote_addr=peer_addr
            )
            local_addr = transport.get_extra_info('sockname')
            self._protocols[(peer_addr, local_addr)] = proto
            return proto

    async def _create_server(self, local_addr=None, sock=None):
        if sock and local_addr:
            raise ValueError('local_addr and sock are mutually exclusive')
        elif not sock and not local_addr:
            raise ValueError('One of local_addr, sock is mandatory')

        try:
            return self._servers[local_addr]
        except KeyError:
            transport, proto = await self._loop.create_datagram_endpoint(
                lambda: UDP(app=self._app, loop=self._loop),
                sock=sock,
                local_addr=local_addr
            )
            proto_addr = proto.transport.get_extra_info('sockname')
            if sock:
                assert sock.getsockname() == proto_addr
            else:
                assert local_addr == proto_addr

            server = UDPServer(transport)
            self._servers[local_addr] = server
            return server

    async def _dispatch(self, protocol, peer_addr):
        local_addr = protocol.transport.get_extra_info('sockname')
        if (peer_addr, local_addr) not in self._protocols:
            self._protocols[(peer_addr, local_addr)] = protocol
        return await self.create_peer(peer_addr, local_addr)


class WSConnector(BaseConnector):
    async def _create_connection(self, peer_addr, local_addr):
        try:
            return self._protocols[(peer_addr, local_addr)]
        except KeyError:
            websocket = await websockets.connect(peer_addr, subprotocols=['sip'])
            local_addr = (utils.gen_str(12) + '.invalid', None)
            proto = WS(app=self._app, loop=self._loop,
                       local_addr=local_addr,
                       peer_addr=peer_addr,
                       websocket=websocket)
            self._protocols[(peer_addr, local_addr)] = proto
            return proto

    async def _create_server(self, local_addr, sock):
        async def hello(websocket, path):
            proto = WS(app=self._app, loop=self._loop,
                       local_addr=local_addr,
                       peer_addr=websocket.remote_address,
                       websocket=websocket)
            await proto.websocket_pump

        try:
            return self._servers[local_addr]
        except KeyError:
            server = await websockets.serve(hello, local_addr[0], local_addr[1],
                                            subprotocols=['sip'])
            self._servers[local_addr] = server
            return server

    async def _dispatch(self, protocol, peer_addr):
        local_addr = protocol.local_addr
        if (peer_addr, local_addr) not in self._protocols:
            self._protocols[(peer_addr, local_addr)] = protocol
        return await self.create_peer(peer_addr, local_addr)

{'fraction_non_alphanumeric': 0.0947809776012358, 'fraction_numerical': 0.003779101842656957, 'mean_word_length': 1.888915451430393, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 16, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 4, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 31, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '296561', 'n_tokens_mistral': 10027, 'n_tokens_neox': 9669, 'n_words': 2071}