From d7e40a86c66fab538d56faff3e29c110c75511e3 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Sat, 17 Feb 2018 20:23:50 +0200
Subject: [PATCH] Check if bot is still in chat at startup

---
 mautrix_telegram/bot.py       | 53 +++++++++++++++++++++++++++++------
 mautrix_telegram/db.py        |  1 +
 mautrix_telegram/formatter.py |  1 -
 3 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/mautrix_telegram/bot.py b/mautrix_telegram/bot.py
index 0aaa2da8..663d5b45 100644
--- a/mautrix_telegram/bot.py
+++ b/mautrix_telegram/bot.py
@@ -17,6 +17,9 @@
 import logging
 
 from telethon.tl.types import *
+from telethon.errors import ChannelInvalidError, ChannelPrivateError
+from telethon.tl.functions.messages import GetChatsRequest
+from telethon.tl.functions.channels import GetChannelsRequest
 
 from .abstract_user import AbstractUser
 from .db import BotChat
@@ -32,7 +35,7 @@ class Bot(AbstractUser):
         self.token = token
         self.whitelisted = True
         self._init_client()
-        self.chats = {chat.id for chat in BotChat.query.all()}
+        self.chats = {(chat.id, chat.type) for chat in BotChat.query.all()}
 
     async def start(self):
         await super().start()
@@ -45,24 +48,56 @@ class Bot(AbstractUser):
         info = await self.client.get_me()
         self.tgid = info.id
 
+        chat_ids = [id for (id, type) in self.chats if type == "chat"]
+        response = await self.client(GetChatsRequest(chat_ids))
+        for chat in response.chats:
+            if isinstance(chat, ChatForbidden) or chat.left or chat.deactivated:
+                self.remove_chat(chat.id, "chat")
+
+        channel_ids = [InputChannel(id, 0)
+                       for (id, type) in self.chats
+                       if type == "channel"]
+        for id in channel_ids:
+            try:
+                await self.client(GetChannelsRequest([id]))
+            except (ChannelPrivateError, ChannelInvalidError):
+                self.remove_chat(id.channel_id, "channel")
+
+    def add_chat(self, id, type):
+        entry = (id, type)
+        if entry not in self.chats:
+            self.chats.add(entry)
+            self.db.add(BotChat(id=id, type=type))
+            self.db.commit()
+
+    def remove_chat(self, id, type):
+        self.chats.remove((id, type))
+        self.db.delete(BotChat.query.get(id))
+        self.db.commit()
+
     async def update(self, update):
         if not isinstance(update, (UpdateNewMessage, UpdateNewChannelMessage)):
             return
         elif not isinstance(update.message, MessageService):
             return
-        action = update.message.action
+
         to_id = update.message.to_id
-        to_id = to_id.chat_id if isinstance(to_id, PeerChat) else to_id.channel_id
+        if isinstance(to_id, PeerChannel):
+            to_id = to_id.channel_id
+            type = "channel"
+        elif isinstance(to_id, PeerChat):
+            to_id = to_id.chat_id
+            type = "chat"
+        else:
+            return
+
+        action = update.message.action
         if isinstance(action, MessageActionChatAddUser):
             if self.tgid in action.users:
-                self.chats.add(to_id)
-                self.db.add(BotChat(id=to_id))
-                self.db.commit()
+                self.add_chat(to_id, type)
         elif isinstance(action, MessageActionChatDeleteUser):
             if action.user_id == self.tgid:
-                self.chats.remove(to_id)
-                BotChat.query.get(to_id).delete()
-                self.db.commit()
+                self.remove_chat(to_id, type)
 
     def is_in_chat(self, peer_id):
         return peer_id in self.chats
diff --git a/mautrix_telegram/db.py b/mautrix_telegram/db.py
index 21bc3792..fadd4559 100644
--- a/mautrix_telegram/db.py
+++ b/mautrix_telegram/db.py
@@ -100,6 +100,7 @@ class BotChat(Base):
     query = None
     __tablename__ = "bot_chat"
     id = Column(Integer, primary_key=True)
+    type = Column(String, nullable=False)
 
 
 def init(db_session):
diff --git a/mautrix_telegram/formatter.py b/mautrix_telegram/formatter.py
index 8a7ea319..faa03799 100644
--- a/mautrix_telegram/formatter.py
+++ b/mautrix_telegram/formatter.py
@@ -189,7 +189,6 @@ def matrix_reply_to_telegram(content, tg_space, room_id=None):
         reply = content["m.relates_to"]["m.in_reply_to"]
         room_id = room_id or reply["room_id"]
         event_id = reply["event_id"]
-        print(event_id, tg_space, room_id)
         message = DBMessage.query.filter(DBMessage.mxid == event_id,
                                          DBMessage.tg_space == tg_space,
                                          DBMessage.mx_room == room_id).one_or_none()