瀏覽代碼

Adding video link preview cog

main
Rocketsoup 2 週之前
父節點
當前提交
b9c0cf7dcf
共有 5 個文件被更改,包括 245 次插入4 次删除
  1. 12
    4
      main.py
  2. 1
    0
      requirements.txt
  3. 2
    0
      rocketbot/bot.py
  4. 205
    0
      rocketbot/cogs/videopreviewcog.py
  5. 25
    0
      rocketbot/utils.py

+ 12
- 4
main.py 查看文件

5
 Author: Ian Albert (@rocketsoup)
5
 Author: Ian Albert (@rocketsoup)
6
 Date: 2021-11-11
6
 Date: 2021-11-11
7
 """
7
 """
8
-import asyncio
8
+import sys
9
 
9
 
10
-from config import CONFIG
11
-from rocketbot.bot import start_bot
12
-from rocketbot.utils import bot_log
10
+MIN_PYTHON_VERSION = (3, 10, 0)
11
+if sys.version_info < MIN_PYTHON_VERSION:
12
+	ver = sys.version_info
13
+	raise RuntimeError(f'rocketbot requires Python {MIN_PYTHON_VERSION[0]}.{MIN_PYTHON_VERSION[1]}.{MIN_PYTHON_VERSION[2]} '
14
+		f'or greater. Detected {ver[0]}.{ver[1]}.{ver[2]}.')
15
+
16
+import asyncio  # noqa: E402
17
+
18
+from config import CONFIG  # noqa: E402
19
+from rocketbot.bot import start_bot  # noqa: E402
20
+from rocketbot.utils import bot_log  # noqa: E402
13
 
21
 
14
 CURRENT_CONFIG_VERSION = 4
22
 CURRENT_CONFIG_VERSION = 4
15
 if (CONFIG.get('__config_version') or 0) < CURRENT_CONFIG_VERSION:
23
 if (CONFIG.get('__config_version') or 0) < CURRENT_CONFIG_VERSION:

+ 1
- 0
requirements.txt 查看文件

1
 discord.py == 2.6.4
1
 discord.py == 2.6.4
2
+yt-dlp == 2026.02.04

+ 2
- 0
rocketbot/bot.py 查看文件

103
 	from rocketbot.cogs.patterncog import PatternCog
103
 	from rocketbot.cogs.patterncog import PatternCog
104
 	from rocketbot.cogs.urlspamcog import URLSpamCog
104
 	from rocketbot.cogs.urlspamcog import URLSpamCog
105
 	from rocketbot.cogs.usernamecog import UsernamePatternCog
105
 	from rocketbot.cogs.usernamecog import UsernamePatternCog
106
+	from rocketbot.cogs.videopreviewcog import VideoPreviewCog
106
 
107
 
107
 	bot_log(None, None, 'Bot initializing...')
108
 	bot_log(None, None, 'Bot initializing...')
108
 
109
 
121
 	await rocketbot.add_cog(PatternCog(rocketbot))
122
 	await rocketbot.add_cog(PatternCog(rocketbot))
122
 	await rocketbot.add_cog(URLSpamCog(rocketbot))
123
 	await rocketbot.add_cog(URLSpamCog(rocketbot))
123
 	await rocketbot.add_cog(UsernamePatternCog(rocketbot))
124
 	await rocketbot.add_cog(UsernamePatternCog(rocketbot))
125
+	await rocketbot.add_cog(VideoPreviewCog(rocketbot))
124
 
126
 
125
 	await rocketbot.start(CONFIG['client_token'], reconnect=True)
127
 	await rocketbot.start(CONFIG['client_token'], reconnect=True)
126
 	print('\nBot aborted')
128
 	print('\nBot aborted')

+ 205
- 0
rocketbot/cogs/videopreviewcog.py 查看文件

