From fe9b968ad84f202dbfdaaebd6ce994d6daf2a8c9 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Fri, 8 Dec 2017 23:40:29 +0100
Subject: [PATCH] Add basic update handler

---
 pyrogram/client/client.py   |  7 +++++++
 pyrogram/session/session.py | 11 ++++++++---
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py
index 0e70fb5b..90facb8e 100644
--- a/pyrogram/client/client.py
+++ b/pyrogram/client/client.py
@@ -81,6 +81,12 @@ class Client:
         self.config = None
         self.session = None
 
+        self.update_handler = None
+
+    # TODO: Better update handler
+    def set_update_handler(self, callback):
+        self.update_handler = callback
+
     def send(self, data: Object):
         return self.session.send(data)
 
@@ -249,6 +255,7 @@ class Client:
         self.load_session(self.session_name)
 
         self.session = Session(self.dc_id, self.test_mode, self.auth_key, self.config.api_id)
+        self.session.update_handler = self.update_handler
 
         terms = self.session.start()
 
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
index 1a46ce56..584bd29c 100644
--- a/pyrogram/session/session.py
+++ b/pyrogram/session/session.py
@@ -96,6 +96,8 @@ class Session:
 
         self.is_connected = Event()
 
+        self.update_handler = None
+
         self.total_connections = 0
         self.total_messages = 0
         self.total_bytes = 0
@@ -246,19 +248,22 @@ class Session:
                 self.pending_acks.add(i.body.answer_msg_id)
                 continue
 
+            if isinstance(i.body, types.NewSessionCreated):
+                continue
+
             msg_id = None
 
             if isinstance(i.body, (types.BadMsgNotification, types.BadServerSalt)):
                 msg_id = i.body.bad_msg_id
-
             elif isinstance(i.body, types.RpcResult):
                 msg_id = i.body.req_msg_id
-
             elif isinstance(i.body, types.Pong):
                 msg_id = i.body.msg_id
-
             elif isinstance(i.body, core.FutureSalts):
                 msg_id = i.body.req_msg_id
+            else:
+                if self.update_handler:
+                    self.update_handler(i.body)
 
             if msg_id in self.results:
                 self.results[msg_id].value = getattr(i.body, "result", i.body)