From d045bc4368e3eefd45f9a446a96c12fcae9427d7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?poljar=20=28Damir=20Jeli=C4=87=29?= <poljar@termina.org.uk>
Date: Tue, 13 Feb 2018 13:06:30 +0100
Subject: [PATCH] Construct the login event while parsing the json.

---
 matrix/api.py      | 30 +++++++++++++++++-------------
 matrix/events.py   | 41 ++++++++++++++++++++++++++++++++++++++++-
 matrix/messages.py | 26 +++++---------------------
 3 files changed, 62 insertions(+), 35 deletions(-)

diff --git a/matrix/api.py b/matrix/api.py
index 39eb185..a6ecb1d 100644
--- a/matrix/api.py
+++ b/matrix/api.py
@@ -20,6 +20,7 @@ from builtins import str
 import time
 import json
 from enum import Enum, unique
+from functools import partial
 
 try:
     from urllib import quote, urlencode
@@ -232,10 +233,11 @@ class MatrixMessage():
         self.creation_time = time.time()  # type: float
         self.send_time = None             # type: float
         self.receive_time = None          # type: float
+        self.event = None
 
         self.request = request_func(**func_args)
 
-    def decode_body(self):
+    def decode_body(self, server):
         try:
             self.decoded_response = json.loads(
                 self.response.body,
@@ -245,9 +247,6 @@ class MatrixMessage():
         except Exception as error:
             return (False, error)
 
-    def to_event():
-        pass
-
 
 class MatrixLoginMessage(MatrixMessage):
     def __init__(self, client, user, password, device_name, device_id=None):
@@ -267,18 +266,23 @@ class MatrixLoginMessage(MatrixMessage):
             data
         )
 
-    def to_event(self, server):
-        response = self.decoded_response
+    def decode_body(self, server):
+        object_hook = partial(
+            MatrixEvents.MatrixLoginEvent.from_dict,
+            server
+        )
 
         try:
-            access_token = response["access_token"]
-            user_id = response["user_id"]
-
-            return (
-                True,
-                MatrixEvents.MatrixLoginEvent(server, user_id, access_token)
+            event = json.loads(
+                self.response.body,
+                encoding='utf-8',
+                object_hook=object_hook
             )
-        except KeyError as error:
+            self.event = event
+
+            return (True, None)
+
+        except json.decoder.JSONDecodeError as error:
             return (False, error)
 
 
diff --git a/matrix/events.py b/matrix/events.py
index 283dae4..04710fa 100644
--- a/matrix/events.py
+++ b/matrix/events.py
@@ -28,7 +28,24 @@ class MatrixEvent():
         pass
 
 
-class MatrixLoginEvent():
+class MatrixErrorEvent(MatrixEvent):
+    def __init__(self, server, error_message, fatal=False):
+        self.error_message = error_message
+        self.fatal = fatal
+        MatrixEvent.__init__(self, server)
+
+    def execute(self):
+        message = ("{prefix}matrix: {error}").format(
+            prefix=W.prefix("error"),
+            error=self.error_message)
+
+        W.prnt(self.server.server_buffer, message)
+
+        if self.fatal:
+            self.server.disconnect(reconnect=False)
+
+
+class MatrixLoginEvent(MatrixEvent):
     def __init__(self, server, user_id, access_token):
         self.user_id = user_id
         self.access_token = access_token
@@ -40,3 +57,25 @@ class MatrixLoginEvent():
         self.server.client.access_token = self.access_token
 
         self.server.sync()
+
+    @classmethod
+    def from_dict(cls, server, parsed_dict):
+        try:
+            return cls(
+                server,
+                parsed_dict["user_id"],
+                parsed_dict["access_token"]
+            )
+        except KeyError:
+            try:
+                message = "Error logging in: {}.".format(parsed_dict["error"])
+                return MatrixErrorEvent(
+                    server,
+                    message,
+                    fatal=True
+                )
+            except KeyError:
+                return MatrixErrorEvent(
+                    server,
+                    "Error logging in: Invalid JSON response from server.",
+                    fatal=True)
diff --git a/matrix/messages.py b/matrix/messages.py
index 18d7a72..5085aae 100644
--- a/matrix/messages.py
+++ b/matrix/messages.py
@@ -682,14 +682,8 @@ def matrix_handle_message(
     response = message.decoded_response
 
     if message_type is MessageType.LOGIN:
-        ret, event = message.to_event(server)
-
-        if ret:
-            event.execute()
-        else:
-            message = ("{prefix}Error while parsing login response.")
-            W.prnt(server.server_buffer, message)
-            server.disconnect(reconnect=False)
+        event = message.event
+        event.execute()
 
     elif message_type is MessageType.SYNC:
         next_batch = response['next_batch']
@@ -765,7 +759,7 @@ def handle_http_response(server, message):
 
     if ('content-type' in message.response.headers and
             message.response.headers['content-type'] == 'application/json'):
-        ret, error = message.decode_body()
+        ret, error = message.decode_body(server)
 
         if not ret:
             # TODO try to resend the message if decoding has failed?
@@ -791,19 +785,9 @@ def handle_http_response(server, message):
 
     elif status_code == 403:
         if message.type == MessageType.LOGIN:
-            response = message.response.decoded_response
-            reason = ("." if not response or not response["error"] else
-                      ": {r}.".format(r=response["error"]))
+            event = message.event
+            event.execute()
 
-            error_message = ("{prefix}Login error{reason}").format(
-                prefix=W.prefix("error"),
-                reason=reason)
-            server_buffer_prnt(server, error_message)
-
-            W.unhook(server.timer_hook)
-            server.timer_hook = None
-
-            server.disconnect()
         elif message.type == MessageType.TOPIC:
             response = message.decoded_response
             reason = ("." if not response or not response["error"] else