1
+import asyncio
2
+import json
3
+import re
4
+import subprocess
5
+
6
+from discord import Message
7
+from discord.ext.commands import Cog
8
+
9
+from rocketbot.cogs.basecog import BaseCog
10
+from rocketbot.cogsetting import CogSetting
11
+from rocketbot.utils import (
12
+	blockquote_markdown,
13
+	format_bytes,
14
+	suppress_markdown_url_previews,
15
+)
16
+
17
+
18
+def filter_video_format(format: dict) -> bool:
19
+	if format.get('resolution') == 'audio only':
20
+		return False
21
+	if format.get('format_note') == 'DASH audio':
22
+		return False
23
+	return True
24
+
25
+def rank_video_format(format: dict) -> tuple:
26
+	content = 0
27
+	if format.get('resolution') == 'audio only':
28
+		content = 1
29
+	elif format.get('format_note') == 'DASH audio':
30
+		content = 1
31
+	elif format.get('format_note') == 'DASH video':
32
+		content = 2
33
+	else:
34
+		content = 3  # both I guess! multiplexed formats don't seem clearly labeled
35
+	res = (format.get('width') or 0) + (format.get('height') or 0)
36
+	return (content, res)
37
+
38
+class MessageLink:
39
+	url: str
40
+	spoiler: bool = False
41
+	link_type: str = 'unknown'
42
+
43
+	def __init__(self, url: str, link_type: str, spoiler: bool = False):
44
+		self.url = url
45
+		self.link_type = link_type
46
+		self.spoiler = spoiler
47
+
48
+class VideoPreviewCog(BaseCog, name='Video Link Previews'):
49
+	SETTING_ENABLED = CogSetting(
50
+		'enabled',
51
+		bool,
52
+		default_value=False,
53
+		brief='Video link previews',
54
+		description='Whether links to certain social media videos should show previews.',
55
+	)
56
+	SETTING_INSTAGRAM = CogSetting(
57
+		'instagram',
58
+		bool,
59
+		default_value=False,
60
+		brief='whether to show video previews for Instagram links',
61
+		description='For both regular posts and reels',
62
+	)
63
+	SETTING_FACEBOOK = CogSetting(
64
+		'facebook',
65
+		bool,
66
+		default_value=False,
67
+		brief='whether to show video previews for Facebook links',
68
+	)
69
+	SETTING_TWITTER = CogSetting(
70
+		'twitter',
71
+		bool,
72
+		default_value=False,
73
+		brief='whether to show video previews for Twitter links',
74
+	)
75
+
76
+	REGEX_INSTAGRAM_POST = r'https?:\/\/(?:www\.)?instagram\.com\/(?:p|reel)\/[a-zA-Z0-9_-]+\/?'
77
+	REGEX_FACEBOOK_POST = r'https?:\/\/(?:www\.)?facebook\.com\/share\/[rv]\/[a-zA-Z0-9_-]+\/?'
78
+	REGEX_TWITTER_POST = r'https?:\/\/(?:twitter|x)\.com\/[a-zA-Z0-9_-]+/status/[0-9]+'
79
+
80
+	REGEX_SPOILERS = '\|\|.+\|\|'
81
+
82
+	# Best video and best audio, mp4 format with m4a audio
83
+	FORMATS = 'bv*[ext=mp4]+ba[ext=m4a]/' \
84
+		'b[ext=mp4]/' \
85
+		'bv*[ext=mp4]+ba[ext=m4a]/' \
86
+		'b[ext=mp4]/' \
87
+		'bv*+ba/' \
88
+		'b'
89
+
90
+	def __init__(self, bot):
91
+		super().__init__(
92
+			bot,
93
+			config_prefix='linkpreview',
94
+			short_description='Manages video link preview behavior.',
95
+		)
96
+		Self = VideoPreviewCog
97
+		self.add_setting(Self.SETTING_ENABLED)
98
+		self.add_setting(Self.SETTING_INSTAGRAM)
99
+		self.add_setting(Self.SETTING_FACEBOOK)
100
+		self.add_setting(Self.SETTING_TWITTER)
101
+
102
+	@Cog.listener()
103
+	async def on_message(self, message: Message):
104
+		"""Event listener"""
105
+		if message.author is None or \
106
+				message.author.bot or \
107
+				message.guild is None or \
108
+				message.channel is None or \
109
+				message.content is None:
110
+			return
111
+		if not self.get_guild_setting(message.guild, self.SETTING_ENABLED):
112
+			return
113
+		links = self._get_previewable_links(message)
114
+		if len(links) == 0:
115
+			return
116
+		await self._wait_for_preview(message, links)
117
+
118
+
119
+		# TODO: Make this just link to the raw video file if possible (yt-dlp --get-url)
120
+
121
+	def _get_previewable_links(self, message: Message) -> list[MessageLink]:
122
+		Self = VideoPreviewCog
123
+		links: list[MessageLink] = []
124
+		content: str = message.content
125
+		has_spoilers = re.match(Self.REGEX_SPOILERS, content) is not None
126
+		if self.get_guild_setting(message.guild, Self.SETTING_INSTAGRAM):
127
+			for link in re.findall(Self.REGEX_INSTAGRAM_POST, content):
128
+				links.append(MessageLink(link, 'instagram', has_spoilers))
129
+		if self.get_guild_setting(message.guild, Self.SETTING_FACEBOOK):
130
+			for link in re.findall(Self.REGEX_FACEBOOK_POST, content):
131
+				links.append(MessageLink(link, 'facebook', has_spoilers))
132
+		if self.get_guild_setting(message.guild, Self.SETTING_TWITTER):
133
+			for link in re.findall(Self.REGEX_TWITTER_POST, content):
134
+				links.append(MessageLink(link, 'twitter', has_spoilers))
135
+		# TODO: Custom patterns
136
+		return links
137
+
138
+	async def _wait_for_preview(self, message: Message, links: list[MessageLink]):
139
+		await asyncio.sleep(3)
140
+		# Look for embeds already showing the video
141
+		self.log(message.guild, "Checking message for embeds")
142
+		for embed in message.embeds:
143
+			if embed.video.url:
144
+				# If there's any video, skip downloading any previews
145
+				self.log(message.guild, "Message already has a video. Skipping this message.")
146
+				return
147
+		await self._fetch_previews(message, links)
148
+
149
+	async def _fetch_previews(self, message: Message, links: list[MessageLink]):
150
+		promises = []
151
+		for link in links:
152
+			promises.append(self._fetch_preview(message, link))
153
+		await asyncio.gather(*promises)
154
+
155
+	async def _fetch_preview(self, message: Message, link: MessageLink):
156
+		result = subprocess.run(
157
+			[
158
+				'yt-dlp',
159
+				'--skip-download',
160
+				'--dump-single-json',
161
+				link.url,
162
+			],
163
+			stdout=subprocess.PIPE,
164
+			stderr=subprocess.PIPE,
165
+			universal_newlines=True
166
+		)
167
+		if result.returncode != 0:
168
+			self.log(message.guild, "Fetching link info JSON failed. Skipping preview.")
169
+			self.log(message.guild, result.stderr)
170
+			return
171
+		try:
172
+			info: dict = json.loads(result.stdout)
173
+		except Exception as e:
174
+			self.log(message.guild, f"Error parsing info.json. Skipping preview. {e}")
175
+			return
176
+		description = info.get('description') or ''
177
+		formats: list[dict] = info.get('formats') or []
178
+		self.log(message.guild, f"Found {len(formats)} formats")
179
+		formats = list(filter(filter_video_format, formats))
180
+		self.log(message.guild, f"Filtered to {len(formats)} formats")
181
+		sorted_formats: list[dict] = sorted(formats, key=rank_video_format, reverse=True)
182
+		if len(sorted_formats) == 0:
183
+			self.log(message.guild, f"No eligible formats for URL {link.url}")
184
+			return
185
+		best_format: dict = sorted_formats[0]
186
+		self.log(message.guild, f"Best format is id {best_format.get('format_id')}")
187
+		video_url: str = best_format.get('url')
188
+		link_description: str = "video"
189
+		if (best_format.get('width') or 0) > 0 and (best_format.get('height') or 0) > 0:
190
+			link_description += f", {best_format.get('width')}×{best_format.get('height')}"
191
+		if (best_format.get('filesize') or 0) > 0:
192
+			link_description += f", {format_bytes(best_format.get('filesize'))}"
193
+		elif (best_format.get('filesize_approx') or 0) > 0:
194
+			link_description += f", {format_bytes(best_format.get('filesize_approx'))}"
195
+		content = "Hmm, video preview didn't load. Let's try this."
196
+		if len(description) > 0:
197
+			content += "\n\n" + blockquote_markdown(suppress_markdown_url_previews(description)) + "\n"
198
+		if link.spoiler:
199
+			content += f"\n||[{link_description}]({video_url})||"
200
+		else:
201
+			content += f"\n[{link_description}]({video_url})"
202
+		await message.reply(
203
+			content,
204
+			mention_author=False
205
+		)

