main
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3"""Set permissions via environment variables.
4
5CTYPE: Chat Type. Enum: {GROUP, CHANNEL, BOT, PRIVATE}
6CID: Chat ID
7TID: Thread ID
8UID: User ID
9
10CONFIG_GLOBAL: json string config for global permission
11CONFIG_C{CID}: json string config for {CID}
12CONFIG_C{CID}_T{TID}: json string config for {CID}-{TID}
13CONFIG_U{UID}: json string config for {UID}
14CONFIG_C{CID}_U{UID}: json string config for {CID} and {UID}
15CONFIG_C{CID}_T{TID}_U{UID}: json string config for {CID}-{TID} and {UID}
16
17If there are multiple envvars for the same message, they will be applied in the order of:
18Global Level -> Chat Type Level -> Chat Level -> Chat_Thread Level -> User Level -> Chat_User Level -> Chat_Thread_User Level
19"""
20
21import json
22import os
23
24from glom import glom
25from loguru import logger
26from pyrogram.types import Message
27
28from config import cache
29from messages.modify import message_modify
30from messages.parser import get_thread_id
31from utils import slim_cid
32
33
34def set_permission(message: Message) -> dict:
35 """Set permissions for this message."""
36 message = message_modify(message)
37 ctype = message.chat.type.name if message.chat and message.chat.type else ""
38 ctype = ctype.upper().removeprefix("SUPER") # SUPERGROUP -> GROUP
39 cid = glom(message, "chat.id", default=0) or 0
40 cid = slim_cid(cid) # remove -100 prefix
41 uid = glom(message, "from_user.id", default=0) or 0
42 tid = get_thread_id(message)
43
44 # Default permission
45 permissions = {
46 "disabled": False, # this is the switch for all permissions
47 "prepend_sender_user": False, # prepend: 👤[@username](tg://user?id={uid})//
48 "need_prefix": True, # need /dl for social media preview
49 "ai_text": True,
50 "ai_img": True,
51 "ai_video": True,
52 "ai_summarize": True,
53 "chat_summary": True,
54 "asr": True,
55 "audio_extract": True,
56 "danmu": True,
57 "subtitle": True,
58 "wget": True,
59 "ocr": True,
60 "price": True,
61 "convert_img": True,
62 "tts": True,
63 "ytb": True,
64 "arxiv": True,
65 "google_search": True,
66 "show_progress": True,
67 "detail_progress": True,
68 "social": True,
69 "douyin": True,
70 "tiktok": True,
71 "instagram": True,
72 "twitter": True,
73 "weibo": True,
74 "xhs": True,
75 "v2ex": True,
76 "music163": True,
77 "spotify": True,
78 "github": True,
79 "wechat": True,
80 "reddit": True,
81 "ytdlp": True,
82 "ytdlp_bilibili": True,
83 "ytdlp_youtube": True,
84 "history": True,
85 "favorite": True,
86 "convert_chinese": True,
87 "quotly": True,
88 "tmdb": True,
89 "ffmpeg": True,
90 "watermark": True,
91 }
92 # Customization
93 # skip for service message
94 if message.service:
95 permissions["disabled"] = True
96
97 # Set permissions via environment variables
98 permissions |= global_permissions() # CONFIG_GLOBAL
99 permissions |= chat_type_permissions(ctype) # CONFIG_GROUP/CHANNEL/BOT/PRIVATE = 0
100 permissions |= chat_permissions(cid) # CONFIG_C{CID} = {"ai": false}
101 permissions |= chat_thread_permissions(cid, tid) # CONFIG_C{CID}_T{TID} = {"ai": false}
102 permissions |= user_permissions(uid) # CONFIG_U{UID} = {"ai": false}
103
104 permissions |= chat_user_permissions(cid, uid) # CONFIG_C{CID}_U{UID} = {"ai": false}
105 permissions |= chat_thread_user_permissions(cid, tid, uid) # CONFIG_C{CID}_T{TID}_U{UID} = {"ai": false}
106 return permissions
107
108
109@cache.memoize(ttl=0)
110def chat_type_permissions(ctype: str) -> dict:
111 """Set chat type permissions."""
112 if not os.getenv(f"CONFIG_{ctype}"):
113 return {}
114 try:
115 permission = json.loads(os.environ[f"CONFIG_{ctype}"])
116 except json.JSONDecodeError:
117 logger.error(f"Invalid JSON in CONFIG_{ctype} config")
118 return {}
119 logger.success(f"ChatType={ctype} permissions: {permission}")
120 return permission
121
122
123@cache.memoize(ttl=0)
124def global_permissions() -> dict:
125 """Set global permissions.
126
127 CONFIG_GLOBAL={"prepend_sender_user": false, "ai": false, "twitter_provider": "vxtwitter-fxtwitter"}
128 # set prepend_sender_user to False
129 # set ai to False
130 # set twitter_provider to "vxtwitter-fxtwitter"
131 """
132 if not os.getenv("CONFIG_GLOBAL"):
133 return {}
134 try:
135 permission = json.loads(os.environ["CONFIG_GLOBAL"])
136 except json.JSONDecodeError:
137 logger.error("Invalid JSON in CONFIG_GLOBAL config")
138 return {}
139 logger.success(f"Global permissions: {permission}")
140 return permission
141
142
143@cache.memoize(ttl=0)
144def chat_permissions(cid: str) -> dict:
145 """Set chat permissions.
146
147 CONFIG_C{CID}={"ai": false, "twitter_provider": "vxtwitter-fxtwitter"}
148 # config for chat id: {CID}
149 """
150 if not os.getenv(f"CONFIG_C{cid}"):
151 return {}
152 try:
153 permission = json.loads(os.environ[f"CONFIG_C{cid}"])
154 except json.JSONDecodeError:
155 logger.error(f"Invalid JSON in CONFIG_C{cid} config")
156 return {}
157 logger.success(f"CID={cid} permissions: {permission}")
158 return permission
159
160
161@cache.memoize(ttl=0)
162def chat_thread_permissions(cid: str, tid: str) -> dict:
163 """Set chat thread permissions.
164
165 CONFIG_C{CID}_T{TID}={"ai": false, "twitter_provider": "vxtwitter-fxtwitter"}
166 # config for chat {CID}, thread id: {TID}
167 """
168 if not os.getenv(f"CONFIG_C{cid}_T{tid}"):
169 return {}
170 try:
171 permission = json.loads(os.environ[f"CONFIG_C{cid}_T{tid}"])
172 except json.JSONDecodeError:
173 logger.error(f"Invalid JSON in CONFIG_C{cid}_T{tid} config")
174 return {}
175 logger.success(f"CID={cid}_TID={tid} permissions: {permission}")
176 return permission
177
178
179def user_permissions(uid: str | int) -> dict:
180 """Set user permissions.
181
182 CONFIG_U{UID}={"ai": false, "twitter_provider": "vxtwitter-fxtwitter"}
183 # config for user {UID}
184 """
185 if not os.getenv(f"CONFIG_U{uid}"):
186 return {}
187 try:
188 permission = json.loads(os.environ[f"CONFIG_U{uid}"])
189 except json.JSONDecodeError:
190 logger.error(f"Invalid JSON in CONFIG_U{uid} config")
191 return {}
192 logger.success(f"UID={uid} permissions: {permission}")
193 return permission
194
195
196def chat_user_permissions(cid: str, uid: str | int) -> dict:
197 """Set chat user permissions.
198
199 CONFIG_C{CID}_U{UID}={"ai": false, "twitter_provider": "vxtwitter-fxtwitter"}
200 # config for chat {CID}, user {UID}
201 """
202 if not os.getenv(f"CONFIG_C{cid}_U{uid}"):
203 return {}
204 try:
205 permission = json.loads(os.environ[f"CONFIG_C{cid}_U{uid}"])
206 except json.JSONDecodeError:
207 logger.error(f"Invalid JSON in CONFIG_C{cid}_U{uid} config")
208 return {}
209 logger.success(f"CID={cid}_UID={uid} permissions: {permission}")
210 return permission
211
212
213def chat_thread_user_permissions(cid: str, tid: str | int, uid: str | int) -> dict:
214 """Set chat thread user permissions.
215
216 CONFIG_C{CID}_T{TID}_U{UID}={"ai": false, "twitter_provider": "vxtwitter-fxtwitter"}
217 # config for chat {CID}, thread id: {TID}, user {UID}
218 """
219 if not os.getenv(f"CONFIG_C{cid}_T{tid}_U{uid}"):
220 return {}
221 try:
222 permission = json.loads(os.environ[f"CONFIG_C{cid}_T{tid}_U{uid}"])
223 except json.JSONDecodeError:
224 logger.error(f"Invalid JSON in CONFIG_C{cid}_T{tid}_U{uid} config")
225 return {}
226 logger.success(f"CID={cid}_TID={tid}_UID={uid} permissions: {permission}")
227 return permission