You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

258 lines
9.2 KiB

4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
  1. # Copyright (c) 2020 Tulir Asokan
  2. #
  3. # This Source Code Form is subject to the terms of the Mozilla Public
  4. # License, v. 2.0. If a copy of the MPL was not distributed with this
  5. # file, You can obtain one at http://mozilla.org/MPL/2.0/.
  6. from typing import Dict, Optional, TYPE_CHECKING
  7. from io import BytesIO
  8. import argparse
  9. import os.path
  10. import asyncio
  11. import json
  12. import re
  13. from aiohttp import ClientSession
  14. from yarl import URL
  15. from PIL import Image
  16. from telethon import TelegramClient
  17. from telethon.tl.functions.messages import GetAllStickersRequest, GetStickerSetRequest
  18. from telethon.tl.types.messages import AllStickers
  19. from telethon.tl.types import InputStickerSetShortName, Document, DocumentAttributeSticker
  20. from telethon.tl.types.messages import StickerSet as StickerSetFull
  21. parser = argparse.ArgumentParser()
  22. parser.add_argument("--list", help="List your saved sticker packs", action="store_true")
  23. parser.add_argument("--session", help="Telethon session file name", default="sticker-import")
  24. parser.add_argument("--config", help="Path to JSON file with Matrix homeserver and access_token",
  25. type=str, default="config.json")
  26. parser.add_argument("--output-dir", help="Directory to write packs to", default="web/packs/",
  27. type=str)
  28. parser.add_argument("pack", help="Sticker pack URLs to import", action="append", nargs="*")
  29. args = parser.parse_args()
  30. async def whoami(url: URL, access_token: str) -> str:
  31. headers = {"Authorization": f"Bearer {access_token}"}
  32. async with ClientSession() as sess, sess.get(url, headers=headers) as resp:
  33. resp.raise_for_status()
  34. user_id = (await resp.json())["user_id"]
  35. print(f"Access token validated (user ID: {user_id})")
  36. return user_id
  37. try:
  38. with open(args.config) as config_file:
  39. config = json.load(config_file)
  40. homeserver_url = config["homeserver"]
  41. access_token = config["access_token"]
  42. except FileNotFoundError:
  43. print("Matrix config file not found. Please enter your homeserver and access token.")
  44. homeserver_url = input("Homeserver URL: ")
  45. access_token = input("Access token: ")
  46. whoami_url = URL(homeserver_url) / "_matrix" / "client" / "r0" / "account" / "whoami"
  47. user_id = asyncio.run(whoami(whoami_url, access_token))
  48. with open(args.config, "w") as config_file:
  49. json.dump({
  50. "homeserver": homeserver_url,
  51. "user_id": user_id,
  52. "access_token": access_token
  53. }, config_file)
  54. print(f"Wrote config to {args.config}")
  55. upload_url = URL(homeserver_url) / "_matrix" / "media" / "r0" / "upload"
  56. async def upload(data: bytes, mimetype: str, filename: str) -> str:
  57. url = upload_url.with_query({"filename": filename})
  58. headers = {"Content-Type": mimetype, "Authorization": f"Bearer {access_token}"}
  59. async with ClientSession() as sess, sess.post(url, data=data, headers=headers) as resp:
  60. return (await resp.json())["content_uri"]
  61. if TYPE_CHECKING:
  62. from typing import TypedDict
  63. class MatrixMediaInfo(TypedDict):
  64. w: int
  65. h: int
  66. size: int
  67. mimetype: str
  68. thumbnail_url: Optional[str]
  69. thumbnail_info: Optional['MatrixMediaInfo']
  70. class MatrixStickerInfo(TypedDict, total=False):
  71. body: str
  72. url: str
  73. info: MatrixMediaInfo
  74. id: str
  75. def convert_image(data: bytes) -> (bytes, int, int):
  76. image: Image.Image = Image.open(BytesIO(data)).convert("RGBA")
  77. new_file = BytesIO()
  78. image.save(new_file, "png")
  79. w, h = image.size
  80. return new_file.getvalue(), w, h
  81. async def reupload_document(client: TelegramClient, document: Document) -> 'MatrixStickerInfo':
  82. print(f"Reuploading {document.id}", end="", flush=True)
  83. data = await client.download_media(document, file=bytes)
  84. print(".", end="", flush=True)
  85. data, width, height = convert_image(data)
  86. print(".", end="", flush=True)
  87. mxc = await upload(data, "image/png", f"{document.id}.png")
  88. print(".", flush=True)
  89. if width > 256 or height > 256:
  90. # Set the width and height to lower values so clients wouldn't show them as huge images
  91. if width > height:
  92. height = int(height / (width / 256))
  93. width = 256
  94. else:
  95. width = int(width / (height / 256))
  96. height = 256
  97. return {
  98. "body": "",
  99. "url": mxc,
  100. "info": {
  101. "w": width,
  102. "h": height,
  103. "size": len(data),
  104. "mimetype": "image/png",
  105. # Element iOS compatibility hack
  106. "thumbnail_url": mxc,
  107. "thumbnail_info": {
  108. "w": width,
  109. "h": height,
  110. "size": len(data),
  111. "mimetype": "image/png",
  112. },
  113. },
  114. }
  115. def add_to_index(name: str) -> None:
  116. index_path = os.path.join(args.output_dir, "index.json")
  117. try:
  118. with open(index_path) as index_file:
  119. index_data = json.load(index_file)
  120. except (FileNotFoundError, json.JSONDecodeError):
  121. index_data = {"packs": [], "homeserver_url": homeserver_url}
  122. if name not in index_data["packs"]:
  123. index_data["packs"].append(name)
  124. with open(index_path, "w") as index_file:
  125. json.dump(index_data, index_file, indent=" ")
  126. print(f"Added {name} to {index_path}")
  127. def add_meta(document: Document, info: 'MatrixStickerInfo', pack: StickerSetFull) -> None:
  128. for attr in document.attributes:
  129. if isinstance(attr, DocumentAttributeSticker):
  130. info["body"] = attr.alt
  131. info["id"] = str(document.id)
  132. info["net.maunium.telegram.sticker"] = {
  133. "pack": {
  134. "id": str(pack.set.id),
  135. "short_name": pack.set.short_name,
  136. },
  137. "id": str(document.id),
  138. "emoticons": [],
  139. }
  140. async def reupload_pack(client: TelegramClient, pack: StickerSetFull) -> None:
  141. if pack.set.animated:
  142. print("Animated stickerpacks are currently not supported")
  143. return
  144. pack_path = os.path.join(args.output_dir, f"{pack.set.short_name}.json")
  145. try:
  146. os.mkdir(os.path.dirname(pack_path))
  147. except FileExistsError:
  148. pass
  149. print(f"Reuploading {pack.set.title} with {pack.set.count} stickers "
  150. f"and writing output to {pack_path}")
  151. already_uploaded = {}
  152. try:
  153. with open(pack_path) as pack_file:
  154. existing_pack = json.load(pack_file)
  155. already_uploaded = {int(sticker["net.maunium.telegram.sticker"]["id"]): sticker
  156. for sticker in existing_pack["stickers"]}
  157. print(f"Found {len(already_uploaded)} already reuploaded stickers")
  158. except FileNotFoundError:
  159. pass
  160. reuploaded_documents: Dict[int, 'MatrixStickerInfo'] = {}
  161. for document in pack.documents:
  162. try:
  163. reuploaded_documents[document.id] = already_uploaded[document.id]
  164. print(f"Skipped reuploading {document.id}")
  165. except KeyError:
  166. reuploaded_documents[document.id] = await reupload_document(client, document)
  167. # Always ensure the body and telegram metadata is correct
  168. add_meta(document, reuploaded_documents[document.id], pack)
  169. for sticker in pack.packs:
  170. if not sticker.emoticon:
  171. continue
  172. for document_id in sticker.documents:
  173. doc = reuploaded_documents[document_id]
  174. # If there was no sticker metadata, use the first emoji we find
  175. if doc["body"] == "":
  176. doc["body"] = sticker.emoticon
  177. doc["net.maunium.telegram.sticker"]["emoticons"].append(sticker.emoticon)
  178. with open(pack_path, "w") as pack_file:
  179. json.dump({
  180. "title": pack.set.title,
  181. "short_name": pack.set.short_name,
  182. "id": str(pack.set.id),
  183. "hash": str(pack.set.hash),
  184. "stickers": list(reuploaded_documents.values()),
  185. }, pack_file, ensure_ascii=False)
  186. print(f"Saved {pack.set.title} as {pack.set.short_name}.json")
  187. add_to_index(os.path.basename(pack_path))
  188. pack_url_regex = re.compile(r"^(?:(?:https?://)?(?:t|telegram)\.(?:me|dog)/addstickers/)?"
  189. r"([A-Za-z0-9-_]+)"
  190. r"(?:\.json)?$")
  191. async def main():
  192. client = TelegramClient(args.session, 298751, "cb676d6bae20553c9996996a8f52b4d7")
  193. await client.start()
  194. if args.list:
  195. stickers: AllStickers = await client(GetAllStickersRequest(hash=0))
  196. index = 1
  197. width = len(str(stickers.sets))
  198. print("Your saved sticker packs:")
  199. for saved_pack in stickers.sets:
  200. print(f"{index:>{width}}. {saved_pack.title} "
  201. f"(t.me/addstickers/{saved_pack.short_name})")
  202. elif args.pack[0]:
  203. input_packs = []
  204. for pack_url in args.pack[0]:
  205. match = pack_url_regex.match(pack_url)
  206. if not match:
  207. print(f"'{pack_url}' doesn't look like a sticker pack URL")
  208. return
  209. input_packs.append(InputStickerSetShortName(short_name=match.group(1)))
  210. for input_pack in input_packs:
  211. pack: StickerSetFull = await client(GetStickerSetRequest(input_pack))
  212. await reupload_pack(client, pack)
  213. else:
  214. parser.print_help()
  215. await client.disconnect()
  216. asyncio.get_event_loop().run_until_complete(main())