From 0093acb5789b721d5072eb5f833d845ff953259b Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Sun, 11 Aug 2019 12:01:55 +0300
Subject: [PATCH] Move manhole state to main bridge object

---
 mautrix_telegram/__main__.py         |  8 ++++++-
 mautrix_telegram/commands/manhole.py | 35 ++++++++++++++++------------
 2 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/mautrix_telegram/__main__.py b/mautrix_telegram/__main__.py
index 8ddb27b8..1a7e07d0 100644
--- a/mautrix_telegram/__main__.py
+++ b/mautrix_telegram/__main__.py
@@ -13,8 +13,8 @@
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
+from typing import Optional
 from itertools import chain
-import sys
 
 from alchemysession import AlchemySessionContainer
 
@@ -23,6 +23,7 @@ from mautrix.bridge.db import Base
 
 from .web.provisioning import ProvisioningAPI
 from .web.public import PublicBridgeWebsite
+from .commands.manhole import ManholeState
 from .abstract_user import init as init_abstract_user
 from .bot import Bot, init as init_bot
 from .config import Config
@@ -55,6 +56,7 @@ class TelegramBridge(Bridge):
     config: Config
     session_container: AlchemySessionContainer
     bot: Bot
+    manhole: Optional[ManholeState]
 
     def prepare_db(self) -> None:
         super().prepare_db()
@@ -87,6 +89,7 @@ class TelegramBridge(Bridge):
         context = Context(self.az, self.config, self.loop, self.session_container, self, self.bot)
         self._prepare_website(context)
         self.matrix = context.mx = MatrixHandler(context)
+        self.manhole = None
 
         init_abstract_user(context)
         init_formatter(context)
@@ -100,6 +103,9 @@ class TelegramBridge(Bridge):
         for puppet in Puppet.by_custom_mxid.values():
             puppet.stop()
         self.shutdown_actions = (user.stop() for user in User.by_tgid.values())
+        if self.manhole:
+            self.manhole.close()
+            self.manhole = None
 
 
 TelegramBridge().run()
diff --git a/mautrix_telegram/commands/manhole.py b/mautrix_telegram/commands/manhole.py
index df7e7ab9..da5f7a3d 100644
--- a/mautrix_telegram/commands/manhole.py
+++ b/mautrix_telegram/commands/manhole.py
@@ -18,6 +18,8 @@ import asyncio
 import sys
 import os
 
+from attr import dataclass
+
 from telethon import __version__ as __telethon_version__
 
 from mautrix import __version__ as __mautrix_version__
@@ -28,8 +30,9 @@ from .. import __version__
 from . import command_handler, CommandEvent, SECTION_ADMIN
 
 
-class State:
-    manhole: Optional[asyncio.AbstractServer] = None
+@dataclass
+class ManholeState:
+    server: Optional[asyncio.AbstractServer] = None
     opened_by: Optional[UserID] = None
     close: Optional[Callable[[], None]] = None
 
@@ -41,8 +44,8 @@ async def manhole(evt: CommandEvent) -> None:
         await evt.reply("The manhole has been disabled in the config.")
         return
 
-    if State.manhole:
-        await evt.reply(f"There's an existing manhole opened by {State.opened_by}")
+    if evt.bridge.manhole:
+        await evt.reply(f"There's an existing manhole opened by {evt.bridge.manhole.opened_by}")
         return
 
     from ..portal import Portal
@@ -60,30 +63,32 @@ async def manhole(evt: CommandEvent) -> None:
     path = evt.config["manhole.path"]
 
     evt.log.info(f"{evt.sender.mxid} opened a manhole.")
-    State.manhole, State.close = await start_manhole(path=path, banner=banner, namespace=namespace,
+    server, close = await start_manhole(path=path, banner=banner, namespace=namespace,
                                                      loop=evt.loop)
-    State.opened_by = evt.sender.mxid
+    evt.bridge.manhole = ManholeState(server=server, opened_by=evt.sender.mxid, close=close)
     await evt.reply(f"Opened manhole at unix://{path}")
-    await State.manhole.wait_closed()
+    await server.wait_closed()
+    evt.bridge.manhole = None
     try:
         os.unlink(path)
     except FileNotFoundError:
         pass
     evt.log.info(f"{evt.sender.mxid}'s manhole was closed.")
-    await evt.reply("Your manhole was closed.")
+    try:
+        await evt.reply("Your manhole was closed.")
+    except AttributeError as e:
+        evt.log.warning(f"Failed to send manhole close notification: {e}")
 
 
 @command_handler(needs_auth=False, needs_admin=True, help_section=SECTION_ADMIN,
                  help_text="Close an open manhole.")
 async def close_manhole(evt: CommandEvent) -> None:
-    if not State.manhole:
+    if not evt.bridge.manhole:
         await evt.reply("There is no open manhole.")
         return
 
-    opened_by = State.opened_by
-    State.close()
-    State.manhole = None
-    State.close = None
-    State.opened_by = None
-    if opened_by != evt.sender:
+    opened_by = evt.bridge.manhole.opened_by
+    evt.bridge.manhole.close()
+    evt.bridge.manhole = None
+    if opened_by != evt.sender.mxid:
         await evt.reply(f"Closed manhole opened by {opened_by}")