Add support for logging in with a bot. Fixes #155

This commit is contained in:
Tulir Asokan
2018-06-23 00:44:41 +03:00
parent 17aefd02da
commit 25b1adf626
9 changed files with 116 additions and 45 deletions
+17 -10
View File
@@ -36,12 +36,16 @@ class AbstractUser:
az = None
def __init__(self):
self.connected = False
self.whitelisted = False
self.client = None
self.tgid = None
self.mxid = None
self.is_relaybot = False
self.is_bot = False
@property
def connected(self):
return self.client and self.client.is_connected()
def _init_client(self):
self.log.debug(f"Initializing client for {self.name}")
@@ -72,6 +76,8 @@ class AbstractUser:
self.log.exception("Failed to handle Telegram update")
async def _get_dialogs(self, limit=None):
if self.is_bot:
return
dialogs = await self.client.get_dialogs(limit=limit)
return [dialog.entity for dialog in dialogs if (
not isinstance(dialog.entity, (User, ChatForbidden, ChannelForbidden))
@@ -85,32 +91,33 @@ class AbstractUser:
async def is_logged_in(self):
return self.client and await self.client.is_user_authorized()
async def has_full_access(self):
return await self.is_logged_in() and self.whitelisted
async def has_full_access(self, allow_bot=False):
return self.whitelisted and (not self.is_bot or allow_bot) and await self.is_logged_in()
async def start(self):
if not self.client:
self._init_client()
self.connected = await self.client.connect()
await self.client.connect()
self.log.debug("%s connected: %s", self.mxid, self.connected)
return self
async def ensure_started(self, even_if_no_session=False):
if not self.whitelisted:
return self
self.log.info("CONNECTING USER %s, connected=%s, even_if_no_session=%s, session_count=%s",
self.mxid, self.connected, even_if_no_session,
self.session_container.Session.query.filter(
self.session_container.Session.session_id == self.mxid).count())
self.log.debug("ensure_started(%s, connected=%s, even_if_no_session=%s, session_count=%s)",
self.mxid, self.connected, even_if_no_session,
self.session_container.Session.query.filter(
self.session_container.Session.session_id == self.mxid).count())
should_connect = (even_if_no_session or
self.session_container.Session.query.filter(
self.session_container.Session.session_id == self.mxid).count() > 0)
if not self.connected and should_connect:
return await self.start()
await self.start()
return self
def stop(self):
self.client.disconnect()
self.client = None
self.connected = False
# region Telegram update handling
+1
View File
@@ -42,6 +42,7 @@ class Bot(AbstractUser):
self.whitelisted = True
self.username = None
self.is_relaybot = True
self.is_bot = True
self.chats = {chat.id: chat.type for chat in BotChat.query.all()}
self.tg_whitelist = []
self.whitelist_group_admins = config["bridge.relaybot.whitelist_group_admins"] or False
+1 -2
View File
@@ -104,8 +104,7 @@ class Puppet(Base):
class BotChat(Base):
query = None
__tablename__ = "bot_chat"
bot_id = Column(Integer, primary_key=True, default=0)
chat_id = Column(Integer, primary_key=True)
id = Column(Integer, primary_key=True)
type = Column(String, nullable=False)
+4 -4
View File
@@ -113,7 +113,7 @@ class MatrixHandler:
return
await user.ensure_started()
portal = Portal.get_by_mxid(room)
if user and await user.has_full_access() and portal:
if user and await user.has_full_access(allow_bot=True) and portal:
await portal.invite_telegram(inviter, user)
return
@@ -218,13 +218,13 @@ class MatrixHandler:
async def handle_power_levels(self, room, sender, new, old):
portal = Portal.get_by_mxid(room)
sender = await User.get_by_mxid(sender).ensure_started()
if await sender.has_full_access() and portal:
if await sender.has_full_access(allow_bot=True) and portal:
await portal.handle_matrix_power_levels(sender, new["users"], old["users"])
async def handle_room_meta(self, type, room, sender, content):
portal = Portal.get_by_mxid(room)
sender = await User.get_by_mxid(sender).ensure_started()
if await sender.has_full_access() and portal:
if await sender.has_full_access(allow_bot=True) and portal:
handler, content_key = {
"m.room.name": (portal.handle_matrix_title, "name"),
"m.room.topic": (portal.handle_matrix_about, "topic"),
@@ -237,7 +237,7 @@ class MatrixHandler:
async def handle_room_pin(self, room, sender, new_events, old_events):
portal = Portal.get_by_mxid(room)
sender = await User.get_by_mxid(sender).ensure_started()
if await sender.has_full_access() and portal:
if await sender.has_full_access(allow_bot=True) and portal:
events = new_events - old_events
if len(events) > 0:
# New event pinned, set that as pinned in Telegram.
+24 -12
View File
@@ -24,7 +24,6 @@ import mimetypes
import unicodedata
import hashlib
import logging
import re
import magic
from sqlalchemy.exc import IntegrityError, InvalidRequestError
@@ -346,12 +345,21 @@ class Portal:
return None
return self.alias_template.format(groupname=username)
def add_bot_chat(self, entity):
if self.bot and entity.id == self.bot.tgid:
self.bot.add_chat(self.tgid, self.peer_type)
return
user = u.User.get_by_tgid(entity.id)
if user and user.is_bot:
user.register_portal(self)
async def sync_telegram_users(self, source, users):
allowed_tgids = set()
for entity in users:
puppet = p.Puppet.get(entity.id)
if self.bot and puppet.tgid == self.bot.tgid:
self.bot.add_chat(self.tgid, self.peer_type)
if entity.bot:
self.add_bot_chat(entity)
allowed_tgids.add(entity.id)
await puppet.intent.ensure_joined(self.mxid)
await puppet.update_info(source, entity)
@@ -380,6 +388,9 @@ class Portal:
"User had left this Telegram chat.")
continue
mx_user = u.User.get_by_mxid(user, create=False)
if mx_user and mx_user.is_bot and mx_user.tgid not in allowed_tgids:
mx_user.unregister_portal(self)
if mx_user and not self.has_bot and mx_user.tgid not in allowed_tgids:
await self.main_intent.kick(self.mxid, mx_user.mxid,
"You had left this Telegram chat.")
@@ -560,7 +571,8 @@ class Portal:
if p.Puppet.get_id_from_mxid(member) or member == self.main_intent.mxid:
continue
user = await u.User.get_by_mxid(member).ensure_started()
if (has_bot and user.relaybot_whitelisted) or await user.has_full_access():
if (has_bot and user.relaybot_whitelisted) or await user.has_full_access(
allow_bot=True):
authenticated.append(user)
return authenticated
@@ -607,7 +619,7 @@ class Portal:
return ""
async def leave_matrix(self, user, source, event_id):
if not await user.is_logged_in():
if await user.needs_relaybot(self):
async with self.require_send_lock(self.bot.tgid):
response = await self.bot.client.send_message(
self.peer, f"__{user.displayname} left the room.__", markdown=True)
@@ -639,7 +651,7 @@ class Portal:
await user.client(LeaveChannelRequest(channel=channel))
async def join_matrix(self, user, event_id):
if not await user.is_logged_in():
if await user.needs_relaybot(self):
async with self.require_send_lock(self.bot.tgid):
response = await self.bot.client.send_message(
self.peer, f"__{user.displayname} joined the room.__", markdown=True)
@@ -654,19 +666,19 @@ class Portal:
pass
@staticmethod
def _preprocess_matrix_message(sender, is_logged_in, message):
def _preprocess_matrix_message(sender, use_relaybot, message):
msgtype = message["msgtype"]
if msgtype == "m.emote":
if "formatted_body" in message:
tpl = config["bridge.message_formats.m_emote.html"]
tpl_args = dict(sender_display_name=sender.displayname,
message=message['formatted_body'])
message=message['formatted_body'])
message["formatted_body"] = Template(tpl).safe_substitute(tpl_args)
tpl = config["bridge.message_formats.m_emote.plain"]
tpl_args = dict(sender_display_name=sender.displayname, message=message['body'])
message["body"] = Template(tpl).safe_substitute(tpl_args)
message["msgtype"] = "m.text"
elif not is_logged_in:
elif not use_relaybot:
html = message["formatted_body"] if "formatted_body" in message else None
text = message["body"]
if msgtype == "m.text":
@@ -675,7 +687,7 @@ class Portal:
if not html:
html = escape_html(text)
tpl = config["bridge.message_formats.m_text.html"]
tpl_args = dict(sender_display_name=sender.displayname,message=html)
tpl_args = dict(sender_display_name=sender.displayname, message=html)
html = Template(tpl).safe_substitute(tpl_args)
tpl = config["bridge.message_formats.m_text.plain"]
tpl_args = dict(sender_display_name=sender.displayname, message=text)
@@ -813,7 +825,7 @@ class Portal:
self.db.commit()
async def handle_matrix_message(self, sender, message, event_id):
logged_in = await sender.is_logged_in()
logged_in = not await sender.needs_relaybot(self)
client = sender.client if logged_in else self.bot.client
sender_id = sender.tgid if logged_in else self.bot.tgid
space = (self.tgid if self.peer_type == "channel" # Channels have their own ID space
@@ -850,7 +862,7 @@ class Portal:
pass
async def handle_matrix_deletion(self, deleter, event_id):
deleter = deleter if await deleter.is_logged_in() else self.bot
deleter = deleter if await deleter.needs_relaybot(self) else self.bot
space = self.tgid if self.peer_type == "channel" else deleter.tgid
message = DBMessage.query.filter(DBMessage.mxid == event_id,
DBMessage.tg_space == space,
+18 -2
View File
@@ -45,15 +45,16 @@ class PublicBridgeWebsite:
async def get_login(self, request):
user = (User.get_by_mxid(request.rel_url.query["mxid"], create=False)
if "mxid" in request.rel_url.query else None)
state = "token" if request.rel_url.query.get("mode", "") == "bot" else "request"
if not user:
return self.render_login(
mxid=request.rel_url.query["mxid"] if "mxid" in request.rel_url.query else None,
state="request")
state=state)
elif not user.whitelisted:
return self.render_login(mxid=user.mxid, error="You are not whitelisted.", status=403)
await user.ensure_started()
if not await user.is_logged_in():
return self.render_login(mxid=user.mxid, state="request")
return self.render_login(mxid=user.mxid, state=state)
return self.render_login(mxid=user.mxid, username=user.username)
@@ -62,6 +63,19 @@ class PublicBridgeWebsite:
text=self.login.render(username=username, state=state, error=error,
message=message, mxid=mxid))
async def post_login_token(self, user, token):
try:
user_info = await user.client.sign_in(bot_token=token)
asyncio.ensure_future(user.post_login(user_info), loop=self.loop)
if user.command_status and user.command_status["action"] == "Login":
user.command_status = None
return self.render_login(mxid=user.mxid, state="logged-in", status=200,
username=user_info.username)
except Exception:
self.log.exception("Error sending bot token")
return self.render_login(mxid=user.mxid, state="token", status=500,
error="Internal server error while sending token.")
async def post_login_phone(self, user, phone):
try:
await user.client.sign_in(phone or "+123")
@@ -153,6 +167,8 @@ class PublicBridgeWebsite:
if "phone" in data:
return await self.post_login_phone(user, data["phone"])
elif "token" in data:
return await self.post_login_token(user, data["token"])
elif "code" in data:
resp = await self.post_login_code(user, data["code"],
password_in_data="password" in data)
+36 -8
View File
@@ -29,6 +29,25 @@ along with this program. If not, see <https://www.gnu.org/licenses/>.
<link rel="stylesheet"
href="//cdn.rawgit.com/milligram/milligram/master/dist/milligram.min.css">
<link rel="stylesheet" href="login.css"/>
<script>
function switchToBotLogin() {
const params = new URLSearchParams(location.search.slice(1))
params.set("mode", "bot")
location.search = "?" + params.toString()
console.log(location.search)
}
function goBack() {
let params = new URLSearchParams(location.search.slice(1))
const mxid = params.get("mxid")
params = new URLSearchParams()
if (mxid) {
params.set("mxid", mxid)
}
location.replace(location.href.split("?")[0] + "?" + params.toString())
}
</script>
</head>
<body>
<main class="container">
@@ -40,6 +59,13 @@ along with this program. If not, see <https://www.gnu.org/licenses/>.
You can now close this page.
You should be invited to Telegram portals on Matrix momentarily.
</p>
% elif state == "bot-logged-in":
<h1>Logged in successfully!</h1>
<p>
Logged in as @${username}.
You can now close this page.
You should be invited to Telegram portals on Matrix momentarily.
</p>
% else:
<h1>You're already logged in!</h1>
<p>
@@ -67,24 +93,26 @@ along with this program. If not, see <https://www.gnu.org/licenses/>.
<label for="value">Phone number</label>
<input type="tel" id="value" name="phone" placeholder="Enter phone number"/>
<button type="submit">Request code</button>
<button class="button-clear" type="button" onclick="switchToBotLogin()">
Use bot token
</button>
% elif state == "token":
<label for="value">Bot token</label>
<input type="text" id="value" name="token" placeholder="Enter bot API token"/>
<button type="submit">Sign in</button>
% elif state == "code":
<label for="value">Phone code</label>
<input type="number" id="value" name="code" placeholder="Enter phone code"/>
<button type="submit">Sign in</button>
<div class="float-right">
<button class="button-clear" type="button"
onclick="location.replace(location.href)">
Go back
</button>
</div>
% elif state == "password":
<label for="value">Password</label>
<input type="password" id="value" name="password"
placeholder="Enter password"/>
<button type="submit">Sign in</button>
% endif
% if state != "request":
<div class="float-right">
<button class="button-clear" type="button"
onclick="location.replace(location.href)">
<button class="button-clear" type="button" onclick="goBack()">
Go back
</button>
</div>
-1
View File
@@ -19,7 +19,6 @@ import re
import logging
from telethon.tl.types import UserProfilePhoto
from telethon.errors import LocationInvalidError
from .db import Puppet as DBPuppet
from . import util
+15 -6
View File
@@ -36,10 +36,11 @@ class User(AbstractUser):
by_tgid = {}
def __init__(self, mxid, tgid=None, username=None, db_contacts=None, saved_contacts=0,
db_portals=None, db_instance=None):
is_bot=False, db_portals=None, db_instance=None):
super().__init__()
self.mxid = mxid
self.tgid = tgid
self.is_bot = is_bot
self.username = username
self.contacts = []
self.saved_contacts = saved_contacts
@@ -127,7 +128,7 @@ class User(AbstractUser):
@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.portals, db_instance=db_user)
False, db_user.saved_contacts, db_user.portals, db_instance=db_user)
# endregion
# region Telegram connection management
@@ -148,19 +149,23 @@ class User(AbstractUser):
async def post_login(self, info=None):
try:
await self.update_info(info)
await self.sync_dialogs()
await self.sync_contacts()
if not self.is_bot:
await self.sync_dialogs()
await self.sync_contacts()
if config["bridge.catch_up"]:
await self.client.catch_up()
except Exception:
self.log.exception("Failed to run post-login functions")
self.log.exception("Failed to run post-login functions for %s", self.mxid)
# endregion
# region Telegram actions that need custom methods
async def update_info(self, info=None):
async def update_info(self, info: User = None):
info = info or await self.client.get_me()
changed = False
if self.is_bot != info.bot:
self.is_bot = info.bot
changed = True
if self.username != info.username:
self.username = info.username
changed = True
@@ -252,6 +257,10 @@ class User(AbstractUser):
except KeyError:
pass
async def needs_relaybot(self, portal):
return not await self.is_logged_in() or (
self.is_bot and portal.tgid_full not in self.portals)
def _hash_contacts(self):
acc = 0
for id in sorted([self.saved_contacts] + [contact.id for contact in self.contacts]):