Andrew Ferrazzutti
adb7453e1a
- Replace builtin generic type annotations with classes from Typing - Replace union type expressions with Union and Optional
313 lines
9.9 KiB
Python
313 lines
9.9 KiB
Python
# matrix-appservice-kakaotalk - A Matrix-KakaoTalk puppeting bridge.
|
|
# Copyright (C) 2022 Tulir Asokan, Andrew Ferrazzutti
|
|
#
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as published by
|
|
# the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU Affero General Public License for more details.
|
|
#
|
|
# 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 __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, AsyncGenerator, AsyncIterable, Awaitable, Dict, Optional, cast
|
|
from datetime import datetime, timedelta
|
|
import asyncio
|
|
|
|
from yarl import URL
|
|
|
|
from mautrix.appservice import IntentAPI
|
|
from mautrix.bridge import BasePuppet, async_getter_lock
|
|
from mautrix.types import ContentURI, RoomID, SyncToken, UserID
|
|
from mautrix.util import magic
|
|
from mautrix.util.simple_template import SimpleTemplate
|
|
|
|
from . import matrix as m, portal as p, user as u
|
|
from .config import Config
|
|
from .db import Puppet as DBPuppet
|
|
|
|
from .kt.types.bson import Long
|
|
|
|
from .kt.types.api.struct import FriendStruct
|
|
from .kt.client.types import UserInfoUnion
|
|
|
|
if TYPE_CHECKING:
|
|
from .__main__ import KakaoTalkBridge
|
|
|
|
|
|
class Puppet(DBPuppet, BasePuppet):
|
|
mx: m.MatrixHandler
|
|
config: Config
|
|
hs_domain: str
|
|
mxid_template: SimpleTemplate[int]
|
|
|
|
by_ktid: Dict[int, Puppet] = {}
|
|
by_custom_mxid: Dict[UserID, Puppet] = {}
|
|
|
|
_last_info_sync: Optional[datetime]
|
|
|
|
def __init__(
|
|
self,
|
|
ktid: Long,
|
|
name: Optional[str] = None,
|
|
photo_id: Optional[str] = None,
|
|
photo_mxc: Optional[ContentURI] = None,
|
|
name_set: bool = False,
|
|
avatar_set: bool = False,
|
|
is_registered: bool = False,
|
|
custom_mxid: Optional[UserID] = None,
|
|
access_token: Optional[str] = None,
|
|
next_batch: Optional[SyncToken] = None,
|
|
base_url: Optional[URL] = None,
|
|
) -> None:
|
|
super().__init__(
|
|
ktid,
|
|
name,
|
|
photo_id,
|
|
photo_mxc,
|
|
name_set,
|
|
avatar_set,
|
|
is_registered,
|
|
custom_mxid,
|
|
access_token,
|
|
next_batch,
|
|
base_url,
|
|
)
|
|
self._last_info_sync = None
|
|
|
|
self.default_mxid = self.get_mxid_from_id(ktid)
|
|
self.default_mxid_intent = self.az.intent.user(self.default_mxid)
|
|
self.intent = self._fresh_intent()
|
|
|
|
self.log = self.log.getChild(str(self.ktid))
|
|
|
|
@property
|
|
def should_sync(self) -> bool:
|
|
now = datetime.now()
|
|
if not self._last_info_sync or now - self._last_info_sync > timedelta(hours=48):
|
|
self._last_info_sync = now
|
|
return True
|
|
return False
|
|
|
|
async def default_puppet_should_leave_room(self, room_id: RoomID) -> bool:
|
|
portal = await p.Portal.get_by_mxid(room_id)
|
|
return portal and portal.kt_sender != self.ktid
|
|
|
|
async def _leave_rooms_with_default_user(self) -> None:
|
|
await super()._leave_rooms_with_default_user()
|
|
# Make the user join all private chat portals.
|
|
await asyncio.gather(
|
|
*[
|
|
self.intent.ensure_joined(portal.mxid)
|
|
async for portal in p.Portal.get_all_by_receiver(self.ktid)
|
|
if portal.mxid
|
|
]
|
|
)
|
|
|
|
def intent_for(self, portal: p.Portal) -> IntentAPI:
|
|
if portal.kt_sender == self.ktid or (
|
|
portal.backfill_lock.locked and self.config["bridge.backfill.invite_own_puppet"]
|
|
):
|
|
return self.default_mxid_intent
|
|
return self.intent
|
|
|
|
@classmethod
|
|
def init_cls(cls, bridge: KakaoTalkBridge) -> AsyncIterable[Awaitable[None]]:
|
|
cls.config = bridge.config
|
|
cls.loop = bridge.loop
|
|
cls.mx = bridge.matrix
|
|
cls.az = bridge.az
|
|
cls.hs_domain = cls.config["homeserver.domain"]
|
|
cls.mxid_template = SimpleTemplate(
|
|
template=cls.config["bridge.username_template"],
|
|
keyword="userid",
|
|
prefix="@",
|
|
suffix=f":{Puppet.hs_domain}",
|
|
type=int,
|
|
)
|
|
cls.sync_with_custom_puppets = cls.config["bridge.sync_with_custom_puppets"]
|
|
cls.homeserver_url_map = {
|
|
server: URL(url)
|
|
for server, url in cls.config["bridge.double_puppet_server_map"].items()
|
|
}
|
|
cls.allow_discover_url = cls.config["bridge.double_puppet_allow_discovery"]
|
|
cls.login_shared_secret_map = {
|
|
server: secret.encode("utf-8")
|
|
for server, secret in cls.config["bridge.login_shared_secret_map"].items()
|
|
}
|
|
cls.login_device_name = "KakaoTalk Bridge"
|
|
|
|
return (puppet.try_start() async for puppet in Puppet.get_all_with_custom_mxid())
|
|
|
|
# region User info updating
|
|
|
|
def update_info_from_participant(
|
|
self,
|
|
source: u.User,
|
|
info: UserInfoUnion,
|
|
update_avatar: bool = True,
|
|
) -> Awaitable[Puppet]:
|
|
return self._update_info(
|
|
source,
|
|
info.nickname,
|
|
info.profileURL,
|
|
update_avatar
|
|
)
|
|
|
|
def update_info_from_friend(
|
|
self,
|
|
source: u.User,
|
|
info: FriendStruct,
|
|
update_avatar: bool = True,
|
|
) -> Awaitable[Puppet]:
|
|
return self._update_info(
|
|
source,
|
|
info.nickName,
|
|
info.profileImageUrl,
|
|
update_avatar
|
|
)
|
|
|
|
async def _update_info(
|
|
self,
|
|
source: u.User,
|
|
name: str,
|
|
avatar_url: str,
|
|
update_avatar: bool = True,
|
|
) -> Puppet:
|
|
self._last_info_sync = datetime.now()
|
|
try:
|
|
changed = await self._update_name(name)
|
|
if update_avatar:
|
|
changed = await self._update_photo(source, avatar_url) or changed
|
|
if changed:
|
|
await self.save()
|
|
except Exception:
|
|
self.log.exception(f"Failed to update info from source {source.ktid}")
|
|
return self
|
|
|
|
async def _update_name(self, name: str) -> bool:
|
|
name = self.config["bridge.displayname_template"].format(displayname=name)
|
|
if name != self.name or not self.name_set:
|
|
self.name = name
|
|
try:
|
|
await self.default_mxid_intent.set_displayname(self.name)
|
|
self.name_set = True
|
|
except Exception:
|
|
self.log.exception("Failed to set displayname")
|
|
self.name_set = False
|
|
return True
|
|
return False
|
|
|
|
@classmethod
|
|
async def reupload_avatar(
|
|
cls,
|
|
source: u.User,
|
|
intent: IntentAPI,
|
|
url: str,
|
|
ktid: int,
|
|
) -> ContentURI:
|
|
async with source.client.get(url) as resp:
|
|
data = await resp.read()
|
|
mime = magic.mimetype(data)
|
|
return await intent.upload_media(
|
|
data, mime_type=mime, async_upload=cls.config["homeserver.async_media"]
|
|
)
|
|
|
|
async def _update_photo(self, source: u.User, photo_id: str) -> bool:
|
|
if photo_id != self.photo_id or not self.avatar_set:
|
|
self.photo_id = photo_id
|
|
if photo_id:
|
|
self.photo_mxc = await self.reupload_avatar(
|
|
source,
|
|
self.default_mxid_intent,
|
|
photo_id,
|
|
self.ktid,
|
|
)
|
|
else:
|
|
self.photo_mxc = ContentURI("")
|
|
try:
|
|
await self.default_mxid_intent.set_avatar_url(self.photo_mxc)
|
|
self.avatar_set = True
|
|
except Exception:
|
|
self.log.exception("Failed to set avatar")
|
|
self.avatar_set = False
|
|
return True
|
|
return False
|
|
|
|
# endregion
|
|
# region Database getters
|
|
|
|
def _add_to_cache(self) -> None:
|
|
self.by_ktid[self.ktid] = self
|
|
if self.custom_mxid:
|
|
self.by_custom_mxid[self.custom_mxid] = self
|
|
|
|
@classmethod
|
|
@async_getter_lock
|
|
async def get_by_ktid(cls, ktid: int, *, create: bool = True) -> Optional[Puppet]:
|
|
try:
|
|
return cls.by_ktid[int]
|
|
except KeyError:
|
|
pass
|
|
|
|
puppet = cast(cls, await super().get_by_ktid(ktid))
|
|
if puppet:
|
|
puppet._add_to_cache()
|
|
return puppet
|
|
|
|
if create:
|
|
puppet = cls(ktid)
|
|
await puppet.insert()
|
|
puppet._add_to_cache()
|
|
return puppet
|
|
|
|
return None
|
|
|
|
@classmethod
|
|
async def get_by_mxid(cls, mxid: UserID, create: bool = True) -> Optional[Puppet]:
|
|
ktid = cls.get_id_from_mxid(mxid)
|
|
if ktid:
|
|
return await cls.get_by_ktid(ktid, create=create)
|
|
return None
|
|
|
|
@classmethod
|
|
@async_getter_lock
|
|
async def get_by_custom_mxid(cls, mxid: UserID) -> Optional[Puppet]:
|
|
try:
|
|
return cls.by_custom_mxid[mxid]
|
|
except KeyError:
|
|
pass
|
|
|
|
puppet = cast(cls, await super().get_by_custom_mxid(mxid))
|
|
if puppet:
|
|
puppet._add_to_cache()
|
|
return puppet
|
|
|
|
return None
|
|
|
|
@classmethod
|
|
def get_id_from_mxid(cls, mxid: UserID) -> Optional[int]:
|
|
return cls.mxid_template.parse(mxid)
|
|
|
|
@classmethod
|
|
def get_mxid_from_id(cls, ktid: int) -> UserID:
|
|
return UserID(cls.mxid_template.format_full(ktid))
|
|
|
|
@classmethod
|
|
async def get_all_with_custom_mxid(cls) -> AsyncGenerator[Puppet, None]:
|
|
puppets = await super().get_all_with_custom_mxid()
|
|
puppet: cls
|
|
for puppet in puppets:
|
|
try:
|
|
yield cls.by_ktid[puppet.ktid]
|
|
except KeyError:
|
|
puppet._add_to_cache()
|
|
yield puppet
|
|
|
|
# endregion
|