Merge branch 'master' into python3.5

mautrix_telegram/user.py

@@ -22,6 +22,7 @@ from telethon.tl.types import *
 from telethon.tl.types.contacts import ContactsNotModified
 from telethon.tl.types import User as TLUser
 from telethon.tl.functions.contacts import GetContactsRequest, SearchRequest
+from mautrix_appservice import MatrixRequestError
 
 from .db import User as DBUser, Message as DBMessage, Contact as DBContact
 from .tgclient import MautrixTelegramClient
@@ -38,17 +39,20 @@ class User:
     by_mxid = {}
     by_tgid = {}
 
-    def __init__(self, mxid, tgid=None, username=None, db_contacts=None, saved_contacts=0):
+    def __init__(self, mxid, tgid=None, username=None, db_contacts=None, saved_contacts=0,
+                 db_portals=None):
         self.mxid = mxid
         self.tgid = tgid
         self.username = username
         self.contacts = []
         self.saved_contacts = saved_contacts
         self.db_contacts = db_contacts
+        self.portals = {}
+        self.db_portals = db_portals
 
         self.command_status = None
         self.connected = False
-        self.client = None
+        self._init_client()
 
         self.is_admin = self.mxid in config.get("bridge.admins", [])
 
@@ -82,6 +86,19 @@ class User:
         else:
             self.contacts = []
 
+    @property
+    def db_portals(self):
+        return [portal.to_db(merge=False) for _, portal in self.portals.items()]
+
+    @db_portals.setter
+    def db_portals(self, portals):
+        if portals:
+            self.portals = {(portal.tgid, portal.tg_receiver):
+                                po.Portal.get_by_tgid(portal.tgid, portal.tg_receiver)
+                            for portal in portals}
+        else:
+            self.portals = {}
+
     def get_input_entity(self, user):
         return user.client.get_input_entity(InputUser(user_id=self.tgid, access_hash=0))
 
@@ -90,7 +107,8 @@ class User:
     def to_db(self):
         return self.db.merge(
             DBUser(mxid=self.mxid, tgid=self.tgid, tg_username=self.username,
-                   contacts=self.db_contacts, saved_contacts=self.saved_contacts))
+                   contacts=self.db_contacts, saved_contacts=self.saved_contacts,
+                   portals=self.db_portals))
 
     def save(self):
         self.to_db()
@@ -99,12 +117,12 @@ class User:
     @classmethod
     def from_db(cls, db_user):
         return User(db_user.mxid, db_user.tgid, db_user.tg_username, db_user.contacts,
-                    db_user.saved_contacts)
+                    db_user.saved_contacts, db_user.portals)
 
     # endregion
     # region Telegram connection management
 
-    async def start(self):
+    def _init_client(self):
         device = f"{platform.system()} {platform.release()}"
         sysversion = MautrixTelegramClient.__version__
         self.client = MautrixTelegramClient(self.mxid,
@@ -115,6 +133,8 @@ class User:
                                             system_version=sysversion,
                                             device_model=device)
         self.client.add_update_handler(self.update_catch)
+
+    async def start(self):
         self.connected = await self.client.connect()
         if self.logged_in:
             asyncio.ensure_future(self.post_login(), loop=self.loop)
@@ -149,7 +169,14 @@ class User:
             self.save()
 
     async def log_out(self):
-        self.connected = False
+        for _, portal in self.portals.items():
+            try:
+                await portal.main_intent.kick(portal.mxid, self.mxid, "Logged out of Telegram.")
+            except MatrixRequestError:
+                pass
+        self.portals = {}
+        self.contacts = []
+        self.save()
         if self.tgid:
             try:
                 del self.by_tgid[self.tgid]
@@ -157,8 +184,12 @@ class User:
                 pass
             self.tgid = None
             self.save()
-        await self.client.log_out()
-        # TODO kick user from portals
+        ok = await self.client.log_out()
+        if not ok:
+            return False
+        self._init_client()
+        await self.start()
+        return True
 
     def _search_local(self, query, max_results=5, min_similarity=45):
         results = []
@@ -201,9 +232,27 @@ class User:
             if invalid:
                 continue
             portal = po.Portal.get_by_entity(entity)
+            self.portals[portal.tgid_full] = portal
             creators.append(portal.create_matrix_room(self, entity, invites=[self.mxid]))
+        self.save()
         await asyncio.gather(*creators, loop=self.loop)
 
+    def register_portal(self, portal):
+        try:
+            if self.portals[portal.tgid_full] == portal:
+                return
+        except KeyError:
+            pass
+        self.portals[portal.tgid_full] = portal
+        self.save()
+
+    def unregister_portal(self, portal):
+        try:
+            del self.portals[portal.tgid_full]
+            self.save()
+        except KeyError:
+            pass
+
     def _hash_contacts(self):
         acc = 0
         for id in sorted([self.saved_contacts] + [contact.id for contact in self.contacts]):
@@ -233,8 +282,8 @@ class User:
             self.log.exception("Failed to handle Telegram update")
 
     async def update(self, update):
-        if isinstance(update, (UpdateShortChatMessage, UpdateShortMessage, UpdateNewMessage,
-                               UpdateNewChannelMessage)):
+        if isinstance(update, (UpdateShortChatMessage, UpdateShortMessage, UpdateNewChannelMessage,
+                               UpdateNewMessage, UpdateEditMessage, UpdateEditChannelMessage)):
             await self.update_message(update)
         elif isinstance(update, (UpdateChatUserTyping, UpdateUserTyping)):
             await self.update_typing(update)
@@ -317,7 +366,8 @@ class User:
         elif isinstance(update, UpdateShortMessage):
             portal = po.Portal.get_by_tgid(update.user_id, self.tgid, "user")
             sender = pu.Puppet.get(self.tgid if update.out else update.user_id)
-        elif isinstance(update, (UpdateNewMessage, UpdateNewChannelMessage)):
+        elif isinstance(update, (UpdateNewMessage, UpdateNewChannelMessage,
+                                 UpdateEditMessage, UpdateEditChannelMessage)):
             update = update.message
             if isinstance(update.to_id, PeerUser) and not update.out:
                 portal = po.Portal.get_by_tgid(update.from_id, peer_type="user",
@@ -331,8 +381,8 @@ class User:
             return update, None, None
         return update, sender, portal
 
-    async def update_message(self, update):
-        update, sender, portal = self.get_message_details(update)
+    def update_message(self, original_update):
+        update, sender, portal = self.get_message_details(original_update)
 
         if isinstance(update, MessageService):
             if isinstance(update.action, MessageActionChannelMigrateFrom):
@@ -341,10 +391,14 @@ class User:
                 return
             self.log.debug("Handling action %s to %s by %d", update.action, portal.tgid_log,
                            sender.id)
-            await portal.handle_telegram_action(self, sender, update.action)
-        else:
-            self.log.debug("Handling message %s to %s by %d", update, portal.tgid_log, sender.tgid)
-            await portal.handle_telegram_message(self, sender, update)
+            return portal.handle_telegram_action(self, sender, update.action)
+
+        if isinstance(original_update, (UpdateEditMessage, UpdateEditChannelMessage)):
+            self.log.debug("Handling edit %s to %s by %d", update, portal.tgid_log, sender.tgid)
+            return portal.handle_telegram_edit(self, sender, update)
+
+        self.log.debug("Handling message %s to %s by %d", update, portal.tgid_log, sender.tgid)
+        return portal.handle_telegram_message(self, sender, update)
 
     # endregion
     # region Class instance lookup