Beta 1
* Screenshot (unpolished, functional) * Dig * Traceroute * Whois * Ping * yt-dl
This commit is contained in:
parent
8e3fea7f14
commit
4c4ddbf6c1
12 changed files with 1157 additions and 0 deletions
10
.idea/college-bot-2.0.iml
Normal file
10
.idea/college-bot-2.0.iml
Normal file
|
@ -0,0 +1,10 @@
|
||||||
|
<?xml version="1.0" encoding="UTF-8"?>
|
||||||
|
<module type="PYTHON_MODULE" version="4">
|
||||||
|
<component name="NewModuleRootManager">
|
||||||
|
<content url="file://$MODULE_DIR$">
|
||||||
|
<excludeFolder url="file://$MODULE_DIR$/.venv" />
|
||||||
|
</content>
|
||||||
|
<orderEntry type="inheritedJdk" />
|
||||||
|
<orderEntry type="sourceFolder" forTests="false" />
|
||||||
|
</component>
|
||||||
|
</module>
|
6
.idea/inspectionProfiles/profiles_settings.xml
Normal file
6
.idea/inspectionProfiles/profiles_settings.xml
Normal file
|
@ -0,0 +1,6 @@
|
||||||
|
<component name="InspectionProjectProfileManager">
|
||||||
|
<settings>
|
||||||
|
<option name="USE_PROJECT_PROFILE" value="false" />
|
||||||
|
<version value="1.0" />
|
||||||
|
</settings>
|
||||||
|
</component>
|
7
.idea/misc.xml
Normal file
7
.idea/misc.xml
Normal file
|
@ -0,0 +1,7 @@
|
||||||
|
<?xml version="1.0" encoding="UTF-8"?>
|
||||||
|
<project version="4">
|
||||||
|
<component name="Black">
|
||||||
|
<option name="sdkName" value="Python 3.11 (college-bot-2.0)" />
|
||||||
|
</component>
|
||||||
|
<component name="ProjectRootManager" version="2" project-jdk-name="Python 3.11 (college-bot-2.0)" project-jdk-type="Python SDK" />
|
||||||
|
</project>
|
8
.idea/modules.xml
Normal file
8
.idea/modules.xml
Normal file
|
@ -0,0 +1,8 @@
|
||||||
|
<?xml version="1.0" encoding="UTF-8"?>
|
||||||
|
<project version="4">
|
||||||
|
<component name="ProjectModuleManager">
|
||||||
|
<modules>
|
||||||
|
<module fileurl="file://$PROJECT_DIR$/.idea/college-bot-2.0.iml" filepath="$PROJECT_DIR$/.idea/college-bot-2.0.iml" />
|
||||||
|
</modules>
|
||||||
|
</component>
|
||||||
|
</project>
|
6
.idea/vcs.xml
Normal file
6
.idea/vcs.xml
Normal file
|
@ -0,0 +1,6 @@
|
||||||
|
<?xml version="1.0" encoding="UTF-8"?>
|
||||||
|
<project version="4">
|
||||||
|
<component name="VcsDirectoryMappings">
|
||||||
|
<mapping directory="$PROJECT_DIR$" vcs="Git" />
|
||||||
|
</component>
|
||||||
|
</project>
|
102
.idea/workspace.xml
Normal file
102
.idea/workspace.xml
Normal file
|
@ -0,0 +1,102 @@
|
||||||
|
<?xml version="1.0" encoding="UTF-8"?>
|
||||||
|
<project version="4">
|
||||||
|
<component name="AutoImportSettings">
|
||||||
|
<option name="autoReloadType" value="SELECTIVE" />
|
||||||
|
</component>
|
||||||
|
<component name="ChangeListManager">
|
||||||
|
<list default="true" id="aa6d62a8-d64d-4a60-a85f-8d9fa52b6b49" name="Changes" comment="Update gitignore">
|
||||||
|
<change afterPath="$PROJECT_DIR$/.idea/college-bot-2.0.iml" afterDir="false" />
|
||||||
|
<change afterPath="$PROJECT_DIR$/.idea/inspectionProfiles/profiles_settings.xml" afterDir="false" />
|
||||||
|
<change afterPath="$PROJECT_DIR$/.idea/misc.xml" afterDir="false" />
|
||||||
|
<change afterPath="$PROJECT_DIR$/.idea/modules.xml" afterDir="false" />
|
||||||
|
<change afterPath="$PROJECT_DIR$/.idea/vcs.xml" afterDir="false" />
|
||||||
|
<change afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
|
||||||
|
<change afterPath="$PROJECT_DIR$/cogs/net.py" afterDir="false" />
|
||||||
|
<change afterPath="$PROJECT_DIR$/cogs/screenshot.py" afterDir="false" />
|
||||||
|
<change afterPath="$PROJECT_DIR$/cogs/ytdl.py" afterDir="false" />
|
||||||
|
<change afterPath="$PROJECT_DIR$/config.example.toml" afterDir="false" />
|
||||||
|
<change afterPath="$PROJECT_DIR$/main.py" afterDir="false" />
|
||||||
|
<change afterPath="$PROJECT_DIR$/requirements.txt" afterDir="false" />
|
||||||
|
</list>
|
||||||
|
<option name="SHOW_DIALOG" value="false" />
|
||||||
|
<option name="HIGHLIGHT_CONFLICTS" value="true" />
|
||||||
|
<option name="HIGHLIGHT_NON_ACTIVE_CHANGELIST" value="false" />
|
||||||
|
<option name="LAST_RESOLUTION" value="IGNORE" />
|
||||||
|
</component>
|
||||||
|
<component name="FileTemplateManagerImpl">
|
||||||
|
<option name="RECENT_TEMPLATES">
|
||||||
|
<list>
|
||||||
|
<option value="Python Script" />
|
||||||
|
</list>
|
||||||
|
</option>
|
||||||
|
</component>
|
||||||
|
<component name="Git.Settings">
|
||||||
|
<option name="RECENT_GIT_ROOT_PATH" value="$PROJECT_DIR$" />
|
||||||
|
</component>
|
||||||
|
<component name="ProjectColorInfo">{
|
||||||
|
"associatedIndex": 8
|
||||||
|
}</component>
|
||||||
|
<component name="ProjectId" id="2aMdul0tGldqewjmOMm9WhRXLP9" />
|
||||||
|
<component name="ProjectLevelVcsManager" settingsEditedManually="true">
|
||||||
|
<ConfirmationsSetting value="2" id="Add" />
|
||||||
|
</component>
|
||||||
|
<component name="ProjectViewState">
|
||||||
|
<option name="hideEmptyMiddlePackages" value="true" />
|
||||||
|
<option name="showLibraryContents" value="true" />
|
||||||
|
</component>
|
||||||
|
<component name="PropertiesComponent">{
|
||||||
|
"keyToString": {
|
||||||
|
"ASKED_ADD_EXTERNAL_FILES": "true",
|
||||||
|
"RunOnceActivity.OpenProjectViewOnStart": "true",
|
||||||
|
"RunOnceActivity.ShowReadmeOnStart": "true",
|
||||||
|
"git-widget-placeholder": "master",
|
||||||
|
"node.js.detected.package.eslint": "true",
|
||||||
|
"node.js.detected.package.tslint": "true",
|
||||||
|
"node.js.selected.package.eslint": "(autodetect)",
|
||||||
|
"node.js.selected.package.tslint": "(autodetect)",
|
||||||
|
"nodejs_package_manager_path": "npm",
|
||||||
|
"settings.editor.selected.configurable": "settings.sync",
|
||||||
|
"vue.rearranger.settings.migration": "true"
|
||||||
|
}
|
||||||
|
}</component>
|
||||||
|
<component name="SharedIndexes">
|
||||||
|
<attachedChunks>
|
||||||
|
<set>
|
||||||
|
<option value="bundled-python-sdk-50da183f06c8-2887949eec09-com.jetbrains.pycharm.pro.sharedIndexes.bundled-PY-233.13135.95" />
|
||||||
|
</set>
|
||||||
|
</attachedChunks>
|
||||||
|
</component>
|
||||||
|
<component name="SpellCheckerSettings" RuntimeDictionaries="0" Folders="0" CustomDictionaries="0" DefaultDictionary="application-level" UseSingleDictionary="true" transferred="true" />
|
||||||
|
<component name="TaskManager">
|
||||||
|
<task active="true" id="Default" summary="Default task">
|
||||||
|
<changelist id="aa6d62a8-d64d-4a60-a85f-8d9fa52b6b49" name="Changes" comment="" />
|
||||||
|
<created>1704132606619</created>
|
||||||
|
<option name="number" value="Default" />
|
||||||
|
<option name="presentableId" value="Default" />
|
||||||
|
<updated>1704132606619</updated>
|
||||||
|
<workItem from="1704132608091" duration="7680000" />
|
||||||
|
<workItem from="1704140854074" duration="30000" />
|
||||||
|
<workItem from="1704140889356" duration="166000" />
|
||||||
|
<workItem from="1704141061129" duration="2495000" />
|
||||||
|
<workItem from="1704224862150" duration="9650000" />
|
||||||
|
</task>
|
||||||
|
<task id="LOCAL-00001" summary="Update gitignore">
|
||||||
|
<option name="closed" value="true" />
|
||||||
|
<created>1704234449168</created>
|
||||||
|
<option name="number" value="00001" />
|
||||||
|
<option name="presentableId" value="LOCAL-00001" />
|
||||||
|
<option name="project" value="LOCAL" />
|
||||||
|
<updated>1704234449168</updated>
|
||||||
|
</task>
|
||||||
|
<option name="localTasksCounter" value="2" />
|
||||||
|
<servers />
|
||||||
|
</component>
|
||||||
|
<component name="TypeScriptGeneratedFilesManager">
|
||||||
|
<option name="version" value="3" />
|
||||||
|
</component>
|
||||||
|
<component name="VcsManagerConfiguration">
|
||||||
|
<option name="ADD_EXTERNAL_FILES_SILENTLY" value="true" />
|
||||||
|
<MESSAGE value="Update gitignore" />
|
||||||
|
<option name="LAST_COMMIT_MESSAGE" value="Update gitignore" />
|
||||||
|
</component>
|
||||||
|
</project>
|
248
cogs/net.py
Normal file
248
cogs/net.py
Normal file
|
@ -0,0 +1,248 @@
|
||||||
|
import io
|
||||||
|
import os
|
||||||
|
import typing
|
||||||
|
|
||||||
|
import discord
|
||||||
|
from rich.console import Console
|
||||||
|
import time
|
||||||
|
import re
|
||||||
|
from dns import asyncresolver
|
||||||
|
from rich.tree import Tree
|
||||||
|
import asyncio
|
||||||
|
from discord.ext import commands
|
||||||
|
|
||||||
|
|
||||||
|
class NetworkCog(commands.Cog):
|
||||||
|
def __init__(self, bot: commands.Bot):
|
||||||
|
self.bot = bot
|
||||||
|
|
||||||
|
@commands.slash_command()
|
||||||
|
async def ping(self, ctx: discord.ApplicationContext, target: str = None):
|
||||||
|
"""Get the bot's latency, or the network latency to a target."""
|
||||||
|
if target is None:
|
||||||
|
return await ctx.respond(f"Pong! {round(self.bot.latency * 1000)}ms")
|
||||||
|
else:
|
||||||
|
await ctx.defer()
|
||||||
|
process = await asyncio.create_subprocess_exec(
|
||||||
|
"ping",
|
||||||
|
"-c",
|
||||||
|
"5",
|
||||||
|
target,
|
||||||
|
stdout=asyncio.subprocess.PIPE,
|
||||||
|
stderr=asyncio.subprocess.PIPE,
|
||||||
|
)
|
||||||
|
stdout, stderr = await process.communicate()
|
||||||
|
paginator = commands.Paginator()
|
||||||
|
|
||||||
|
for line in stdout.splitlines():
|
||||||
|
paginator.add_line(line.decode("utf-8"))
|
||||||
|
for line in stderr.splitlines():
|
||||||
|
paginator.add_line("[STDERR] " + line.decode("utf-8"))
|
||||||
|
|
||||||
|
for page in paginator.pages:
|
||||||
|
await ctx.respond(page)
|
||||||
|
|
||||||
|
@commands.slash_command()
|
||||||
|
async def whois(self, ctx: discord.ApplicationContext, target: str):
|
||||||
|
"""Get information about a user."""
|
||||||
|
|
||||||
|
async def run_command(with_disclaimer: bool = False):
|
||||||
|
args = [] if with_disclaimer else ["-H"]
|
||||||
|
process = await asyncio.create_subprocess_exec(
|
||||||
|
"whois",
|
||||||
|
*args,
|
||||||
|
target,
|
||||||
|
stdout=asyncio.subprocess.PIPE,
|
||||||
|
stderr=asyncio.subprocess.PIPE,
|
||||||
|
)
|
||||||
|
so, se = await process.communicate()
|
||||||
|
return so, se, process.returncode
|
||||||
|
|
||||||
|
await ctx.defer()
|
||||||
|
paginator = commands.Paginator()
|
||||||
|
redacted = io.BytesIO()
|
||||||
|
stdout, stderr, status = await run_command()
|
||||||
|
|
||||||
|
def decide(ln: str) -> typing.Optional[bool]:
|
||||||
|
if ln.startswith(">>> Last update"):
|
||||||
|
return
|
||||||
|
if "REDACTED" in ln or "Please query the WHOIS server of the owning registrar" in ln or ":" not in ln:
|
||||||
|
return False
|
||||||
|
else:
|
||||||
|
return True
|
||||||
|
|
||||||
|
for line in stdout.decode().splitlines():
|
||||||
|
line = line.strip()
|
||||||
|
if not line:
|
||||||
|
continue
|
||||||
|
a = decide(line)
|
||||||
|
if a:
|
||||||
|
paginator.add_line(line)
|
||||||
|
elif a is None:
|
||||||
|
redacted.write(b"[STDERR] " + line.encode() + b"\n")
|
||||||
|
|
||||||
|
for line in stderr.decode().splitlines():
|
||||||
|
line = line.strip()
|
||||||
|
if not line:
|
||||||
|
continue
|
||||||
|
a = decide(line)
|
||||||
|
if a:
|
||||||
|
paginator.add_line("[STDERR] " + line)
|
||||||
|
elif a is None:
|
||||||
|
redacted.write(b"[STDERR] " + line.encode() + b"\n")
|
||||||
|
|
||||||
|
if not paginator.pages:
|
||||||
|
stdout, stderr, status = await run_command(with_disclaimer=True)
|
||||||
|
if not any((stdout, stderr)):
|
||||||
|
return await ctx.respond(f"No output was returned with status code {status}.")
|
||||||
|
file = io.BytesIO()
|
||||||
|
file.write(stdout)
|
||||||
|
if stderr:
|
||||||
|
file.write(b"\n----- STDERR -----\n")
|
||||||
|
file.write(stderr)
|
||||||
|
file.seek(0)
|
||||||
|
return await ctx.respond(
|
||||||
|
"Seemingly all output was filtered. Returning raw command output.",
|
||||||
|
file=discord.File(file, "whois.txt")
|
||||||
|
)
|
||||||
|
|
||||||
|
for page in paginator.pages:
|
||||||
|
await ctx.respond(page)
|
||||||
|
if redacted.getvalue():
|
||||||
|
redacted.seek(0)
|
||||||
|
await ctx.respond(file=discord.File(redacted, "redacted.txt"))
|
||||||
|
|
||||||
|
@commands.slash_command()
|
||||||
|
async def dig(
|
||||||
|
self,
|
||||||
|
ctx: discord.ApplicationContext,
|
||||||
|
domain: str,
|
||||||
|
_type: discord.Option(
|
||||||
|
str,
|
||||||
|
name="type",
|
||||||
|
default="A",
|
||||||
|
choices=[
|
||||||
|
"A",
|
||||||
|
"AAAA",
|
||||||
|
"ANY",
|
||||||
|
"AXFR",
|
||||||
|
"CNAME",
|
||||||
|
"HINFO",
|
||||||
|
"LOC",
|
||||||
|
"MX",
|
||||||
|
"NS",
|
||||||
|
"PTR",
|
||||||
|
"SOA",
|
||||||
|
"SRV",
|
||||||
|
"TXT",
|
||||||
|
],
|
||||||
|
),
|
||||||
|
):
|
||||||
|
"""Looks up a domain name"""
|
||||||
|
await ctx.defer()
|
||||||
|
if re.search(r"\s+", domain):
|
||||||
|
return await ctx.respond("Domain name cannot contain spaces.")
|
||||||
|
try:
|
||||||
|
response = await asyncresolver.resolve(
|
||||||
|
domain,
|
||||||
|
_type.upper(),
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
return await ctx.respond(f"Error: {e}")
|
||||||
|
res = response
|
||||||
|
tree = Tree(f"DNS Lookup for {domain}")
|
||||||
|
for record in res:
|
||||||
|
record_tree = tree.add(f"{record.rdtype.name} Record")
|
||||||
|
record_tree.add(f"Name: {res.name}")
|
||||||
|
record_tree.add(f"Value: {record.to_text()}")
|
||||||
|
console = Console()
|
||||||
|
with console.capture() as capture:
|
||||||
|
console.print(tree)
|
||||||
|
text = capture.get()
|
||||||
|
paginator = commands.Paginator(prefix="```", suffix="```")
|
||||||
|
for line in text.splitlines():
|
||||||
|
paginator.add_line(line)
|
||||||
|
paginator.add_line(empty=True)
|
||||||
|
paginator.add_line(f"Exit code: {0}")
|
||||||
|
paginator.add_line(f"DNS Server used: {res.nameserver}")
|
||||||
|
for page in paginator.pages:
|
||||||
|
await ctx.respond(page)
|
||||||
|
|
||||||
|
@commands.slash_command()
|
||||||
|
async def traceroute(
|
||||||
|
self,
|
||||||
|
ctx: discord.ApplicationContext,
|
||||||
|
url: str,
|
||||||
|
port: discord.Option(int, description="Port to use", default=None),
|
||||||
|
ping_type: discord.Option(
|
||||||
|
str,
|
||||||
|
name="ping-type",
|
||||||
|
description="Type of ping to use. See `traceroute --help`",
|
||||||
|
choices=["icmp", "tcp", "udp", "udplite", "dccp", "default"],
|
||||||
|
default="default",
|
||||||
|
),
|
||||||
|
use_ip_version: discord.Option(
|
||||||
|
str, name="ip-version", description="IP version to use.", choices=["ipv4", "ipv6"], default="ipv4"
|
||||||
|
),
|
||||||
|
max_ttl: discord.Option(int, name="ttl", description="Max number of hops", default=30),
|
||||||
|
):
|
||||||
|
"""Performs a traceroute request."""
|
||||||
|
await ctx.defer()
|
||||||
|
if re.search(r"\s+", url):
|
||||||
|
return await ctx.respond("URL cannot contain spaces.")
|
||||||
|
|
||||||
|
args = ["sudo", "-E", "-n", "traceroute"]
|
||||||
|
flags = {
|
||||||
|
"ping_type": {
|
||||||
|
"icmp": "-I",
|
||||||
|
"tcp": "-T",
|
||||||
|
"udp": "-U",
|
||||||
|
"udplite": "-UL",
|
||||||
|
"dccp": "-D",
|
||||||
|
},
|
||||||
|
"use_ip_version": {"ipv4": "-4", "ipv6": "-6"},
|
||||||
|
}
|
||||||
|
|
||||||
|
if ping_type == "default" or os.getuid() == 0:
|
||||||
|
args = args[3:] # removes sudo
|
||||||
|
else:
|
||||||
|
args.append(flags["ping_type"][ping_type])
|
||||||
|
args.append(flags["use_ip_version"][use_ip_version])
|
||||||
|
args.append("-m")
|
||||||
|
args.append(str(max_ttl))
|
||||||
|
if port is not None:
|
||||||
|
args.append("-p")
|
||||||
|
args.append(str(port))
|
||||||
|
args.append(url)
|
||||||
|
paginator = commands.Paginator()
|
||||||
|
paginator.add_line(f"Running command: {' '.join(args[3 if args[0] == 'sudo' else 0:])}")
|
||||||
|
paginator.add_line(empty=True)
|
||||||
|
try:
|
||||||
|
start = time.time_ns()
|
||||||
|
process = await asyncio.create_subprocess_exec(
|
||||||
|
args[0],
|
||||||
|
*args[1:],
|
||||||
|
stdout=asyncio.subprocess.PIPE,
|
||||||
|
stderr=asyncio.subprocess.PIPE,
|
||||||
|
)
|
||||||
|
await process.wait()
|
||||||
|
stdout, stderr = await process.communicate()
|
||||||
|
end = time.time_ns()
|
||||||
|
time_taken_in_ms = (end - start) / 1000000
|
||||||
|
if stdout:
|
||||||
|
for line in stdout.splitlines():
|
||||||
|
paginator.add_line(line.decode())
|
||||||
|
if stderr:
|
||||||
|
for line in stderr.splitlines():
|
||||||
|
paginator.add_line(line.decode())
|
||||||
|
paginator.add_line(empty=True)
|
||||||
|
paginator.add_line(f"Exit code: {process.returncode}")
|
||||||
|
paginator.add_line(f"Time taken: {time_taken_in_ms:,.1f}ms")
|
||||||
|
except Exception as e:
|
||||||
|
paginator.add_line(f"Error: {e}")
|
||||||
|
for page in paginator.pages:
|
||||||
|
await ctx.respond(page)
|
||||||
|
|
||||||
|
|
||||||
|
def setup(bot):
|
||||||
|
bot.add_cog(NetworkCog(bot))
|
333
cogs/screenshot.py
Normal file
333
cogs/screenshot.py
Normal file
|
@ -0,0 +1,333 @@
|
||||||
|
import datetime
|
||||||
|
import io
|
||||||
|
import logging
|
||||||
|
import shutil
|
||||||
|
import tempfile
|
||||||
|
import time
|
||||||
|
import zipfile
|
||||||
|
from urllib.parse import urlparse
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
import discord
|
||||||
|
from discord.ext import commands
|
||||||
|
import asyncio
|
||||||
|
import aiohttp
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from selenium import webdriver
|
||||||
|
from selenium.common.exceptions import WebDriverException
|
||||||
|
from selenium.webdriver.chrome.options import Options as ChromeOptions
|
||||||
|
from selenium.webdriver.chrome.service import Service as ChromeService
|
||||||
|
|
||||||
|
|
||||||
|
class ScreenshotCog(commands.Cog):
|
||||||
|
def __init__(self, bot: commands.Bot):
|
||||||
|
self.bot = bot
|
||||||
|
|
||||||
|
self.chrome_options = ChromeOptions()
|
||||||
|
self.chrome_options.add_argument("--headless")
|
||||||
|
self.chrome_options.add_argument("--disable-dev-shm-usage")
|
||||||
|
# self.chrome_options.add_argument("--disable-gpu")
|
||||||
|
self.chrome_options.add_argument("--disable-extensions")
|
||||||
|
self.chrome_options.add_argument("--incognito")
|
||||||
|
|
||||||
|
prefs = {
|
||||||
|
# "download.open_pdf_in_system_reader": False,
|
||||||
|
# "download.prompt_for_download": True,
|
||||||
|
# "download.default_directory": "/dev/null",
|
||||||
|
# "plugins.always_open_pdf_externally": False,
|
||||||
|
"download_restrictions": 3,
|
||||||
|
}
|
||||||
|
self.chrome_options.add_experimental_option(
|
||||||
|
"prefs", prefs
|
||||||
|
)
|
||||||
|
|
||||||
|
self.dir = Path(__file__).parent.parent / "chrome"
|
||||||
|
self.dir.mkdir(mode=0o775, exist_ok=True)
|
||||||
|
|
||||||
|
self.chrome_dir = self.dir / "chrome-headless-shell-linux64"
|
||||||
|
self.chrome_bin = self.chrome_dir / "chrome-headless-shell"
|
||||||
|
self.chromedriver_dir = self.dir / "chromedriver-linux64"
|
||||||
|
self.chromedriver_bin = self.chromedriver_dir / "chromedriver"
|
||||||
|
|
||||||
|
self.chrome_options.binary_location = str(self.chrome_bin.resolve())
|
||||||
|
|
||||||
|
self.log = logging.getLogger("jimmy.cogs.screenshot")
|
||||||
|
|
||||||
|
def clear_directories(self):
|
||||||
|
shutil.rmtree(self.chrome_dir, ignore_errors=True)
|
||||||
|
shutil.rmtree(self.chromedriver_dir, ignore_errors=True)
|
||||||
|
self.chrome_dir.mkdir(mode=0o775, exist_ok=True)
|
||||||
|
self.chromedriver_dir.mkdir(mode=0o775, exist_ok=True)
|
||||||
|
|
||||||
|
async def download_latest_chrome(self, current: str = None, *, channel: str = "Stable"):
|
||||||
|
async with aiohttp.ClientSession(raise_for_status=True) as session:
|
||||||
|
async with session.get(
|
||||||
|
"https://googlechromelabs.github.io/chrome-for-testing/last-known-good-versions-with-downloads.json",
|
||||||
|
) as response:
|
||||||
|
versions = await response.json()
|
||||||
|
self.log.debug("Got chrome versions: %r", versions)
|
||||||
|
downloads = versions["channels"][channel]["downloads"]
|
||||||
|
version = versions["channels"][channel]["version"]
|
||||||
|
if version == current:
|
||||||
|
self.log.debug(f"Chrome is up to date ({versions} == {current})")
|
||||||
|
return
|
||||||
|
|
||||||
|
self.log.debug("Downloading chrome...")
|
||||||
|
chrome_zip_url = filter(lambda x: x["platform"] == "linux64", downloads["chrome-headless-shell"])
|
||||||
|
if not chrome_zip_url:
|
||||||
|
self.log.critical("No chrome zip url found for linux64 in %r.", downloads["chrome-headless-shell"])
|
||||||
|
raise RuntimeError("No chrome zip url found for linux64.")
|
||||||
|
chrome_zip_url = next(chrome_zip_url)["url"]
|
||||||
|
self.log.debug("Chrome zip url: %s", chrome_zip_url)
|
||||||
|
|
||||||
|
self.clear_directories()
|
||||||
|
|
||||||
|
chrome_target = (self.chrome_dir.parent / f"chrome-download-{version}.zip")
|
||||||
|
chromedriver_target = (self.chromedriver_dir.parent / f"chromedriver-download-{version}.zip")
|
||||||
|
if chrome_target.exists():
|
||||||
|
chrome_target.unlink()
|
||||||
|
if chromedriver_target.exists():
|
||||||
|
chromedriver_target.unlink()
|
||||||
|
with chrome_target.open("wb+") as file:
|
||||||
|
async with session.get(chrome_zip_url) as response:
|
||||||
|
async for data in response.content.iter_any():
|
||||||
|
self.log.debug("Read %d bytes from chrome zip.", len(data))
|
||||||
|
file.write(data)
|
||||||
|
|
||||||
|
self.log.debug("Extracting chrome...")
|
||||||
|
with zipfile.ZipFile(chrome_target) as zip_file:
|
||||||
|
await asyncio.get_event_loop().run_in_executor(
|
||||||
|
None,
|
||||||
|
zip_file.extractall,
|
||||||
|
self.chrome_dir.parent
|
||||||
|
)
|
||||||
|
self.log.debug("Finished extracting chrome.")
|
||||||
|
|
||||||
|
self.log.debug("Downloading chromedriver...")
|
||||||
|
chromedriver_zip_url = filter(lambda x: x["platform"] == "linux64", downloads["chromedriver"])
|
||||||
|
if not chromedriver_zip_url:
|
||||||
|
self.log.critical("No chromedriver zip url found for linux64 in %r.", downloads["chromedriver"])
|
||||||
|
raise RuntimeError("No chromedriver zip url found for linux64.")
|
||||||
|
chromedriver_zip_url = next(chromedriver_zip_url)["url"]
|
||||||
|
self.log.debug("Chromedriver zip url: %s", chromedriver_zip_url)
|
||||||
|
|
||||||
|
with chromedriver_target.open("wb+") as file:
|
||||||
|
async with session.get(chromedriver_zip_url) as response:
|
||||||
|
async for data in response.content.iter_any():
|
||||||
|
self.log.debug("Read %d bytes from chromedriver zip.", len(data))
|
||||||
|
file.write(data)
|
||||||
|
|
||||||
|
self.log.debug("Extracting chromedriver...")
|
||||||
|
with zipfile.ZipFile(chromedriver_target) as zip_file:
|
||||||
|
await asyncio.get_event_loop().run_in_executor(
|
||||||
|
None,
|
||||||
|
zip_file.extractall,
|
||||||
|
self.chromedriver_dir.parent
|
||||||
|
)
|
||||||
|
self.log.debug("Finished extracting chromedriver.")
|
||||||
|
|
||||||
|
self.log.debug("Making binaries executable.")
|
||||||
|
await asyncio.get_event_loop().run_in_executor(
|
||||||
|
None,
|
||||||
|
self.chrome_bin.chmod,
|
||||||
|
0o775
|
||||||
|
)
|
||||||
|
await asyncio.get_event_loop().run_in_executor(
|
||||||
|
None,
|
||||||
|
self.chromedriver_bin.chmod,
|
||||||
|
0o775
|
||||||
|
)
|
||||||
|
self.log.debug("Finished making binaries executable.")
|
||||||
|
|
||||||
|
async def get_version(self, full: bool = False) -> str:
|
||||||
|
proc = await asyncio.create_subprocess_exec(
|
||||||
|
str(self.chromedriver_bin),
|
||||||
|
"--version",
|
||||||
|
stdout=asyncio.subprocess.PIPE,
|
||||||
|
stderr=asyncio.subprocess.PIPE
|
||||||
|
)
|
||||||
|
stdout, stderr = await proc.communicate()
|
||||||
|
if proc.returncode != 0:
|
||||||
|
raise RuntimeError(f"Error getting chromedriver version: {stderr.decode()}")
|
||||||
|
if full:
|
||||||
|
return stdout.decode().strip()
|
||||||
|
return stdout.decode().strip().split(" ")[1]
|
||||||
|
|
||||||
|
async def is_up_to_date(self, channel: str = "Stable"):
|
||||||
|
try:
|
||||||
|
current = await self.get_version(False)
|
||||||
|
except (RuntimeError, FileNotFoundError):
|
||||||
|
return False
|
||||||
|
async with aiohttp.ClientSession(raise_for_status=True) as session:
|
||||||
|
async with session.get(
|
||||||
|
"https://googlechromelabs.github.io/chrome-for-testing/last-known-good-versions-with-downloads.json",
|
||||||
|
) as response:
|
||||||
|
versions = await response.json()
|
||||||
|
self.log.debug("Got chrome versions: %r", versions)
|
||||||
|
version = versions["channels"][channel]["version"]
|
||||||
|
if version == current:
|
||||||
|
self.log.debug(f"Chrome is up to date ({versions} == {current})")
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
@commands.command(name="update-chrome")
|
||||||
|
async def update_chrome(self, ctx: commands.Context, channel: str = "Stable"):
|
||||||
|
channel = channel.title()
|
||||||
|
if await self.is_up_to_date(channel):
|
||||||
|
await ctx.reply("Chrome is already up to date. Updating anyway.")
|
||||||
|
async with ctx.channel.typing():
|
||||||
|
try:
|
||||||
|
await self.download_latest_chrome(channel)
|
||||||
|
except RuntimeError as e:
|
||||||
|
return await ctx.reply(f"\N{cross mark} Error downloading chrome: {e}")
|
||||||
|
|
||||||
|
chrome_okay = self.chrome_bin.exists() and self.chrome_bin.is_file()
|
||||||
|
chromedriver_okay = self.chromedriver_bin.exists() and self.chromedriver_bin.is_file()
|
||||||
|
|
||||||
|
try:
|
||||||
|
chromedriver_version = await self.get_version(True)
|
||||||
|
except RuntimeError as e:
|
||||||
|
chromedriver_version = str(e)
|
||||||
|
|
||||||
|
return await ctx.reply(
|
||||||
|
f"\N{white heavy check mark} Done.\n"
|
||||||
|
f"CHANNEL: {channel}\n"
|
||||||
|
f"CHROME OKAY: {chrome_okay}\n"
|
||||||
|
f"CHROMEDRIVER OKAY: {chromedriver_okay}\n"
|
||||||
|
f"CHROMEDRIVER VERSION: {chromedriver_version}"
|
||||||
|
)
|
||||||
|
|
||||||
|
def compress_png(self, input_file: io.BytesIO) -> io.BytesIO:
|
||||||
|
img = Image.open(input_file)
|
||||||
|
img = img.convert("RGB")
|
||||||
|
with tempfile.NamedTemporaryFile(suffix=".webp") as file:
|
||||||
|
quality = 100
|
||||||
|
while quality > 0:
|
||||||
|
if quality == 100:
|
||||||
|
quality_r = 99
|
||||||
|
else:
|
||||||
|
quality_r = quality
|
||||||
|
self.log.debug("Compressing image with quality %d%%", quality_r)
|
||||||
|
img.save(file.name, "webp", quality=quality_r)
|
||||||
|
file.seek(0)
|
||||||
|
value = io.BytesIO(file.read())
|
||||||
|
if len(value.getvalue()) <= 24 * 1024 * 1024:
|
||||||
|
self.log.debug("%d%% was sufficient.", quality_r)
|
||||||
|
break
|
||||||
|
quality -= 15
|
||||||
|
else:
|
||||||
|
raise RuntimeError("Couldn't compress image.")
|
||||||
|
return value
|
||||||
|
|
||||||
|
@commands.slash_command()
|
||||||
|
async def screenshot(
|
||||||
|
self,
|
||||||
|
ctx: discord.ApplicationContext,
|
||||||
|
url: str,
|
||||||
|
load_timeout: int = 10,
|
||||||
|
render_timeout: int = None,
|
||||||
|
eager: bool = None,
|
||||||
|
resolution: str = "1920x1080"
|
||||||
|
):
|
||||||
|
"""Screenshots a webpage."""
|
||||||
|
await ctx.defer()
|
||||||
|
|
||||||
|
if eager is None:
|
||||||
|
eager = render_timeout is None
|
||||||
|
if render_timeout is None:
|
||||||
|
render_timeout = 30 if eager else 10
|
||||||
|
if not url.startswith("http"):
|
||||||
|
url = "https://" + url
|
||||||
|
parsed = urlparse(url)
|
||||||
|
await ctx.respond("Initialising...")
|
||||||
|
|
||||||
|
if not all(map(lambda x: x.exists() and x.is_file(), (self.chrome_bin, self.chromedriver_bin))):
|
||||||
|
await ctx.edit(content="Chrome is not installed, downloading. This may take a minute.")
|
||||||
|
await self.download_latest_chrome()
|
||||||
|
await ctx.edit(content="Initialising...")
|
||||||
|
elif not await self.is_up_to_date():
|
||||||
|
await ctx.edit(content="Updating chrome. This may take a minute.")
|
||||||
|
await self.download_latest_chrome()
|
||||||
|
await ctx.edit(content="Initialising...")
|
||||||
|
|
||||||
|
start_init = time.time()
|
||||||
|
service = await asyncio.to_thread(ChromeService, str(self.chromedriver_bin))
|
||||||
|
driver: webdriver.Chrome = await asyncio.to_thread(
|
||||||
|
webdriver.Chrome,
|
||||||
|
service=service,
|
||||||
|
options=self.chrome_options
|
||||||
|
)
|
||||||
|
driver.set_page_load_timeout(load_timeout)
|
||||||
|
if resolution:
|
||||||
|
try:
|
||||||
|
width, height = map(int, resolution.split("x"))
|
||||||
|
driver.set_window_size(width, height)
|
||||||
|
if height > 4320 or width > 7680:
|
||||||
|
return await ctx.respond("Invalid resolution. Max resolution is 7680x4320 (8K).")
|
||||||
|
except ValueError:
|
||||||
|
return await ctx.respond("Invalid resolution. please provide width x height, e.g. 1920x1080")
|
||||||
|
if eager:
|
||||||
|
driver.implicitly_wait(render_timeout)
|
||||||
|
end_init = time.time()
|
||||||
|
|
||||||
|
await ctx.edit(content=("Loading webpage..." if not eager else "Loading & screenshotting webpage..."))
|
||||||
|
start_request = time.time()
|
||||||
|
await asyncio.to_thread(driver.get, url)
|
||||||
|
end_request = time.time()
|
||||||
|
|
||||||
|
if not eager:
|
||||||
|
now = discord.utils.utcnow()
|
||||||
|
expires = now + datetime.timedelta(seconds=render_timeout)
|
||||||
|
await ctx.edit(content=f"Rendering (expires {discord.utils.format_dt(expires, 'R')})...")
|
||||||
|
start_wait = time.time()
|
||||||
|
await asyncio.sleep(render_timeout)
|
||||||
|
end_wait = time.time()
|
||||||
|
else:
|
||||||
|
start_wait = end_wait = 1
|
||||||
|
|
||||||
|
await ctx.edit(content="Saving screenshot...")
|
||||||
|
start_save = time.time()
|
||||||
|
ss = await asyncio.to_thread(driver.get_screenshot_as_png)
|
||||||
|
file = io.BytesIO()
|
||||||
|
await asyncio.to_thread(file.write, ss)
|
||||||
|
file.seek(0)
|
||||||
|
end_save = time.time()
|
||||||
|
|
||||||
|
if len(await asyncio.to_thread(file.getvalue)) > 24 * 1024 * 1024:
|
||||||
|
start_compress = time.time()
|
||||||
|
file = await asyncio.to_thread(self.compress_png, file)
|
||||||
|
fn = "screenshot.webp"
|
||||||
|
end_compress = time.time()
|
||||||
|
else:
|
||||||
|
fn = "screenshot.png"
|
||||||
|
start_compress = end_compress = 1
|
||||||
|
|
||||||
|
await ctx.edit(content="Cleaning up...")
|
||||||
|
start_cleanup = time.time()
|
||||||
|
await asyncio.to_thread(driver.quit)
|
||||||
|
end_cleanup = time.time()
|
||||||
|
|
||||||
|
screenshot_size_mb = round(len(await asyncio.to_thread(file.getvalue)) / 1024 / 1024, 2)
|
||||||
|
|
||||||
|
def seconds(start: float, end: float) -> float:
|
||||||
|
return round(end - start, 2)
|
||||||
|
|
||||||
|
embed = discord.Embed(
|
||||||
|
title=f"Screenshot of {parsed.hostname}",
|
||||||
|
description=f"Init time: {seconds(start_init, end_init)}s\n"
|
||||||
|
f"Request time: {seconds(start_request, end_request)}s\n"
|
||||||
|
f"Wait time: {seconds(start_wait, end_wait)}s\n"
|
||||||
|
f"Save time: {seconds(start_save, end_save)}s\n"
|
||||||
|
f"Compress time: {seconds(start_compress, end_compress)}s\n"
|
||||||
|
f"Cleanup time: {seconds(start_cleanup, end_cleanup)}s\n"
|
||||||
|
f"Screenshot size: {screenshot_size_mb}MB\n",
|
||||||
|
colour=discord.Colour.dark_theme(),
|
||||||
|
timestamp=discord.utils.utcnow()
|
||||||
|
)
|
||||||
|
embed.set_image(url="attachment://" + fn)
|
||||||
|
return await ctx.edit(content=None, embed=embed, file=discord.File(file, filename=fn))
|
||||||
|
|
||||||
|
|
||||||
|
def setup(bot):
|
||||||
|
bot.add_cog(ScreenshotCog(bot))
|
335
cogs/ytdl.py
Normal file
335
cogs/ytdl.py
Normal file
|
@ -0,0 +1,335 @@
|
||||||
|
import asyncio
|
||||||
|
import functools
|
||||||
|
import logging
|
||||||
|
import textwrap
|
||||||
|
import typing
|
||||||
|
|
||||||
|
import discord
|
||||||
|
import yt_dlp
|
||||||
|
from urllib.parse import urlparse
|
||||||
|
from pathlib import Path
|
||||||
|
import tempfile
|
||||||
|
from discord.ext import commands
|
||||||
|
|
||||||
|
|
||||||
|
COOKIES_TXT = Path.cwd() / "cookies.txt"
|
||||||
|
|
||||||
|
|
||||||
|
class YTDLCog(commands.Cog):
|
||||||
|
def __init__(self, bot: commands.Bot) -> None:
|
||||||
|
self.bot = bot
|
||||||
|
self.log = logging.getLogger("jimmy.cogs.ytdl")
|
||||||
|
self.common_formats = {
|
||||||
|
"144p": "17", # mp4 (h264+aac) v
|
||||||
|
"240p": "133+139",
|
||||||
|
"360p": "18",
|
||||||
|
"480p": "135+139",
|
||||||
|
"720p": "22",
|
||||||
|
"1080p": "137+140",
|
||||||
|
"1440p": "248+251", # webm (vp9+opus) v
|
||||||
|
"2160p": "313+251",
|
||||||
|
"mp3": "ba[filesize<25M]",
|
||||||
|
"m4a": "ba[ext=m4a][filesize<25M]",
|
||||||
|
"opus": "ba[ext=webm][filesize<25M]",
|
||||||
|
"vorbis": "ba[ext=webm][filesize<25M]",
|
||||||
|
"ogg": "ba[ext=webm][filesize<25M]",
|
||||||
|
}
|
||||||
|
self.default_options = {
|
||||||
|
"noplaylist": True,
|
||||||
|
"nocheckcertificate": True,
|
||||||
|
"no_color": True,
|
||||||
|
"noprogress": True,
|
||||||
|
"logger": self.log,
|
||||||
|
"format": "((bv+ba/b)[vcodec!=h265][vcodec!=av01][filesize<15M]/b[filesize<=15M]/b)",
|
||||||
|
"outtmpl": f"%(title).50s.%(ext)s",
|
||||||
|
"format_sort": [
|
||||||
|
"vcodec:h264",
|
||||||
|
"acodec:aac",
|
||||||
|
"vcodec:vp9",
|
||||||
|
"acodec:opus",
|
||||||
|
"acodec:vorbis",
|
||||||
|
"vcodec:vp8",
|
||||||
|
"ext",
|
||||||
|
],
|
||||||
|
"merge_output_format": "webm/mp4/mov/m4a/oga/ogg/mp3/mka/mkv",
|
||||||
|
"source_address": "0.0.0.0",
|
||||||
|
"concurrent_fragment_downloads": 4,
|
||||||
|
"max_filesize": (25 * 1024 * 1024) - 256
|
||||||
|
}
|
||||||
|
self.colours = {
|
||||||
|
"youtube.com": 0xff0000,
|
||||||
|
"tiktok.com": 0x25F5EF,
|
||||||
|
"instagram.com": 0xe1306c,
|
||||||
|
"shronk.net": 0xFFF952
|
||||||
|
}
|
||||||
|
|
||||||
|
@commands.slash_command(name="yt-dl")
|
||||||
|
@commands.max_concurrency(1, wait=False)
|
||||||
|
# @commands.bot_has_permissions(send_messages=True, embed_links=True, attach_files=True)
|
||||||
|
async def yt_dl_command(
|
||||||
|
self,
|
||||||
|
ctx: discord.ApplicationContext,
|
||||||
|
url: typing.Annotated[
|
||||||
|
str,
|
||||||
|
discord.Option(
|
||||||
|
str,
|
||||||
|
description="The URL to download from.",
|
||||||
|
required=True
|
||||||
|
)
|
||||||
|
],
|
||||||
|
user_format: typing.Annotated[
|
||||||
|
typing.Optional[str],
|
||||||
|
discord.Option(
|
||||||
|
str,
|
||||||
|
name="format",
|
||||||
|
description="The name of the format to download. Can also specify resolutions for youtube.",
|
||||||
|
required=False,
|
||||||
|
default=None
|
||||||
|
)
|
||||||
|
],
|
||||||
|
audio_only: typing.Annotated[
|
||||||
|
bool,
|
||||||
|
discord.Option(
|
||||||
|
bool,
|
||||||
|
name="audio-only",
|
||||||
|
description="Whether to convert result into an m4a file. Overwrites `format` if True.",
|
||||||
|
required=False,
|
||||||
|
default=False,
|
||||||
|
)
|
||||||
|
],
|
||||||
|
snip: typing.Annotated[
|
||||||
|
typing.Optional[str],
|
||||||
|
discord.Option(
|
||||||
|
str,
|
||||||
|
description="A start and end position to trim. e.g. 00:00:00-00:10:00.",
|
||||||
|
required=False
|
||||||
|
)
|
||||||
|
]
|
||||||
|
):
|
||||||
|
"""Runs yt-dlp and outputs into discord."""
|
||||||
|
await ctx.defer()
|
||||||
|
options = self.default_options.copy()
|
||||||
|
description = ""
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory(prefix="jimmy-ytdl-") as temp_dir:
|
||||||
|
temp_dir = Path(temp_dir)
|
||||||
|
paths = {
|
||||||
|
target: str(temp_dir)
|
||||||
|
for target in (
|
||||||
|
"home",
|
||||||
|
"temp",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
chosen_format = self.default_options["format"]
|
||||||
|
if user_format:
|
||||||
|
if user_format in self.common_formats:
|
||||||
|
chosen_format = self.common_formats[user_format]
|
||||||
|
else:
|
||||||
|
chosen_format = user_format
|
||||||
|
|
||||||
|
if audio_only:
|
||||||
|
# Overwrite format here to be best audio under 25 megabytes.
|
||||||
|
chosen_format = "ba[filesize<20M]"
|
||||||
|
# Also force sorting by the best audio bitrate first.
|
||||||
|
options["format_sort"] = [
|
||||||
|
"abr",
|
||||||
|
"br"
|
||||||
|
]
|
||||||
|
options["postprocessors"] = [
|
||||||
|
{"key": "FFmpegExtractAudio", "preferredquality": "96", "preferredcodec": "best"}
|
||||||
|
]
|
||||||
|
options["format"] = chosen_format
|
||||||
|
options["paths"] = paths
|
||||||
|
|
||||||
|
with yt_dlp.YoutubeDL(options) as downloader:
|
||||||
|
await ctx.respond(
|
||||||
|
embed=discord.Embed().set_footer(text="Downloading (step 1/10)")
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
# noinspection PyTypeChecker
|
||||||
|
extracted_info = await asyncio.to_thread(downloader.extract_info, url, download=False)
|
||||||
|
except yt_dlp.utils.DownloadError as e:
|
||||||
|
title = "error"
|
||||||
|
description = str(e)
|
||||||
|
thumbnail_url = webpage_url = None
|
||||||
|
else:
|
||||||
|
title = extracted_info.get("title", url)
|
||||||
|
title = textwrap.shorten(title, 100)
|
||||||
|
thumbnail_url = extracted_info.get("thumbnail") or None
|
||||||
|
webpage_url = extracted_info.get("webpage_url") or None
|
||||||
|
|
||||||
|
chosen_format = extracted_info.get("format")
|
||||||
|
chosen_format_id = extracted_info.get("format_id")
|
||||||
|
final_extension = extracted_info.get("ext")
|
||||||
|
format_note = extracted_info.get("format_note", "%s (%s)" % (chosen_format, chosen_format_id))
|
||||||
|
resolution = extracted_info.get("resolution")
|
||||||
|
fps = extracted_info.get("fps")
|
||||||
|
vcodec = extracted_info.get("vcodec")
|
||||||
|
acodec = extracted_info.get("acodec")
|
||||||
|
filesize = extracted_info.get("filesize", extracted_info.get("filesize_approx", 1))
|
||||||
|
|
||||||
|
lines = []
|
||||||
|
if chosen_format and chosen_format_id:
|
||||||
|
lines.append(
|
||||||
|
"* Chosen format: `%s` (`%s`)" % (chosen_format, chosen_format_id),
|
||||||
|
)
|
||||||
|
if format_note:
|
||||||
|
lines.append("* Format note: %r" % format_note)
|
||||||
|
if final_extension:
|
||||||
|
lines.append("* File extension: " + final_extension)
|
||||||
|
if resolution:
|
||||||
|
_s = resolution
|
||||||
|
if fps:
|
||||||
|
_s += " @ %s FPS" % fps
|
||||||
|
lines.append("* Resolution: " + _s)
|
||||||
|
if vcodec or acodec:
|
||||||
|
lines.append("%s+%s" % (vcodec or "N/A", acodec or "N/A"))
|
||||||
|
if filesize:
|
||||||
|
lines.append("* Filesize: %s" % yt_dlp.utils.format_bytes(filesize))
|
||||||
|
|
||||||
|
if lines:
|
||||||
|
description += "\n"
|
||||||
|
description += "\n".join(lines)
|
||||||
|
|
||||||
|
domain = urlparse(webpage_url).netloc
|
||||||
|
await ctx.edit(
|
||||||
|
embed=discord.Embed(
|
||||||
|
title=title,
|
||||||
|
description=description,
|
||||||
|
url=webpage_url,
|
||||||
|
colour=self.colours.get(domain, discord.Colour.og_blurple())
|
||||||
|
).set_footer(text="Downloading (step 2/10)").set_thumbnail(url=thumbnail_url)
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
await asyncio.to_thread(functools.partial(downloader.download, [url]))
|
||||||
|
except yt_dlp.DownloadError as e:
|
||||||
|
logging.error(e, exc_info=True)
|
||||||
|
return await ctx.edit(
|
||||||
|
embed=discord.Embed(
|
||||||
|
title="Error",
|
||||||
|
description=f"Download failed:\n```\n{e}\n```",
|
||||||
|
colour=discord.Colour.red(),
|
||||||
|
url=webpage_url,
|
||||||
|
),
|
||||||
|
delete_after=120,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
file = next(temp_dir.glob("*." + extracted_info["ext"]))
|
||||||
|
except StopIteration:
|
||||||
|
return await ctx.edit(
|
||||||
|
embed=discord.Embed(
|
||||||
|
title="Error",
|
||||||
|
description="Failed to locate downloaded video file.\n"
|
||||||
|
f"Files: {', '.join(list(map(str, temp_dir.iterdir())))}",
|
||||||
|
colour=discord.Colour.red(),
|
||||||
|
url=webpage_url
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
if snip:
|
||||||
|
try:
|
||||||
|
trim_start, trim_end = snip.split("-")
|
||||||
|
except ValueError:
|
||||||
|
trim_start, trim_end = snip, None
|
||||||
|
trim_start = trim_start or "00:00:00"
|
||||||
|
trim_end = trim_end or extracted_info.get("duration_string", "00:30:00")
|
||||||
|
new_file = temp_dir / ("output." + file.suffix)
|
||||||
|
args = [
|
||||||
|
"-hwaccel",
|
||||||
|
"auto",
|
||||||
|
"-ss",
|
||||||
|
trim_start,
|
||||||
|
"-i",
|
||||||
|
str(file),
|
||||||
|
"-to",
|
||||||
|
trim_end,
|
||||||
|
"-preset",
|
||||||
|
"faster",
|
||||||
|
"-crf",
|
||||||
|
"28",
|
||||||
|
"-deadline",
|
||||||
|
"realtime",
|
||||||
|
"-cpu-used",
|
||||||
|
"5",
|
||||||
|
"-movflags",
|
||||||
|
"faststart",
|
||||||
|
"-b:a",
|
||||||
|
"48k",
|
||||||
|
"-y",
|
||||||
|
"-strict",
|
||||||
|
"2",
|
||||||
|
str(new_file)
|
||||||
|
]
|
||||||
|
async with ctx.channel.typing():
|
||||||
|
await ctx.edit(
|
||||||
|
embed=discord.Embed(
|
||||||
|
title=f"Trimming from {trim_start} to {trim_end}.",
|
||||||
|
description="Please wait, this may take a couple of minutes.",
|
||||||
|
colour=discord.Colour.og_blurple(),
|
||||||
|
timestamp=discord.utils.utcnow()
|
||||||
|
)
|
||||||
|
)
|
||||||
|
process = await asyncio.create_subprocess_exec(
|
||||||
|
"ffmpeg",
|
||||||
|
*args,
|
||||||
|
stdout=asyncio.subprocess.PIPE,
|
||||||
|
stderr=asyncio.subprocess.PIPE
|
||||||
|
)
|
||||||
|
stdout, stderr = await process.communicate()
|
||||||
|
if process.returncode != 0:
|
||||||
|
return await ctx.edit(
|
||||||
|
embed=discord.Embed(
|
||||||
|
title="Error",
|
||||||
|
description=f"Trimming failed:\n```\n{stderr.decode()}\n```",
|
||||||
|
colour=discord.Colour.red(),
|
||||||
|
url=webpage_url
|
||||||
|
)
|
||||||
|
)
|
||||||
|
file = new_file
|
||||||
|
|
||||||
|
stat = file.stat()
|
||||||
|
size_bytes = stat.st_size
|
||||||
|
if size_bytes >= ((25 * 1024 * 1024) - 256):
|
||||||
|
return await ctx.edit(
|
||||||
|
embed=discord.Embed(
|
||||||
|
title="Error",
|
||||||
|
description=f"File is too large to upload ({round(size_bytes / 1024 / 1024)}MB).",
|
||||||
|
colour=discord.Colour.red(),
|
||||||
|
url=webpage_url
|
||||||
|
)
|
||||||
|
)
|
||||||
|
size_megabits = (size_bytes * 8) / 1024 / 1024
|
||||||
|
eta_seconds = size_megabits / 20
|
||||||
|
upload_file = await asyncio.to_thread(discord.File, file, filename=file.name)
|
||||||
|
await ctx.edit(
|
||||||
|
embed=discord.Embed(
|
||||||
|
title="Uploading...",
|
||||||
|
description=f"ETA <t:{int(eta_seconds + discord.utils.utcnow().timestamp()) + 2}:R>",
|
||||||
|
colour=discord.Colour.og_blurple(),
|
||||||
|
timestamp=discord.utils.utcnow()
|
||||||
|
)
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
await ctx.edit(
|
||||||
|
file=upload_file,
|
||||||
|
embed=discord.Embed(
|
||||||
|
title=f"Downloaded {title}!",
|
||||||
|
colour=discord.Colour.green(),
|
||||||
|
timestamp=discord.utils.utcnow(),
|
||||||
|
url=webpage_url
|
||||||
|
)
|
||||||
|
)
|
||||||
|
except discord.HTTPException as e:
|
||||||
|
self.log.error(e, exc_info=True)
|
||||||
|
return await ctx.edit(
|
||||||
|
embed=discord.Embed(
|
||||||
|
title="Error",
|
||||||
|
description=f"Upload failed:\n```\n{e}\n```",
|
||||||
|
colour=discord.Colour.red(),
|
||||||
|
url=webpage_url
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
def setup(bot):
|
||||||
|
bot.add_cog(YTDLCog(bot))
|
2
config.example.toml
Normal file
2
config.example.toml
Normal file
|
@ -0,0 +1,2 @@
|
||||||
|
[jimmy]
|
||||||
|
token = "foo"
|
91
main.py
Normal file
91
main.py
Normal file
|
@ -0,0 +1,91 @@
|
||||||
|
import datetime
|
||||||
|
import traceback
|
||||||
|
|
||||||
|
import toml
|
||||||
|
import discord
|
||||||
|
import logging
|
||||||
|
from pathlib import Path
|
||||||
|
from rich.logging import RichHandler
|
||||||
|
from logging import FileHandler
|
||||||
|
from discord.ext import commands
|
||||||
|
|
||||||
|
|
||||||
|
log = logging.getLogger("jimmy")
|
||||||
|
|
||||||
|
try:
|
||||||
|
CONFIG = toml.load('config.toml')
|
||||||
|
except FileNotFoundError:
|
||||||
|
log.critical("Unable to locate config.toml.", exc_info=True)
|
||||||
|
raise
|
||||||
|
|
||||||
|
logging.basicConfig(
|
||||||
|
format="%(asctime)s %(levelname)s %(name)s %(message)s",
|
||||||
|
datefmt="%Y-%m-%d %H:%M:%S",
|
||||||
|
level=CONFIG.get("logging", {}).get("level", "INFO"),
|
||||||
|
handlers=[
|
||||||
|
RichHandler(
|
||||||
|
level=CONFIG.get("logging", {}).get("level", "INFO"),
|
||||||
|
show_time=False,
|
||||||
|
show_path=False,
|
||||||
|
markup=True
|
||||||
|
),
|
||||||
|
FileHandler(
|
||||||
|
filename=CONFIG.get("logging", {}).get("file", "jimmy.log"),
|
||||||
|
mode="a",
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
bot = commands.Bot(
|
||||||
|
command_prefix=commands.when_mentioned_or("h!", "H!"),
|
||||||
|
case_insensitive=True,
|
||||||
|
strip_after_prefix=True,
|
||||||
|
debug_guilds=CONFIG["jimmy"].get("debug_guilds")
|
||||||
|
)
|
||||||
|
|
||||||
|
bot.load_extension("cogs.ytdl")
|
||||||
|
bot.load_extension("cogs.net")
|
||||||
|
bot.load_extension("cogs.screenshot")
|
||||||
|
|
||||||
|
|
||||||
|
@bot.event
|
||||||
|
async def on_ready():
|
||||||
|
log.info(f"Logged in as {bot.user} ({bot.user.id})")
|
||||||
|
|
||||||
|
|
||||||
|
@bot.listen()
|
||||||
|
async def on_application_command(ctx: discord.ApplicationContext):
|
||||||
|
log.info(f"Received command [b]{ctx.command}[/] from {ctx.author} in {ctx.guild}")
|
||||||
|
ctx.start_time = discord.utils.utcnow()
|
||||||
|
|
||||||
|
|
||||||
|
@bot.listen()
|
||||||
|
async def on_application_command_error(ctx: discord.ApplicationContext, exc: Exception):
|
||||||
|
log.error(f"Error in {ctx.command} from {ctx.author} in {ctx.guild}", exc_info=exc)
|
||||||
|
if isinstance(exc, commands.CommandOnCooldown):
|
||||||
|
expires = discord.utils.utcnow() + datetime.timedelta(seconds=exc.retry_after)
|
||||||
|
await ctx.respond(f"Command on cooldown. Try again {discord.utils.format_dt(expires, style='R')}.")
|
||||||
|
elif isinstance(exc, commands.MaxConcurrencyReached):
|
||||||
|
await ctx.respond("You've reached the maximum number of concurrent uses for this command.")
|
||||||
|
else:
|
||||||
|
if await bot.is_owner(ctx.author):
|
||||||
|
paginator = commands.Paginator(prefix="```py")
|
||||||
|
for line in traceback.format_exception(type(exc), exc, exc.__traceback__):
|
||||||
|
paginator.add_line(line[:1990])
|
||||||
|
for page in paginator.pages:
|
||||||
|
await ctx.respond(page)
|
||||||
|
else:
|
||||||
|
await ctx.respond(f"An error occurred while processing your command. Please try again later.\n"
|
||||||
|
f"{exc}")
|
||||||
|
|
||||||
|
|
||||||
|
@bot.listen()
|
||||||
|
async def on_application_command_completion(ctx: discord.ApplicationContext):
|
||||||
|
time_taken = discord.utils.utcnow() - ctx.start_time
|
||||||
|
log.info(
|
||||||
|
f"Completed command [b]{ctx.command}[/] from {ctx.author} in "
|
||||||
|
f"{ctx.guild} in {time_taken.total_seconds():.2f} seconds."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
bot.run(CONFIG["jimmy"]["token"])
|
9
requirements.txt
Normal file
9
requirements.txt
Normal file
|
@ -0,0 +1,9 @@
|
||||||
|
wheel>=0.42
|
||||||
|
setuptools>=69
|
||||||
|
yt-dlp @ https://github.com/yt-dlp/yt-dlp/archive/master.tar.gz
|
||||||
|
py-cord==2.4.1
|
||||||
|
httpx==0.26
|
||||||
|
psycopg==3.1.16
|
||||||
|
toml==0.10.2
|
||||||
|
pillow==10.2
|
||||||
|
selenium==4.16
|
Loading…
Reference in a new issue