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.

237 lines
8.4 KiB

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, TypedDict, Optional
  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
  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. class MatrixMediaInfo(TypedDict):
  62. w: int
  63. h: int
  64. size: int
  65. mimetype: str
  66. thumbnail_url: Optional[str]
  67. thumbnail_info: Optional['MatrixMediaInfo']
  68. class MatrixStickerInfo(TypedDict, total=False):
  69. body: str
  70. url: str
  71. info: MatrixMediaInfo
  72. def convert_image(data: bytes) -> (bytes, int, int):
  73. image: Image.Image = Image.open(BytesIO(data)).convert("RGBA")
  74. new_file = BytesIO()
  75. image.save(new_file, "png")
  76. w, h = image.size
  77. return new_file.getvalue(), w, h
  78. async def reupload_document(client: TelegramClient, document: Document) -> MatrixStickerInfo:
  79. print(f"Reuploading {document.id}", end="", flush=True)
  80. data = await client.download_media(document, file=bytes)
  81. print(".", end="", flush=True)
  82. data, width, height = convert_image(data)
  83. print(".", end="", flush=True)
  84. mxc = await upload(data, "image/png", f"{document.id}.png")
  85. print(".", flush=True)
  86. if width > 256 or height > 256:
  87. # Set the width and height to lower values so clients wouldn't show them as huge images
  88. if width > height:
  89. height = int(height / (width / 256))
  90. width = 256
  91. else:
  92. width = int(width / (height / 256))
  93. height = 256
  94. return {
  95. "body": "",
  96. "url": mxc,
  97. "info": {
  98. "w": width,
  99. "h": height,
  100. "size": len(data),
  101. "mimetype": "image/png",
  102. # Element iOS compatibility hack
  103. "thumbnail_url": mxc,
  104. "thumbnail_info": {
  105. "w": width,
  106. "h": height,
  107. "size": len(data),
  108. "mimetype": "image/png",
  109. },
  110. },
  111. }
  112. def add_to_index(name: str) -> None:
  113. index_path = os.path.join(args.output_dir, "index.json")
  114. try:
  115. with open(index_path) as index_file:
  116. index_data = json.load(index_file)
  117. except (FileNotFoundError, json.JSONDecodeError):
  118. index_data = {"packs": [], "homeserver_url": homeserver_url}
  119. if name not in index_data["packs"]:
  120. index_data["packs"].append(name)
  121. with open(index_path, "w") as index_file:
  122. json.dump(index_data, index_file, indent=" ")
  123. print(f"Added {name} to {index_path}")
  124. async def reupload_pack(client: TelegramClient, pack: StickerSetFull) -> None:
  125. if pack.set.animated:
  126. print("Animated stickerpacks are currently not supported")
  127. return
  128. pack_path = os.path.join(args.output_dir, f"{pack.set.short_name}.json")
  129. try:
  130. os.mkdir(os.path.dirname(pack_path))
  131. except FileExistsError:
  132. pass
  133. print(f"Reuploading {pack.set.title} with {pack.set.count} stickers "
  134. f"and writing output to {pack_path}")
  135. already_uploaded = {}
  136. try:
  137. with open(pack_path) as pack_file:
  138. existing_pack = json.load(pack_file)
  139. already_uploaded = {sticker["net.maunium.telegram.sticker"]["id"]: sticker
  140. for sticker in existing_pack["stickers"]}
  141. print(f"Found {len(already_uploaded)} already reuploaded stickers")
  142. except FileNotFoundError:
  143. pass
  144. reuploaded_documents: Dict[int, MatrixStickerInfo] = {}
  145. for document in pack.documents:
  146. try:
  147. reuploaded_documents[document.id] = already_uploaded[document.id]
  148. print(f"Skipped reuploading {document.id}")
  149. except KeyError:
  150. reuploaded_documents[document.id] = await reupload_document(client, document)
  151. for sticker in pack.packs:
  152. for document_id in sticker.documents:
  153. doc = reuploaded_documents[document_id]
  154. doc["body"] = sticker.emoticon
  155. doc["net.maunium.telegram.sticker"] = {
  156. "pack": {
  157. "id": pack.set.id,
  158. "short_name": pack.set.short_name,
  159. },
  160. "id": document_id,
  161. "emoticon": sticker.emoticon,
  162. }
  163. with open(pack_path, "w") as pack_file:
  164. json.dump({
  165. "title": pack.set.title,
  166. "short_name": pack.set.short_name,
  167. "id": pack.set.id,
  168. "hash": pack.set.hash,
  169. "stickers": list(reuploaded_documents.values()),
  170. }, pack_file, ensure_ascii=False)
  171. print(f"Saved {pack.set.title} as {pack.set.short_name}.json")
  172. add_to_index(os.path.basename(pack_path))
  173. pack_url_regex = re.compile(r"^(?:(?:https?://)?(?:t|telegram)\.(?:me|dog)/addstickers/)?"
  174. r"([A-Za-z0-9-_]+)"
  175. r"(?:\.json)?$")
  176. async def main():
  177. client = TelegramClient(args.session, 298751, "cb676d6bae20553c9996996a8f52b4d7")
  178. await client.start()
  179. if args.list:
  180. stickers: AllStickers = await client(GetAllStickersRequest(hash=0))
  181. index = 1
  182. width = len(str(stickers.sets))
  183. print("Your saved sticker packs:")
  184. for saved_pack in stickers.sets:
  185. print(f"{index:>{width}}. {saved_pack.title} "
  186. f"(t.me/addstickers/{saved_pack.short_name})")
  187. elif args.pack[0]:
  188. input_packs = []
  189. for pack_url in args.pack[0]:
  190. match = pack_url_regex.match(pack_url)
  191. if not match:
  192. print(f"'{pack_url}' doesn't look like a sticker pack URL")
  193. return
  194. input_packs.append(InputStickerSetShortName(short_name=match.group(1)))
  195. for input_pack in input_packs:
  196. pack: StickerSetFull = await client(GetStickerSetRequest(input_pack))
  197. await reupload_pack(client, pack)
  198. else:
  199. parser.print_help()
  200. asyncio.run(main())