From 3c22ab7bd1ffe59a20720d0615116766eb53b068 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Fri, 14 Oct 2022 13:55:43 +0300
Subject: [PATCH] Try to automatically detect when data export is accepted

---
 mautrix_telegram/abstract_user.py |  7 +++++++
 mautrix_telegram/user.py          | 10 +++++++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/mautrix_telegram/abstract_user.py b/mautrix_telegram/abstract_user.py
index 617003d3..31dc86f5 100644
--- a/mautrix_telegram/abstract_user.py
+++ b/mautrix_telegram/abstract_user.py
@@ -541,6 +541,8 @@ class AbstractUser(ABC):
                 sender = await pu.Puppet.get_by_tgid(self.tgid)
             elif isinstance(update.from_id, (PeerUser, PeerChannel)):
                 sender = await pu.Puppet.get_by_peer(update.from_id)
+            elif isinstance(update.peer_id, PeerUser):
+                sender = await pu.Puppet.get_by_peer(update.peer_id)
             else:
                 sender = None
         else:
@@ -622,6 +624,9 @@ class AbstractUser(ABC):
             )
             await portal.create_matrix_room(self, chan, invites=[self.mxid])
 
+    async def _check_server_notice_edit(self, message: Message) -> None:
+        pass
+
     async def update_message(self, original_update: UpdateMessage) -> None:
         update, sender, portal = await self.get_message_details(original_update)
         if not portal:
@@ -673,6 +678,8 @@ class AbstractUser(ABC):
             return await portal.handle_telegram_action(self, sender, update)
 
         if isinstance(original_update, (UpdateEditMessage, UpdateEditChannelMessage)):
+            if sender and sender.tgid == 777000:
+                await self._check_server_notice_edit(update)
             return await portal.handle_telegram_edit(self, sender, update)
         return await portal.handle_telegram_message(self, sender, update)
 
diff --git a/mautrix_telegram/user.py b/mautrix_telegram/user.py
index c0b7937d..9aa77afe 100644
--- a/mautrix_telegram/user.py
+++ b/mautrix_telegram/user.py
@@ -37,9 +37,9 @@ from telethon.tl.types import (
     Chat,
     ChatForbidden,
     InputUserSelf,
+    Message,
     NotifyPeer,
     PeerUser,
-    TypeChat,
     TypeUpdate,
     UpdateFolderPeers,
     UpdateNewChannelMessage,
@@ -403,6 +403,14 @@ class User(DBUser, AbstractUser, BaseUser):
                     self.log.exception("Error in takeout backfill loop, retrying in an hour")
                     await asyncio.sleep(3600)
 
+    async def _check_server_notice_edit(self, message: Message) -> None:
+        if "Data export request" in message.message and "Accepted" in message.message:
+            self.log.debug(
+                f"Received an edit to message {message.id} that looks like the data export"
+                " was accepted, marking takeout as retriable"
+            )
+            self.takeout_retry_immediate.set()
+
     async def _takeout_and_backfill(self, first_req: Backfill, first_attempt: bool = True) -> None:
         self.takeout_retry_immediate.clear()
         self.takeout_requested = True