# Source code for asyncssh.listener

    # Copyright (c) 2013-2021 by Ron Frederick <ronf@timeheart.net> and others.
    #
    # This program and the accompanying materials are made available under
    # the terms of the Eclipse Public License v2.0 which accompanies this
    # distribution and is available at:
    #
    #     http://www.eclipse.org/legal/epl-2.0/
    #
    # This program may also be made available under the following secondary
    # licenses when the conditions for such availability set forth in the
    # Eclipse Public License v2.0 are satisfied:
    #
    #    GNU General Public License, Version 2.0, or any later versions of
    #    that license
    #
    # SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-or-later
    #
    # Contributors:
    #     Ron Frederick - initial implementation, API, and documentation

    """SSH listeners"""

    import asyncio
    import errno
    import socket
    from types import TracebackType
    from typing import TYPE_CHECKING, AnyStr, Callable, Generic, List, Optional
    from typing import Sequence, Tuple, Type, Union

    from .forward import SSHForwarderCoro
    from .forward import SSHLocalPortForwarder, SSHLocalPathForwarder
    from .misc import HostPort, MaybeAwait
    from .session import SSHTCPSession, SSHUNIXSession
    from .socks import SSHSOCKSForwarder

    if TYPE_CHECKING:
        # pylint: disable=cyclic-import
        from .channel import SSHTCPChannel, SSHUNIXChannel
        from .connection import SSHConnection, SSHClientConnection

    _LocalListenerFactory = Callable[[], asyncio.BaseProtocol]

    ListenKey = Union[HostPort, str]

    TCPListenerFactory = Callable[[str, int], MaybeAwait[SSHTCPSession[AnyStr]]]
    UNIXListenerFactory = Callable[[], MaybeAwait[SSHUNIXSession[AnyStr]]]

    class SSHListener:
        """SSH listener for inbound connections"""

        def __init__(self) -> None:
            self._tunnel: Optional['SSHConnection'] = None

        async def __aenter__(self) -> 'SSHListener':
            return self

        async def __aexit__(self, _exc_type: Optional[Type[BaseException]],
                            _exc_value: Optional[BaseException],
                            _traceback: Optional[TracebackType]) -> bool:
            self.close()
            await self.wait_closed()
            return False

        def get_port(self) -> int:
            """Return the port number being listened on

               This method returns the port number that the remote listener
               was bound to. When the requested remote listening port is `0`
               to indicate a dynamic port, this method can be called to
               determine what listening port was selected. This function
               only applies to TCP listeners.

               :returns: The port number being listened on

            """

            # pylint: disable=no-self-use

            return 0

        def set_tunnel(self, tunnel: 'SSHConnection') -> None:
            """Set tunnel associated with listener"""

            self._tunnel = tunnel

        def close(self) -> None:
            """Stop listening for new connections

               This method can be called to stop listening for connections.
               Existing connections will remain open.

            """

            if self._tunnel:
                self._tunnel.close()

        async def wait_closed(self) -> None:
            """Wait for the listener to close

               This method is a coroutine which waits for the associated
               listeners to be closed.

            """

            if self._tunnel:
                await self._tunnel.wait_closed()
                self._tunnel = None

    class SSHClientListener(SSHListener):
        """Client listener used to accept inbound forwarded connections"""

        def __init__(self, conn: 'SSHClientConnection', encoding: Optional[str],
                     errors: str, window: int, max_pktsize: int):
            super().__init__()

            self._conn: Optional['SSHClientConnection'] = conn
            self._encoding = encoding
            self._errors = errors
            self._window = window
            self._max_pktsize = max_pktsize
            self._close_event = asyncio.Event()

        async def _close(self) -> None:
            """Close this listener"""

            self._close_event.set()
            self._conn = None

        def close(self) -> None:
            """Close this listener asynchronously"""

            super().close()

            if self._conn:
                self._conn.create_task(self._close())

        async def wait_closed(self) -> None:
            """Wait for this listener to finish closing"""

            await super().wait_closed()

            await self._close_event.wait()

    class SSHTCPClientListener(SSHClientListener, Generic[AnyStr]):
        """Client listener used to accept inbound forwarded TCP connections"""

        def __init__(self, conn: 'SSHClientConnection',
                     session_factory: TCPListenerFactory[AnyStr],
                     listen_host: str, listen_port: int,
                     encoding: Optional[str], errors: str,
                     window: int, max_pktsize: int):
            super().__init__(conn, encoding, errors, window, max_pktsize)

            self._session_factory: TCPListenerFactory[AnyStr] = session_factory
            self._listen_host = listen_host
            self._listen_port = listen_port

        async def _close(self) -> None:
            """Close this listener"""

            if self._conn: # pragma: no branch
                await self._conn.close_client_tcp_listener(self._listen_host,
                                                           self._listen_port)

            await super()._close()

        def process_connection(self, orig_host: str, orig_port: int) -> \
                Tuple['SSHTCPChannel[AnyStr]',
                      MaybeAwait[SSHTCPSession[AnyStr]]]:
            """Process a forwarded TCP connection"""

            assert self._conn is not None

            chan = self._conn.create_tcp_channel(self._encoding, self._errors,
                                                 self._window, self._max_pktsize)

            chan.set_inbound_peer_names(self._listen_host, self._listen_port,
                                        orig_host, orig_port)

            return chan, self._session_factory(orig_host, orig_port)

        def get_port(self) -> int:
            """Return the port number being listened on"""

            return self._listen_port

    class SSHUNIXClientListener(SSHClientListener, Generic[AnyStr]):
        """Client listener used to accept inbound forwarded UNIX connections"""

        def __init__(self, conn: 'SSHClientConnection',
                     session_factory: UNIXListenerFactory[AnyStr],
                     listen_path: str, encoding: Optional[str],
                     errors: str, window: int, max_pktsize: int):
            super().__init__(conn, encoding, errors, window, max_pktsize)

            self._session_factory: UNIXListenerFactory[AnyStr] = session_factory
            self._listen_path = listen_path

        async def _close(self) -> None:
            """Close this listener"""

            if self._conn: # pragma: no branch
                await self._conn.close_client_unix_listener(self._listen_path)

            await super()._close()

        def process_connection(self) -> \
                Tuple['SSHUNIXChannel[AnyStr]',
                      MaybeAwait[SSHUNIXSession[AnyStr]]]:
            """Process a forwarded UNIX connection"""

            assert self._conn is not None

            chan = self._conn.create_unix_channel(self._encoding, self._errors,
                                                  self._window, self._max_pktsize)

            chan.set_inbound_peer_names(self._listen_path)

            return chan, self._session_factory()

    class SSHForwardListener(SSHListener):
        """A listener used when forwarding traffic from local ports"""

        def __init__(self, conn: 'SSHConnection',
                     servers: Sequence[asyncio.AbstractServer],
                     listen_key: ListenKey, listen_port: int = 0):
            super().__init__()

            self._conn: Optional['SSHConnection'] = conn
            self._servers = servers
            self._listen_key = listen_key
            self._listen_port = listen_port

        def get_port(self) -> int:
            """Return the port number being listened on"""

            return self._listen_port

        def close(self) -> None:
            """Close this listener"""

            if self._conn: # pragma: no branch
                self._conn.close_forward_listener(self._listen_key)

                for server in self._servers:
                    server.close()

                self._conn = None

        async def wait_closed(self) -> None:
            """Wait for this listener to finish closing"""

            await super().wait_closed()

            for server in self._servers:
                await server.wait_closed()

            self._servers = []

    async def create_tcp_local_listener(
            conn: 'SSHConnection', loop: asyncio.AbstractEventLoop,
            protocol_factory: _LocalListenerFactory, listen_host: str,
            listen_port: int) -> 'SSHForwardListener':
        """Create a listener to forward traffic from a local TCP port over SSH"""

        addrinfo = await loop.getaddrinfo(listen_host or None, listen_port,
                                          family=socket.AF_UNSPEC,
                                          type=socket.SOCK_STREAM,
                                          flags=socket.AI_PASSIVE)

        if not addrinfo: # pragma: no cover
            raise OSError('getaddrinfo() returned empty list')

        servers: List[asyncio.AbstractServer] = []

        for family, socktype, proto, _, sa in addrinfo:
            try:
                sock = socket.socket(family, socktype, proto)
            except OSError: # pragma: no cover
                continue

            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True)

            if family == socket.AF_INET6:
                try:
                    sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, True)
                except AttributeError: # pragma: no cover
                    pass

            if sa[1] == 0:
                sa = sa[:1] + (listen_port,) + sa[2:] # type: ignore

            try:
                sock.bind(sa)
            except (OSError, OverflowError) as exc:
                sock.close()

                for server in servers:
                    server.close()

                if isinstance(exc, OverflowError): # pragma: no cover
                    exc.errno = errno.EOVERFLOW # type: ignore
                    exc.strerror = str(exc) # type: ignore

                # pylint: disable=no-member
                raise OSError(exc.errno, 'error while attempting ' # type: ignore
                              'to bind on address %r: %s' %
                              (sa, exc.strerror)) from None # type: ignore

            if listen_port == 0:
                listen_port = sock.getsockname()[1]
                conn.logger.debug1('Assigning dynamic port %d', listen_port)

            server = await loop.create_server(protocol_factory, sock=sock)
            servers.append(server)

        listen_key = listen_host or '', listen_port
        return SSHForwardListener(conn, servers, listen_key, listen_port)

    async def create_tcp_forward_listener(conn: 'SSHConnection',
                                          loop: asyncio.AbstractEventLoop,
                                          coro: SSHForwarderCoro, listen_host: str,
                                          listen_port: int) -> \
            'SSHForwardListener':
        """Create a listener to forward traffic from a local TCP port over SSH"""

        def protocol_factory() -> asyncio.BaseProtocol:
            """Start a port forwarder for each new local connection"""

            return SSHLocalPortForwarder(conn, coro)

        return await create_tcp_local_listener(conn, loop, protocol_factory,
                                               listen_host, listen_port)

    async def create_unix_forward_listener(conn: 'SSHConnection',
                                           loop: asyncio.AbstractEventLoop,
                                           coro: SSHForwarderCoro,
                                           listen_path: str) -> \
            'SSHForwardListener':
        """Create a listener to forward a local UNIX domain socket over SSH"""

        def protocol_factory() -> asyncio.BaseProtocol:
            """Start a path forwarder for each new local connection"""

            return SSHLocalPathForwarder(conn, coro)

        server = await loop.create_unix_server(protocol_factory, listen_path)

        return SSHForwardListener(conn, [server], listen_path)

    async def create_socks_listener(conn: 'SSHConnection',
                                    loop: asyncio.AbstractEventLoop,
                                    coro: SSHForwarderCoro, listen_host: str,
                                    listen_port: int) -> SSHForwardListener:
        """Create a SOCKS listener to forward traffic over SSH"""

        def protocol_factory() -> asyncio.BaseProtocol:
            """Start a port forwarder for each new SOCKS connection"""

            return SSHSOCKSForwarder(conn, coro)

        return await create_tcp_local_listener(conn, loop, protocol_factory,
                                               listen_host, listen_port)