matrix-appservice-kakaotalk/matrix_appservice_kakaotalk/puppet.py

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