+ 25
- 0
rocketbot/utils.py 查看文件

181
 	"""Indents a block of Markdown by one level."""
181
 	"""Indents a block of Markdown by one level."""
182
 	return '    ' + (markdown.replace('\n', '\n    '))
182
 	return '    ' + (markdown.replace('\n', '\n    '))
183
 
183
 
184
+def suppress_markdown_url_previews(markdown: str) -> str:
185
+	"""Finds URLs in markdown and encloses them in <...> to suppress the preview."""
186
+	return re.sub(r'(?<!<)(https?://\S+)(?!>)', '<\\1>', markdown)
187
+
188
+def format_bytes(size: int) -> str:
189
+	"""Formats s size in bytes to a human readable description (e.g. "3.2 KiB")"""
190
+	if size < 0:
191
+		size = 0
192
+	kib = 1024
193
+	mib = kib * kib
194
+	gib = mib * kib
195
+	if size < kib:
196
+		return f"{size:,} bytes"
197
+	if size < 10 * kib:
198
+		return f"{size/kib:,.1f} KiB"
199
+	if size < mib:
200
+		return f"{size/kib:,.0f} KiB"
201
+	if size < 10 * mib:
202
+		return f"{size/mib:,.1f} MiB"
203
+	if size < gib:
204
+		return f"{size/mib:,.0f} MiB"
205
+	if size < 10 * gib:
206
+		return f"{size/gib:,.1f} GiB"
207
+	return f"{size/gib:,.0f} GiB"
208
+
184
 MOD_PERMISSIONS: Permissions = Permissions(Permissions.manage_messages.flag)
209
 MOD_PERMISSIONS: Permissions = Permissions(Permissions.manage_messages.flag)
185
 
210
 
186
 class TimeDeltaTransformer(Transformer):
211
 class TimeDeltaTransformer(Transformer):

Loading…
取消
儲存