async def register_client(self, reader, writer):
client = b""
try:
while True:
# ProtocolException will be raised if client disconnect during this
client = await self.get_client(reader, writer)
# check nick
if client == None:
writer.write(
protocol.dump(
protocol.new_status(
410,
"Nick should contains only letters or digits")
).encode("utf8"))
writer.write(b'\n')
continue
self._clients[client.nick] # check this nick in clients dict
writer.write(
protocol.dump(protocol.new_status(409)).encode("utf8"))
writer.write(b'\n')
except KeyError:
writer.write(protocol.dump(protocol.new_status()).encode("utf8"))
writer.write(b'\n')
self._clients[client.nick] = client
logging.info("User {} connected.".format(client.nick))
self.broadcast_clients()
await self.client_handler(client)
async def accept_connection(self, reader, writer):
try:
writer.write(protocol.dump(protocol.new_status()).encode("utf8"))
writer.write(b'\n')
try:
await self.register_client(reader, writer)
except protocol.ProtocolException as pe:
writer.write(
protocol.dump(protocol.new_status(418,
str(pe))).encode("utf8"))
writer.write(b'\n')
except BrokenPipeError:
return
async def get_client(self, reader, writer):
user_data = (await reader.readline()).decode("utf8")
_type, data = protocol.load(user_data)
try:
if not re.match(r"^[A-Za-z0-9]+$", data["nick"]):
return None
client = user.User(data, reader, writer)
except KeyError:
writer.write(
protocol.dump(
protocol.new_status(
418,
'Registration request should contains "nick" field')).
encode("utf8"))
return client
async def registration(self):
error = ""
while error != None:
print("Enter your nick: ", end='')
sys.stdout.flush()
self._nick = await self.ainput()
self._writer.write(protocol.dump(protocol.form_service({"nick": self._nick})).encode("utf8"))
self._writer.write(b'\n')
nick_response = await self._reader.readline()
_type, response = protocol.load(nick_response.decode("utf8"))
error = protocol.check_error(response)
if error != None:
print(protocol.check_error(response))
print("Welcome to chat, {}!\n".format(self._nick))
handle_stdin_task = self._loop.create_task(self.handle_stdin())
handle_message_task = self._loop.create_task(self.handle_message())
def send_payload(self, data, noack=False):
self.prepare_sock()
encoded_payload = protocol.dump(data)
tried = 0
while tried < 3:
try:
tried += 1
self.sock.send(encoded_payload)
self.logger.debug("Sendto {}: {}".format(
self.ip, encoded_payload))
if noack:
return
recv = self.sock.recv(1024)
self.logger.debug("Recv {}: {}".format(self.ip, recv))
answer = protocol.load(recv)
# TODO: check status
return answer
except Exception as e:
print(e)
# self.connect() # FIXME
continue
raise Exception("Fail to connect to server")
def message_handler(self, message):
print("{}: {}\n".format(message["nick"], message["message"]))
sys.stdout.flush()
self._writer.write(protocol.dump(protocol.new_status(message_id=message["messageId"])).encode("utf8"))
self._writer.write(b'\n')
async def handle_stdin(self):
while True:
message = await self.ainput()
self._writer.write(protocol.dump(protocol.new_message(message)).encode("utf8"))
self._writer.write(b'\n')
self._non_deliver_messages += 1
def send_service(self, service):
json_service = protocol.dump(service)
self._send_raw(json_service)
def send_message(self, message):
json_message = protocol.dump(protocol.new_message(message))
self._send_raw(json_message)