college-bot-v1/cogs/other.py

807 lines
31 KiB
Python
Raw Normal View History

2022-12-01 13:34:26 +00:00
import asyncio
2023-06-03 14:43:31 +01:00
import glob
2023-12-05 21:12:33 +00:00
import hashlib
2022-12-01 12:22:11 +00:00
import io
2023-04-27 11:13:25 +01:00
import json
2023-12-05 21:12:33 +00:00
import logging
2022-12-01 12:22:11 +00:00
import os
2023-12-05 21:12:33 +00:00
import pathlib
2023-01-03 13:56:23 +00:00
import random
import re
2023-11-04 17:18:39 +00:00
import shutil
2023-03-14 12:39:57 +00:00
import tempfile
2023-01-03 15:17:09 +00:00
import textwrap
2023-11-29 11:53:42 +00:00
import typing
2023-04-29 02:25:19 +01:00
from functools import partial
2023-01-03 13:56:23 +00:00
from pathlib import Path
2024-04-14 19:44:30 +01:00
from time import time
from typing import Dict, Literal, Tuple
2023-01-03 14:43:49 +00:00
from urllib.parse import urlparse
2022-11-13 23:16:47 +00:00
import aiohttp
2024-04-14 19:44:30 +01:00
import config
2023-01-03 13:56:23 +00:00
import discord
2023-11-04 17:18:39 +00:00
import httpx
2023-12-05 21:12:33 +00:00
import openai
2023-01-03 13:56:23 +00:00
import psutil
2023-12-05 21:12:33 +00:00
import pydub
2023-11-04 17:18:39 +00:00
import pytesseract
import pyttsx3
from PIL import Image
2024-04-14 19:44:30 +01:00
from discord.ext import commands
2023-01-03 13:56:23 +00:00
from selenium import webdriver
from selenium.webdriver.chrome.options import Options as ChromeOptions
from selenium.webdriver.chrome.service import Service as ChromeService
from selenium.webdriver.firefox.options import Options as FirefoxOptions
from selenium.webdriver.firefox.service import Service as FirefoxService
2023-11-04 17:18:39 +00:00
2024-04-14 19:44:30 +01:00
from utils import Timer
2022-11-13 23:16:47 +00:00
2023-04-28 21:31:00 +01:00
try:
from config import proxy
except ImportError:
proxy = None
try:
from config import proxies
except ImportError:
if proxy:
proxies = [proxy] * 2
else:
proxies = []
2023-08-16 01:25:59 +01:00
try:
_engine = pyttsx3.init()
# noinspection PyTypeChecker
VOICES = [x.id for x in _engine.getProperty("voices")]
del _engine
except Exception as _pyttsx3_err:
2023-12-05 17:38:39 +00:00
logging.error("Failed to load pyttsx3: %r", _pyttsx3_err, exc_info=True)
2023-08-16 01:25:59 +01:00
pyttsx3 = None
VOICES = []
2023-03-20 14:48:23 +00:00
2022-11-13 23:16:47 +00:00
2023-12-05 21:12:33 +00:00
async def ollama_stream_reader(response: httpx.Response) -> typing.AsyncGenerator[dict[str, str | int | bool], None]:
async for chunk in response.aiter_lines():
# Each line is a JSON string
2023-11-10 23:36:18 +00:00
try:
loaded = json.loads(chunk)
2023-11-10 23:43:32 +00:00
yield loaded
2023-11-10 23:40:27 +00:00
except json.JSONDecodeError as e:
2023-12-05 17:38:39 +00:00
logging.warning("Failed to decode chunk %r: %r", chunk, e)
2023-11-10 23:36:18 +00:00
pass
2023-11-10 22:30:15 +00:00
2023-04-28 21:31:00 +01:00
def format_autocomplete(ctx: discord.AutocompleteContext):
url = ctx.options.get("url", os.urandom(6).hex())
self: "OtherCog" = ctx.bot.cogs["OtherCog"] # type: ignore
if url in self._fmt_cache:
2023-04-29 02:59:17 +01:00
suitable = []
for _format_key in self._fmt_cache[url]:
_format = self._fmt_cache[url][_format_key]
_format_nice = _format["format"]
if ctx.value.lower() in _format_nice.lower():
suitable.append(_format_nice)
return suitable
2023-04-28 21:31:00 +01:00
try:
parsed = urlparse(url, allow_fragments=True)
except ValueError:
pass
else:
if parsed.scheme in ("http", "https") and parsed.netloc:
self._fmt_queue.put_nowait(url)
2023-04-29 02:27:18 +01:00
return []
2023-04-28 21:31:00 +01:00
# noinspection DuplicatedCode
2022-11-13 23:16:47 +00:00
class OtherCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
self.lock = asyncio.Lock()
2023-12-05 21:48:46 +00:00
self.transcribe_lock = asyncio.Lock()
2023-03-20 14:39:22 +00:00
self.http = httpx.AsyncClient()
2023-04-28 21:31:00 +01:00
self._fmt_cache = {}
self._fmt_queue = asyncio.Queue()
self._worker_task = self.bot.loop.create_task(self.cache_population_job())
2023-11-11 18:31:48 +00:00
self.ollama_locks: dict[discord.Message, asyncio.Event] = {}
2023-11-13 20:20:55 +00:00
self.context_cache: dict[str, list[int]] = {}
2023-12-05 17:38:39 +00:00
self.log = logging.getLogger("jimmy.cogs.other")
2023-11-11 18:31:48 +00:00
2023-04-28 21:31:00 +01:00
def cog_unload(self):
self._worker_task.cancel()
async def cache_population_job(self):
while True:
url = await self._fmt_queue.get()
if url not in self._fmt_cache:
await self.list_formats(url, use_proxy=1)
self._fmt_queue.task_done()
async def list_formats(self, url: str, *, use_proxy: int = 0) -> dict:
if url in self._fmt_cache:
return self._fmt_cache[url]
2023-04-29 02:59:17 +01:00
import yt_dlp
class NullLogger:
def debug(self, *args, **kwargs):
pass
def info(self, *args, **kwargs):
pass
def warning(self, *args, **kwargs):
pass
def error(self, *args, **kwargs):
pass
with tempfile.TemporaryDirectory(prefix="jimmy-ytdl", suffix="-info") as tempdir:
with yt_dlp.YoutubeDL(
2023-11-04 17:18:39 +00:00
{
"windowsfilenames": True,
"restrictfilenames": True,
"noplaylist": True,
"nocheckcertificate": True,
"no_color": True,
"noprogress": True,
"logger": NullLogger(),
"paths": {"home": tempdir, "temp": tempdir},
"cookiefile": Path(__file__).parent.parent / "jimmy-cookies.txt",
}
2023-04-29 02:59:17 +01:00
) as downloader:
try:
info = await self.bot.loop.run_in_executor(
2023-11-04 17:18:39 +00:00
None, partial(downloader.extract_info, url, download=False)
2023-04-29 02:59:17 +01:00
)
except yt_dlp.utils.DownloadError:
return {}
info = downloader.sanitize_info(info)
new = {
fmt["format_id"]: {
"id": fmt["format_id"],
"ext": fmt["ext"],
"protocol": fmt["protocol"],
2023-06-01 01:09:05 +01:00
"acodec": fmt.get("acodec", "?"),
"vcodec": fmt.get("vcodec", "?"),
"resolution": fmt.get("resolution", "?x?"),
2023-11-04 17:18:39 +00:00
"filesize": fmt.get("filesize", float("inf")),
"format": fmt.get("format", "?"),
2023-04-29 02:59:17 +01:00
}
for fmt in info["formats"]
}
2023-04-28 21:31:00 +01:00
self._fmt_cache[url] = new
return new
2022-11-14 17:20:31 +00:00
2023-01-03 15:20:50 +00:00
async def screenshot_website(
2023-01-15 19:39:07 +00:00
self,
ctx: discord.ApplicationContext,
website: str,
driver: Literal["chrome", "firefox"],
render_time: int = 10,
2023-02-02 12:16:02 +00:00
load_timeout: int = 30,
2023-11-29 15:48:09 +00:00
window_height: int = 2560,
window_width: int = 1440,
2023-01-15 19:39:07 +00:00
full_screenshot: bool = False,
2023-01-16 15:51:03 +00:00
) -> Tuple[discord.File, str, int, int]:
2023-01-15 19:39:07 +00:00
async def _blocking(*args):
return await self.bot.loop.run_in_executor(None, *args)
def find_driver():
2023-01-16 09:55:45 +00:00
nonlocal driver, driver_path
2023-01-15 19:39:07 +00:00
drivers = {
"firefox": [
"/usr/bin/firefox-esr",
"/usr/bin/firefox",
],
2023-12-24 00:15:36 +00:00
"chrome": ["/usr/bin/chromium", "/usr/bin/google-chrome"],
2023-01-15 19:39:07 +00:00
}
selected_driver = driver
arr = drivers.pop(selected_driver)
for binary in arr:
b = Path(binary).resolve()
if not b.exists():
continue
driver = selected_driver
driver_path = b
break
else:
for key, value in drivers.items():
for binary in value:
b = Path(binary).resolve()
if not b.exists():
continue
driver = key
driver_path = b
break
else:
2023-01-13 22:42:03 +00:00
continue
break
else:
2023-01-15 19:39:07 +00:00
raise RuntimeError("No browser binary.")
return driver, driver_path
driver, driver_path = find_driver()
2023-12-05 17:38:39 +00:00
self.log.info(
2023-01-15 19:39:07 +00:00
"Using driver '{}' with binary '{}' to screenshot '{}', as requested by {}.".format(
driver, driver_path, website, ctx.user
)
)
def _setup():
nonlocal driver
if driver == "chrome":
options = ChromeOptions()
options.add_argument("--headless")
options.add_argument("--disable-dev-shm-usage")
options.add_argument("--disable-gpu")
options.add_argument("--disable-extensions")
options.add_argument("--incognito")
options.binary_location = str(driver_path)
service = ChromeService("/usr/bin/chromedriver")
driver = webdriver.Chrome(service=service, options=options)
driver.set_window_size(window_height, window_width)
2023-01-13 22:42:03 +00:00
else:
2023-01-15 19:39:07 +00:00
options = FirefoxOptions()
options.add_argument("--headless")
options.add_argument("--private-window")
options.add_argument("--safe-mode")
options.add_argument("--new-instance")
options.binary_location = str(driver_path)
service = FirefoxService("/usr/bin/geckodriver")
driver = webdriver.Firefox(service=service, options=options)
driver.set_window_size(window_height, window_width)
return driver, textwrap.shorten(website, 100)
# Is it overkill to cast this to a thread? yes
# Do I give a flying fuck? kinda
# Why am I doing this? I suspect setup is causing a ~10-second block of the event loop
2023-01-16 15:51:03 +00:00
driver_name = driver
start_init = time()
2023-01-15 19:39:07 +00:00
driver, friendly_url = await asyncio.to_thread(_setup)
2023-01-16 15:51:03 +00:00
end_init = time()
2023-12-05 17:38:39 +00:00
self.log.info("Driver '{}' initialised in {} seconds.".format(driver_name, round(end_init - start_init, 2)))
2023-01-03 13:56:23 +00:00
def _edit(content: str):
2023-01-09 14:36:32 +00:00
self.bot.loop.create_task(ctx.interaction.edit_original_response(content=content))
2023-02-02 12:16:02 +00:00
expires = round(time() + load_timeout)
_edit(content=f"Screenshotting <{friendly_url}>... (49%, loading webpage, aborts <t:{expires}:R>)")
await _blocking(driver.set_page_load_timeout, load_timeout)
2023-01-16 15:51:03 +00:00
start = time()
2023-01-15 19:39:07 +00:00
await _blocking(driver.get, website)
2023-01-16 15:51:03 +00:00
end = time()
get_time = round((end - start) * 1000)
render_time_expires = round(time() + render_time)
_edit(content=f"Screenshotting <{friendly_url}>... (66%, stopping render <t:{render_time_expires}:R>)")
2023-01-03 14:29:33 +00:00
await asyncio.sleep(render_time)
_edit(content=f"Screenshotting <{friendly_url}>... (83%, saving screenshot)")
2023-01-03 13:56:23 +00:00
domain = re.sub(r"https?://", "", website)
2023-01-15 19:39:07 +00:00
screenshot_method = driver.get_screenshot_as_png
2023-01-16 15:51:03 +00:00
if full_screenshot and driver_name == "firefox":
2023-01-15 19:39:07 +00:00
screenshot_method = driver.get_full_page_screenshot_as_png
2023-01-16 15:51:03 +00:00
start = time()
2023-01-15 19:39:07 +00:00
data = await _blocking(screenshot_method)
2023-01-03 13:56:23 +00:00
_io = io.BytesIO()
2023-11-29 15:17:21 +00:00
_io.write(data)
2023-01-03 13:56:23 +00:00
_io.seek(0)
2023-01-16 15:51:03 +00:00
end = time()
screenshot_time = round((end - start) * 1000)
2023-01-03 13:56:23 +00:00
driver.quit()
2023-01-16 15:51:03 +00:00
return discord.File(_io, f"{domain}.png"), driver_name, get_time, screenshot_time
2023-01-03 13:56:23 +00:00
2022-12-29 17:41:41 +00:00
@staticmethod
async def get_interface_ip_addresses() -> Dict[str, list[Dict[str, str | bool | int]]]:
addresses = await asyncio.to_thread(psutil.net_if_addrs)
stats = await asyncio.to_thread(psutil.net_if_stats)
result = {}
for key in addresses.keys():
result[key] = []
for ip_addr in addresses[key]:
if ip_addr.broadcast is None:
continue
else:
result[key].append(
{
"ip": ip_addr.address,
"netmask": ip_addr.netmask,
"broadcast": ip_addr.broadcast,
"up": stats[key].isup,
2023-01-03 15:20:50 +00:00
"speed": stats[key].speed,
2022-12-29 17:41:41 +00:00
}
)
return result
2022-11-14 17:20:31 +00:00
@staticmethod
async def get_xkcd(session: aiohttp.ClientSession, n: int) -> dict | None:
async with session.get("https://xkcd.com/{!s}/info.0.json".format(n)) as response:
if response.status == 200:
2022-11-13 23:16:47 +00:00
data = await response.json()
2022-11-14 17:20:31 +00:00
return data
@staticmethod
async def random_xkcd_number(session: aiohttp.ClientSession) -> int:
async with session.get("https://c.xkcd.com/random/comic") as response:
if response.status != 302:
number = random.randint(100, 999)
else:
number = int(response.headers["location"].split("/")[-2])
2022-11-14 17:20:31 +00:00
return number
@staticmethod
async def random_xkcd(session: aiohttp.ClientSession) -> dict | None:
"""Fetches a random XKCD.
Basically a shorthand for random_xkcd_number and get_xkcd.
"""
number = await OtherCog.random_xkcd_number(session)
return await OtherCog.get_xkcd(session, number)
@staticmethod
def get_xkcd_embed(data: dict) -> discord.Embed:
2022-11-13 23:16:47 +00:00
embed = discord.Embed(
title=data["safe_title"], description=data["alt"], color=discord.Colour.embed_background()
2022-11-13 23:16:47 +00:00
)
embed.set_footer(text="XKCD #{!s}".format(data["num"]))
embed.set_image(url=data["img"])
2022-11-14 17:20:31 +00:00
return embed
@staticmethod
async def generate_xkcd(n: int = None) -> discord.Embed:
async with aiohttp.ClientSession() as session:
if n is None:
data = await OtherCog.random_xkcd(session)
n = data["num"]
2022-11-14 17:20:31 +00:00
else:
data = await OtherCog.get_xkcd(session, n)
if data is None:
return discord.Embed(
title="Failed to load XKCD :(", description="Try again later.", color=discord.Colour.red()
2022-11-14 17:20:31 +00:00
).set_footer(text="Attempted to retrieve XKCD #{!s}".format(n))
return OtherCog.get_xkcd_embed(data)
class XKCDGalleryView(discord.ui.View):
def __init__(self, n: int):
super().__init__(timeout=300, disable_on_timeout=True)
self.n = n
2022-11-16 17:28:47 +00:00
def __rich_repr__(self):
yield "n", self.n
yield "message", self.message
@discord.ui.button(label="Previous", style=discord.ButtonStyle.blurple)
2022-11-14 17:20:31 +00:00
async def previous_comic(self, _, interaction: discord.Interaction):
self.n -= 1
await interaction.response.defer()
await interaction.edit_original_response(embed=await OtherCog.generate_xkcd(self.n))
@discord.ui.button(label="Random", style=discord.ButtonStyle.blurple)
2022-11-14 17:20:31 +00:00
async def random_comic(self, _, interaction: discord.Interaction):
await interaction.response.defer()
await interaction.edit_original_response(embed=await OtherCog.generate_xkcd())
self.n = random.randint(1, 999)
@discord.ui.button(label="Next", style=discord.ButtonStyle.blurple)
2022-11-14 17:20:31 +00:00
async def next_comic(self, _, interaction: discord.Interaction):
self.n += 1
await interaction.response.defer()
await interaction.edit_original_response(embed=await OtherCog.generate_xkcd(self.n))
@commands.slash_command()
async def xkcd(self, ctx: discord.ApplicationContext, *, number: int = None):
"""Shows an XKCD comic"""
embed = await self.generate_xkcd(number)
view = self.XKCDGalleryView(number)
return await ctx.respond(embed=embed, view=view)
2022-11-13 23:16:47 +00:00
2022-12-29 17:41:41 +00:00
@commands.command(name="kys", aliases=["kill"])
2022-12-28 21:14:14 +00:00
@commands.is_owner()
async def end_your_life(self, ctx: commands.Context):
await ctx.send(":( okay")
await self.bot.close()
2023-11-04 16:47:14 +00:00
@staticmethod
2023-11-29 10:35:35 +00:00
async def check_proxy(url: str = "socks5://localhost:1090", *, timeout: float = 3.0):
2024-02-21 01:13:39 +00:00
async with httpx.AsyncClient(http2=True, timeout=timeout) as client:
my_ip4 = (await client.get("https://api.ipify.org")).text
real_ips = [my_ip4]
2023-11-04 16:47:14 +00:00
# Check the proxy
2024-02-21 01:13:39 +00:00
async with httpx.AsyncClient(http2=True, proxies=url, timeout=timeout) as client:
2023-11-04 16:47:14 +00:00
try:
2024-02-21 01:13:39 +00:00
response = await client.get(
"https://1.1.1.1/cdn-cgi/trace",
2023-03-16 21:45:01 +00:00
)
2024-02-21 01:13:39 +00:00
response.raise_for_status()
for line in response.text.splitlines():
if line.startswith("ip"):
if any(x in line for x in real_ips):
return 1
except (httpx.TransportError, httpx.HTTPStatusError):
return 2
return 0
2023-11-04 17:18:39 +00:00
2023-03-27 23:16:28 +01:00
@commands.slash_command()
@commands.cooldown(5, 10, commands.BucketType.user)
@commands.max_concurrency(1, commands.BucketType.user)
async def quote(self, ctx: discord.ApplicationContext):
"""Generates a random quote"""
2023-11-04 17:18:39 +00:00
emoji = discord.PartialEmoji(name="loading", animated=True, id=1101463077586735174)
2023-04-28 21:31:00 +01:00
2023-04-28 15:20:15 +01:00
async def get_quote() -> str | discord.File:
try:
response = await self.http.get("https://inspirobot.me/api?generate=true")
except (ConnectionError, httpx.HTTPError, httpx.NetworkError) as e:
return "Failed to get quote. " + str(e)
if response.status_code != 200:
return f"Failed to get quote. Status code: {response.status_code}"
url = response.text
try:
response = await self.http.get(url)
except (ConnectionError, httpx.HTTPError, httpx.NetworkError) as e:
return url
else:
if response.status_code != 200:
return url
x = io.BytesIO(response.content)
x.seek(0)
return discord.File(x, filename="quote.jpg")
class GenerateNewView(discord.ui.View):
2023-04-28 15:39:32 +01:00
def __init__(self):
2023-11-04 17:18:39 +00:00
super().__init__(timeout=300, disable_on_timeout=True)
2023-04-28 15:39:32 +01:00
2023-04-28 15:20:15 +01:00
async def __aenter__(self):
self.disable_all_items()
if self.message:
await self.message.edit(view=self)
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
self.enable_all_items()
if self.message:
await self.message.edit(view=self)
return self
2023-04-28 15:24:28 +01:00
async def interaction_check(self, interaction: discord.Interaction) -> bool:
return interaction.user == ctx.user and interaction.channel == ctx.channel
2023-04-28 15:20:15 +01:00
@discord.ui.button(
label="New Quote",
style=discord.ButtonStyle.green,
2023-11-04 17:18:39 +00:00
emoji=discord.PartialEmoji.from_str("\U000023ed\U0000fe0f"),
2023-04-28 15:20:15 +01:00
)
async def new_quote(self, _, interaction: discord.Interaction):
2023-04-28 15:32:15 +01:00
await interaction.response.defer(invisible=True)
2023-04-28 15:20:15 +01:00
async with self:
2023-04-28 15:38:05 +01:00
followup = await interaction.followup.send(f"{emoji} Generating quote")
2023-04-28 15:20:15 +01:00
new_result = await get_quote()
if isinstance(new_result, discord.File):
2023-04-28 15:40:46 +01:00
return await followup.edit(content=None, file=new_result, view=GenerateNewView())
2023-04-28 15:20:15 +01:00
else:
2023-04-28 15:38:05 +01:00
return await followup.edit(content=new_result, view=GenerateNewView())
2023-04-28 15:20:15 +01:00
@discord.ui.button(
2023-11-04 17:18:39 +00:00
label="Regenerate", style=discord.ButtonStyle.blurple, emoji=discord.PartialEmoji.from_str("\U0001f504")
2023-04-28 15:20:15 +01:00
)
async def regenerate(self, _, interaction: discord.Interaction):
2023-04-28 15:32:15 +01:00
await interaction.response.defer(invisible=True)
2023-04-28 15:20:15 +01:00
async with self:
2023-04-28 15:24:28 +01:00
message = await interaction.original_response()
2023-04-28 15:38:05 +01:00
if "\U00002b50" in [_reaction.emoji for _reaction in message.reactions]:
2023-04-28 15:24:28 +01:00
return await interaction.followup.send(
"\N{cross mark} Message is starred and cannot be regenerated. You can press "
"'New Quote' to generate a new quote instead.",
2023-11-04 17:18:39 +00:00
ephemeral=True,
2023-04-28 15:24:28 +01:00
)
2023-04-28 15:20:15 +01:00
new_result = await get_quote()
if isinstance(new_result, discord.File):
return await interaction.edit_original_response(file=new_result)
else:
return await interaction.edit_original_response(content=new_result)
2023-11-04 17:18:39 +00:00
@discord.ui.button(label="Delete", style=discord.ButtonStyle.red, emoji="\N{wastebasket}\U0000fe0f")
2023-04-28 15:26:09 +01:00
async def delete(self, _, interaction: discord.Interaction):
2023-04-28 15:32:15 +01:00
await interaction.response.defer(invisible=True)
2023-04-28 15:26:09 +01:00
await interaction.delete_original_response()
self.stop()
2023-03-27 23:16:28 +01:00
await ctx.defer()
2023-04-28 15:20:15 +01:00
result = await get_quote()
if isinstance(result, discord.File):
return await ctx.respond(file=result, view=GenerateNewView())
2023-03-27 23:16:28 +01:00
else:
2023-04-28 15:20:15 +01:00
return await ctx.respond(result, view=GenerateNewView())
2023-03-14 12:39:57 +00:00
2024-04-14 19:44:30 +01:00
async def _ocr_core(self, attachment: discord.Attachment) -> tuple[dict[str, float], str]:
timings: dict[str, float] = {}
2023-05-30 19:53:47 +01:00
with Timer() as _t:
2023-05-05 20:08:47 +01:00
data = await attachment.read()
file = io.BytesIO(data)
file.seek(0)
2023-05-30 19:53:47 +01:00
timings["Download attachment"] = _t.total
with Timer() as _t:
2023-05-05 20:08:47 +01:00
img = await self.bot.loop.run_in_executor(None, Image.open, file)
2023-05-30 19:53:47 +01:00
timings["Parse image"] = _t.total
2023-05-05 10:35:17 +01:00
try:
2023-05-30 19:53:47 +01:00
with Timer() as _t:
2023-05-05 20:08:47 +01:00
text = await self.bot.loop.run_in_executor(None, pytesseract.image_to_string, img)
2023-05-30 19:53:47 +01:00
timings["Perform OCR"] = _t.total
2023-05-05 10:35:17 +01:00
except pytesseract.TesseractError as e:
2024-04-14 19:44:30 +01:00
raise RuntimeError(f"Failed to perform OCR: `{e}`")
2023-05-05 10:35:17 +01:00
2024-04-14 19:44:30 +01:00
if len(text) >= 1744:
2023-05-30 19:53:47 +01:00
with Timer() as _t:
2023-05-05 20:08:47 +01:00
try:
2024-04-14 19:44:30 +01:00
file.seek(0)
response = await self.http.post(
"https://paste.nexy7574.co.uk/upload",
data={
"expiration": "1week",
"burn_after": "0",
"syntax_highlight": "none",
"privacy": "unlisted",
},
files={
"file": (attachment.filename, file, attachment.content_type)
2023-11-04 17:18:39 +00:00
},
2024-04-14 19:44:30 +01:00
follow_redirects=False
2023-05-05 20:08:47 +01:00
)
response.raise_for_status()
2024-04-14 19:44:30 +01:00
except httpx.HTTPError as e:
raise RuntimeError(f"Failed to upload OCR content: `{e}`")
2023-05-05 20:08:47 +01:00
else:
2024-04-14 19:44:30 +01:00
text = "View on [paste.nexy7574.co.uk](%s)" % response.next_request.url
timings["Upload text to pastebin"] = _t.total
return timings, text
@commands.slash_command()
@commands.cooldown(1, 30, commands.BucketType.user)
@commands.max_concurrency(1, commands.BucketType.user)
async def ocr(
self,
ctx: discord.ApplicationContext,
attachment: discord.Option(
discord.SlashCommandOptionType.attachment,
description="Image to perform OCR on",
),
):
"""OCRs an image"""
await ctx.defer()
attachment: discord.Attachment
timings, text = await self._ocr_core(attachment)
embed = discord.Embed(
description=text,
colour=discord.Colour.blurple()
)
embed.set_image(url=attachment.url)
with Timer() as _t:
await ctx.respond(
embed=embed
)
timings["Respond (Text)"] = _t.total
2023-05-05 10:35:17 +01:00
if timings:
2023-05-30 19:56:39 +01:00
text = "Timings:\n" + "\n".join("{}: {:.2f}s".format(k.title(), v) for k, v in timings.items())
await ctx.edit(
2024-04-14 19:44:30 +01:00
content=text
)
2023-05-05 10:35:17 +01:00
2024-04-14 19:44:30 +01:00
@commands.message_command(name="Run OCR")
async def message_ocr(self, ctx: discord.ApplicationContext, message: discord.Message):
await ctx.defer()
attachment: discord.Attachment | None
for attachment in message.attachments:
if attachment.content_type.startswith("image/"):
break
else:
return await ctx.respond(":x: No images found in message.", delete_after=30)
timings, text = await self._ocr_core(attachment)
embed = discord.Embed(
title="OCR for " + attachment.filename,
description=text,
colour=discord.Colour.blurple(),
url=message.jump_url
)
embed.set_image(url=attachment.url)
await ctx.respond(
embed=embed
)
2023-05-25 09:51:16 +01:00
@commands.message_command(name="Convert Image to GIF")
async def convert_image_to_gif(self, ctx: discord.ApplicationContext, message: discord.Message):
2023-05-25 09:55:02 +01:00
await ctx.defer()
2023-05-25 09:51:16 +01:00
for attachment in message.attachments:
if attachment.content_type.startswith("image/"):
break
else:
return await ctx.respond("No image found.")
image = attachment
image: discord.Attachment
with tempfile.TemporaryFile("wb+") as f:
await image.save(f)
f.seek(0)
img = await self.bot.loop.run_in_executor(None, Image.open, f)
if img.format.upper() not in ("PNG", "JPEG", "WEBP", "HEIF", "BMP", "TIFF"):
return await ctx.respond("Image must be PNG, JPEG, WEBP, or HEIF.")
with tempfile.TemporaryFile("wb+") as f2:
caller = partial(img.save, f2, format="GIF")
await self.bot.loop.run_in_executor(None, caller)
f2.seek(0)
try:
await ctx.respond(file=discord.File(f2, filename="image.gif"))
except discord.HTTPException as e:
if e.code == 40005:
return await ctx.respond("Image is too large.")
return await ctx.respond(f"Failed to upload: `{e}`")
try:
2023-05-25 09:53:24 +01:00
f2.seek(0)
2023-05-25 09:51:16 +01:00
await ctx.user.send(file=discord.File(f2, filename="image.gif"))
except discord.Forbidden:
return await ctx.respond("Unable to mirror to your DM - am I blocked?", ephemeral=True)
2023-06-03 14:17:34 +01:00
@commands.slash_command()
@commands.cooldown(1, 180, commands.BucketType.user)
@commands.max_concurrency(1, commands.BucketType.user)
async def sherlock(
2023-11-04 17:18:39 +00:00
self, ctx: discord.ApplicationContext, username: str, search_nsfw: bool = False, use_tor: bool = False
2023-06-03 14:17:34 +01:00
):
"""Sherlocks a username."""
2023-06-28 22:19:40 +01:00
# git clone https://github.com/sherlock-project/sherlock.git && cd sherlock && docker build -t sherlock .
2023-06-03 14:43:31 +01:00
if re.search(r"\s", username) is not None:
return await ctx.respond("Username cannot contain spaces.")
2023-06-03 14:22:35 +01:00
async def background_task():
2023-06-03 14:46:36 +01:00
chars = ["|", "/", "-", "\\"]
n = 0
2023-06-03 14:22:35 +01:00
# Every 5 seconds update the embed to show that the command is still running
while True:
2023-06-03 14:57:35 +01:00
await asyncio.sleep(2.5)
2023-06-03 14:23:53 +01:00
elapsed = time() - start_time
2023-06-03 14:28:44 +01:00
embed = discord.Embed(
2023-06-03 14:47:49 +01:00
title="Sherlocking username %s" % chars[n % 4],
2023-06-03 14:28:44 +01:00
description=f"Elapsed: {elapsed:.0f}s",
2023-11-04 17:18:39 +00:00
colour=discord.Colour.dark_theme(),
2023-06-03 14:22:35 +01:00
)
2023-11-04 17:18:39 +00:00
await ctx.edit(embed=embed)
2023-06-03 14:46:36 +01:00
n += 1
2023-06-03 14:22:35 +01:00
2023-06-03 14:17:34 +01:00
await ctx.defer()
# output results to a temporary directory
2023-06-03 15:20:34 +01:00
tempdir = Path("./tmp/sherlock").resolve()
tempdir.mkdir(parents=True, exist_ok=True)
command = [
"docker",
"run",
"--rm",
"-t",
"-v",
f"{tempdir}:/opt/sherlock/results",
"sherlock",
2023-11-04 17:18:39 +00:00
"--folderoutput",
"/opt/sherlock/results",
2023-06-03 15:20:34 +01:00
"--print-found",
2023-11-04 17:18:39 +00:00
"--csv",
2023-06-03 15:20:34 +01:00
]
if search_nsfw:
command.append("--nsfw")
if use_tor:
command.append("--tor")
# Output to result.csv
# Username to search for
command.append(username)
# Run the command
start_time = time()
result = await asyncio.create_subprocess_exec(
*command,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
)
await ctx.respond(embed=discord.Embed(title="Starting..."))
task = asyncio.create_task(background_task())
# Wait for it to finish
stdout, stderr = await result.communicate()
await result.wait()
task.cancel()
# wait for task to exit
try:
await task
except asyncio.CancelledError:
pass
# If it errored, send the error
if result.returncode != 0:
shutil.rmtree(tempdir, ignore_errors=True)
return await ctx.edit(
2023-06-03 14:17:34 +01:00
embed=discord.Embed(
2023-06-03 15:20:34 +01:00
title="Error",
description=f"```ansi\n{stderr.decode()[:4000]}```",
colour=discord.Colour.red(),
)
2023-06-03 14:17:34 +01:00
)
2023-06-03 15:20:34 +01:00
# If it didn't error, send the results
stdout = stdout.decode()
if len(stdout) > 4000:
paginator = commands.Paginator("```ansi", max_size=4000)
for line in stdout.splitlines():
paginator.add_line(line)
desc = paginator.pages[0]
title = "Results (truncated)"
else:
desc = f"```ansi\n{stdout}```"
title = "Results"
files = list(map(discord.File, glob.glob(f"{tempdir}/*")))
await ctx.edit(
files=files,
embed=discord.Embed(
title=title,
description=desc,
colour=discord.Colour.green(),
),
)
shutil.rmtree(tempdir, ignore_errors=True)
2023-06-03 14:17:34 +01:00
2023-12-03 22:42:16 +00:00
@commands.message_command(name="Transcribe")
async def transcribe_message(self, ctx: discord.ApplicationContext, message: discord.Message):
await ctx.defer()
2023-12-05 21:48:46 +00:00
async with self.transcribe_lock:
if not message.attachments:
return await ctx.respond("No attachments found.")
_ft = "wav"
for attachment in message.attachments:
if attachment.content_type.startswith("audio/"):
_ft = attachment.filename.split(".")[-1]
break
else:
return await ctx.respond("No voice messages.")
if getattr(config, "OPENAI_KEY", None) is None:
return await ctx.respond("Service unavailable.")
file_hash = hashlib.sha1(usedforsecurity=False)
file_hash.update(await attachment.read())
file_hash = file_hash.hexdigest()
cache = Path.home() / ".cache" / "lcc-bot" / ("%s-transcript.txt" % file_hash)
cached = False
if not cache.exists():
client = openai.OpenAI(api_key=config.OPENAI_KEY)
with tempfile.NamedTemporaryFile("wb+", suffix=".mp4") as f:
with tempfile.NamedTemporaryFile("wb+", suffix="-" + attachment.filename) as f2:
await attachment.save(f2.name)
f2.seek(0)
seg: pydub.AudioSegment = await asyncio.to_thread(pydub.AudioSegment.from_file, file=f2, format=_ft)
seg = seg.set_channels(1)
await asyncio.to_thread(seg.export, f.name, format="mp4")
f.seek(0)
transcript = await asyncio.to_thread(
client.audio.transcriptions.create, file=pathlib.Path(f.name), model="whisper-1"
)
text = transcript.text
cache.write_text(text)
else:
text = cache.read_text()
cached = True
paginator = commands.Paginator("", "", 4096)
for line in text.splitlines():
paginator.add_line(textwrap.shorten(line, 4096))
embeds = list(map(lambda p: discord.Embed(description=p), paginator.pages))
await ctx.respond(embeds=embeds or [discord.Embed(description="No text found.")])
if await self.bot.is_owner(ctx.user):
await ctx.respond(
("Cached response ({})" if cached else "Uncached response ({})").format(file_hash), ephemeral=True
2023-12-03 22:42:16 +00:00
)
2023-01-03 14:43:49 +00:00
2022-11-13 23:16:47 +00:00
def setup(bot):
bot.add_cog(OtherCog(bot))