import asyncio
import crypt
import sys
import os

from prompt_toolkit.eventloop.defaults import use_asyncio_event_loop
from prompt_toolkit.patch_stdout import patch_stdout
from prompt_toolkit import prompt

import asyncssh

# Tell prompt_toolkit to use the asyncio event loop.
use_asyncio_event_loop()

class ChatClient:
    _clients = []

    def __init__(self, process):
        self._process = process

    @classmethod
    async def handle_client(cls, process):
        await cls(process).interact()

    def write(self, msg):
        self._process.stdout.write(msg)

    def broadcast(self, msg):
        for client in self._clients:
            if client != self:
                client.write(msg)

    async def interact(self):
        self.write('Welcome to chat!nn')

        self.write('Enter your name: ')
        name = (await self._process.stdin.readline()).rstrip('n')

        self.write('n%d other users are connected.nn' % len(self._clients))

        self._clients.append(self)
        self.broadcast('*** %s has entered chat ***n' % name)

        try:
            ### This works, but it isn't using prompt toolkit...
            #async for line in self._process.stdin:
            #    self.broadcast('%s: %s' % (name, line))
            #    self._process.stdout.write('chat# ')
            while True:
                line = await prompt('chat# ', async_=True) # <--- Should use client's stdout
                self.broadcast('{0}: {1}'.format(name, line))
        except asyncssh.BreakReceived:
            pass

        self.broadcast('*** %s has left chat ***n' % name)
        self._clients.remove(self)

passwords = {'guest': '',                 # guest account with no password
             'user123': 'qV2iEadIGV2rw'   # password of 'secretpw'
            }

class CustomSSHServer(asyncssh.SSHServer):
    def connection_made(self, conn):
        print('SSH connection received from %s.' %
                  conn.get_extra_info('peername')[0])

    def connection_lost(self, exc):
        if exc:
            print('SSH connection error: ' + str(exc), file=sys.stderr)
        else:
            print('SSH connection closed.')

    def begin_auth(self, username):
        # If the user's password is the empty string, no auth is required
        return passwords.get(username) != ''

    def password_auth_supported(self):
        return True

    def validate_password(self, username, password):
        pw = passwords.get(username, '*')
        return crypt.crypt(password, pw) == pw

async def start_server():
    await asyncssh.create_server(CustomSSHServer, '', 8022,
                        server_host_keys=['ssh_host_key'],
                        process_factory=ChatClient.handle_client)

loop = asyncio.get_event_loop()

try:
    loop.run_until_complete(start_server())
except (OSError, asyncssh.Error) as exc:
    sys.exit('Error starting server: ' + str(exc))

loop.run_forever()