|
| 1 | +import aiohttp |
| 2 | +import mimetypes |
| 3 | +from maubot import Plugin, MessageEvent |
| 4 | +from maubot.handlers import command |
| 5 | +from mautrix.types import TextMessageEventContent, MessageType, Format |
| 6 | +from bs4 import BeautifulSoup |
| 7 | + |
| 8 | + |
| 9 | +class ImdbBot(Plugin): |
| 10 | + headers = { |
| 11 | + "Sec-GPC": "1", |
| 12 | + "accept-encoding": "gzip, deflate, br, zstd", |
| 13 | + "accept-language": "pl,en-US;q=0.7,en;q=0.3", |
| 14 | + "user-agent": "Mozilla/5.0 (X11; Linux x86_64; rv:138.0) Gecko/20100101 Firefox/138.0", |
| 15 | + "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", |
| 16 | + "referer": "https://duckduckgo.com/" |
| 17 | + } |
| 18 | + |
| 19 | + @command.new(name="imdb", help="IMDb Search") |
| 20 | + @command.argument("title", pass_raw=True, required=True) |
| 21 | + async def search(self, evt: MessageEvent, title: str) -> None: |
| 22 | + await evt.mark_read() |
| 23 | + title = title.strip() |
| 24 | + if not title: |
| 25 | + await evt.respond("Usage: !imdb <title>") |
| 26 | + return |
| 27 | + |
| 28 | + url = await self.web_search("site:imdb.com", title) |
| 29 | + if not url or url == "https://www.imdb.com/": |
| 30 | + await evt.reply(f"Failed to find results for *{title}*") |
| 31 | + return |
| 32 | + |
| 33 | + content = await self.prepare_message(url) |
| 34 | + if content: |
| 35 | + await evt.reply(content) |
| 36 | + else: |
| 37 | + await evt.reply("Something went wrong when I was preparing summary.") |
| 38 | + |
| 39 | + async def web_search(self, modifier: str, query: str) -> str: |
| 40 | + url = f"https://lite.duckduckgo.com/lite/?q=\\+{modifier}+{query}" |
| 41 | + try: |
| 42 | + timeout = aiohttp.ClientTimeout(total=20) |
| 43 | + response = await self.http.get(url, headers=ImdbBot.headers, timeout=timeout, allow_redirects=True, raise_for_status=True) |
| 44 | + result = str(response.url) |
| 45 | + if "duckduckgo.com" not in result: |
| 46 | + return result |
| 47 | + except aiohttp.ClientError as e: |
| 48 | + self.log.error(f"Web search: Connection failed: {e}") |
| 49 | + return "" |
| 50 | + |
| 51 | + async def prepare_message(self, url: str) -> TextMessageEventContent | None: |
| 52 | + timeout = aiohttp.ClientTimeout(total=20) |
| 53 | + try: |
| 54 | + response = await self.http.get(url, headers=ImdbBot.headers, timeout=timeout, raise_for_status=True) |
| 55 | + text = await response.text() |
| 56 | + except aiohttp.ClientError as e: |
| 57 | + self.log.error(f"Scraping IMDb: Connection failed: {e}") |
| 58 | + return None |
| 59 | + |
| 60 | + soup = BeautifulSoup(text, "html.parser") |
| 61 | + if soup.head is None: |
| 62 | + return None |
| 63 | + info = soup.head.find("meta", property="og:title") |
| 64 | + info = info["content"].split("⭐") if info else "" |
| 65 | + title = info[0].strip() if info else "-" |
| 66 | + info = info[1].split("|") if info else ["-", "-"] |
| 67 | + rating = info[0].strip() + "/10" |
| 68 | + tags = info[1].strip() |
| 69 | + video_type = soup.head.find("meta", property="og:type") |
| 70 | + video_type = video_type["content"] if video_type else "" |
| 71 | + video_type = "Movie" if video_type == "video.movie" else "TV Series" |
| 72 | + description = soup.head.find("meta", attrs={"name": "description"}) |
| 73 | + description = description["content"] if description else "-" |
| 74 | + time_age = soup.head.find("meta", property="og:description") |
| 75 | + time_age = time_age["content"].split("|") if time_age else "" |
| 76 | + time = time_age[0].strip() if time_age else "-" |
| 77 | + age = time_age[1].strip() if len(time_age) == 2 else "-" |
| 78 | + image = soup.head.find("meta", property="og:image") |
| 79 | + image = image["content"] if image else "" |
| 80 | + image_uri = "" |
| 81 | + |
| 82 | + try: |
| 83 | + response = await self.http.get(image, raise_for_status=True) |
| 84 | + data = await response.read() |
| 85 | + content_type = response.content_type |
| 86 | + extension = mimetypes.guess_extension(content_type) |
| 87 | + image_uri = await self.client.upload_media( |
| 88 | + data=data, |
| 89 | + mime_type=content_type, |
| 90 | + filename=f"image{extension}", |
| 91 | + size=len(data)) |
| 92 | + except aiohttp.ClientError as e: |
| 93 | + self.log.error(f"Preparing image: Connection failed: {image}: {e}") |
| 94 | + except Exception as e: |
| 95 | + self.log.error(f"Preparing image: Unknown error: {image}: {e}") |
| 96 | + |
| 97 | + body = (f"> ### {title} (url)\n> {description}\n" |
| 98 | + f"> \n" |
| 99 | + f"> > **⭐ Rating:** {rating}\n" |
| 100 | + f"> > **Runtime:** {time}\n" |
| 101 | + f"> > **Age restriction:** {age}\n" |
| 102 | + f"> > **Tags:** {tags}\n" |
| 103 | + f"> > \n" |
| 104 | + f"> > **{video_type} ・ Results from IMDb**") |
| 105 | + |
| 106 | + html = (f"<div>" |
| 107 | + f"<blockquote>" |
| 108 | + f"<a href=\"{url}\">" |
| 109 | + f"<h3>{title}</h3>" |
| 110 | + f"</a>" |
| 111 | + f"<p>{description}</p>" |
| 112 | + f"<blockquote><b>⭐ Rating:</b> {rating}</blockquote>" |
| 113 | + f"<blockquote><b>Runtime:</b> {time}</blockquote>" |
| 114 | + f"<blockquote><b>Age restriction:</b> {age}</blockquote>" |
| 115 | + f"<blockquote><b>Tags:</b> {tags}</blockquote>" |
| 116 | + f"<img src=\"{image_uri}\" width=\"300\" /><br>" |
| 117 | + f"<b><sub>{video_type} ・ Results from IMDb</sub></b>" |
| 118 | + f"</blockquote>" |
| 119 | + f"</div>") |
| 120 | + |
| 121 | + content = TextMessageEventContent( |
| 122 | + msgtype=MessageType.NOTICE, |
| 123 | + format=Format.HTML, |
| 124 | + body=body, |
| 125 | + formatted_body=html) |
| 126 | + return content |
0 commit comments