Skip to content

Commit

Permalink
add nyaa and sukebei torrent search function
Browse files Browse the repository at this point in the history
  • Loading branch information
the-blank-x authored and Chizuru committed Nov 28, 2020
1 parent bf7731f commit 11bd13a
Show file tree
Hide file tree
Showing 6 changed files with 191 additions and 8 deletions.
1 change: 1 addition & 0 deletions Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ RUN pacman -Syu --noconfirm \
jq \
ffmpeg \
python \
fakeroot \
p7zip \
python-pip \
openssl \
Expand Down
4 changes: 2 additions & 2 deletions bot/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import os
import threading
import time

from pyrogram import Client
import aria2p
import telegram.ext as tg
from dotenv import load_dotenv
Expand Down Expand Up @@ -102,7 +102,7 @@ def getConfig(name: str):
USE_SERVICE_ACCOUNTS = False
except KeyError:
USE_SERVICE_ACCOUNTS = False

app = Client('ayanami', api_id=TELEGRAM_API, api_hash=TELEGRAM_HASH, bot_token=BOT_TOKEN)
updater = tg.Updater(token=BOT_TOKEN,use_context=True)
bot = updater.bot
dispatcher = updater.dispatcher
10 changes: 7 additions & 3 deletions bot/__main__.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
import shutil, psutil
import signal
import pickle

from pyrogram import idle
from bot import app
from os import execl, kill, path, remove
from sys import executable
import time
Expand All @@ -13,7 +14,7 @@
from bot.helper.telegram_helper.message_utils import *
from .helper.ext_utils.bot_utils import get_readable_file_size, get_readable_time
from .helper.telegram_helper.filters import CustomFilters
from .modules import authorize, list, cancel_mirror, mirror_status, mirror, clone, watch, shell, eval, anime, stickers
from .modules import authorize, list, cancel_mirror, mirror_status, mirror, clone, watch, shell, eval, anime, stickers, search


@run_async
Expand Down Expand Up @@ -96,6 +97,8 @@ def bot_help(update, context):
/weebhelp - get help for anime, manga and character module.
/stickerhelp - get help for stickers module.
/tshelp - get help for torrent search module.
'''
sendMessage(help_string, context.bot, update)

Expand Down Expand Up @@ -130,5 +133,6 @@ def main():
LOGGER.info("Bot Started!")
signal.signal(signal.SIGINT, fs_utils.exit_clean_up)


app.start()
main()
idle()
23 changes: 23 additions & 0 deletions bot/helper/custom_filters.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
from pyrogram import filters

def callback_data(data):
def func(flt, client, callback_query):
return callback_query.data in flt.data

data = data if isinstance(data, list) else [data]
return filters.create(
func,
'CustomCallbackDataFilter',
data=data
)

def callback_chat(chats):
def func(flt, client, callback_query):
return callback_query.message.chat.id in flt.chats

chats = chats if isinstance(chats, list) else [chats]
return filters.create(
func,
'CustomCallbackChatsFilter',
chats=chats
)
152 changes: 152 additions & 0 deletions bot/modules/search.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,152 @@
import time
import html
import asyncio
import aiohttp
import feedparser
from telegram.ext import run_async, CommandHandler
from telegram import ParseMode
from bot import dispatcher
from urllib.parse import quote as urlencode, urlsplit
from pyrogram import Client, filters
from pyrogram.parser import html as pyrogram_html
from pyrogram.types import InlineKeyboardMarkup, InlineKeyboardButton
from bot.helper import custom_filters
from bot import app
session = aiohttp.ClientSession()
search_lock = asyncio.Lock()
search_info = {False: dict(), True: dict()}
async def return_search(query, page=1, sukebei=False):
page -= 1
query = query.lower().strip()
used_search_info = search_info[sukebei]
async with search_lock:
results, get_time = used_search_info.get(query, (None, 0))
if (time.time() - get_time) > 3600:
results = []
async with session.get(f'https://{"sukebei." if sukebei else ""}nyaa.si/?page=rss&q={urlencode(query)}') as resp:
d = feedparser.parse(await resp.text())
text = ''
a = 0
parser = pyrogram_html.HTML(None)
for i in sorted(d['entries'], key=lambda i: int(i['nyaa_seeders']), reverse=True):
if i['nyaa_size'].startswith('0'):
continue
if not int(i['nyaa_seeders']):
break
link = i['link']
splitted = urlsplit(link)
if splitted.scheme == 'magnet' and splitted.query:
link = f'<code>{link}</code>'
newtext = f'''{a + 1}. {html.escape(i["title"])}
<b>Link:</b> {link}
<b>Size:</b> {i["nyaa_size"]}
<b>Seeders:</b> {i["nyaa_seeders"]}
<b>Leechers:</b> {i["nyaa_leechers"]}
<b>Category:</b> {i["nyaa_category"]}\n\n'''
futtext = text + newtext
if (a and not a % 10) or len((await parser.parse(futtext))['message']) > 4096:
results.append(text)
futtext = newtext
text = futtext
a += 1
results.append(text)
ttl = time.time()
used_search_info[query] = results, ttl
try:
return results[page], len(results), ttl
except IndexError:
return '', len(results), ttl

message_info = dict()
ignore = set()
@app.on_message(filters.command(['ts', 'nyaa', 'nyaasi']))
async def nyaa_search(client, message):
text = message.text.split(' ')
text.pop(0)
query = ' '.join(text)
await init_search(client, message, query, False)

@app.on_message(filters.command(['sts', 'sukebei']))
async def nyaa_search_sukebei(client, message):
text = message.text.split(' ')
text.pop(0)
query = ' '.join(text)
await init_search(client, message, query, True)

async def init_search(client, message, query, sukebei):
result, pages, ttl = await return_search(query, sukebei=sukebei)
if not result:
await message.reply_text('No results found')
else:
buttons = [InlineKeyboardButton(f'1/{pages}', 'nyaa_nop'), InlineKeyboardButton('Next', 'nyaa_next')]
if pages == 1:
buttons.pop()
reply = await message.reply_text(result, reply_markup=InlineKeyboardMarkup([
buttons
]))
message_info[(reply.chat.id, reply.message_id)] = message.from_user.id, ttl, query, 1, pages, sukebei

@app.on_callback_query(custom_filters.callback_data('nyaa_nop'))
async def nyaa_nop(client, callback_query):
await callback_query.answer(cache_time=3600)

callback_lock = asyncio.Lock()
@app.on_callback_query(custom_filters.callback_data(['nyaa_back', 'nyaa_next']))
async def nyaa_callback(client, callback_query):
message = callback_query.message
message_identifier = (message.chat.id, message.message_id)
data = callback_query.data
async with callback_lock:
if message_identifier in ignore:
await callback_query.answer()
return
user_id, ttl, query, current_page, pages, sukebei = message_info.get(message_identifier, (None, 0, None, 0, 0, None))
og_current_page = current_page
if data == 'nyaa_back':
current_page -= 1
elif data == 'nyaa_next':
current_page += 1
if current_page < 1:
current_page = 1
elif current_page > pages:
current_page = pages
ttl_ended = (time.time() - ttl) > 3600
if ttl_ended:
text = getattr(message.text, 'html', 'Search expired')
else:
if callback_query.from_user.id != user_id:
await callback_query.answer('...no', cache_time=3600)
return
text, pages, ttl = await return_search(query, current_page, sukebei)
buttons = [InlineKeyboardButton('Back', 'nyaa_back'), InlineKeyboardButton(f'{current_page}/{pages}', 'nyaa_nop'), InlineKeyboardButton('Next', 'nyaa_next')]
if ttl_ended:
buttons = [InlineKeyboardButton('Search Expired', 'nyaa_nop')]
else:
if current_page == 1:
buttons.pop(0)
if current_page == pages:
buttons.pop()
if ttl_ended or current_page != og_current_page:
await callback_query.edit_message_text(text, reply_markup=InlineKeyboardMarkup([
buttons
]))
message_info[message_identifier] = user_id, ttl, query, current_page, pages, sukebei
if ttl_ended:
ignore.add(message_identifier)
await callback_query.answer()

@run_async
def searchhelp(update, context):
help_string = '''
• /ts <i>[search query]</i>
• /nyaa <i>[search query]</i>
• /nyaasi <i>[search query]</i>
• /sts <i>[search query]</i>
• /sukebei <i>[search query]</i>
'''
update.effective_message.reply_photo("https://telegra.ph/file/87b7878ee9d6273af566f.jpg", help_string, parse_mode=ParseMode.HTML)


SEARCHHELP_HANDLER = CommandHandler("tshelp", searchhelp)
dispatcher.add_handler(SEARCHHELP_HANDLER)
9 changes: 6 additions & 3 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,9 @@ python-dotenv>=0.10
tenacity>=6.0.0
python-magic
beautifulsoup4>=4.8.2,<4.8.10
Pyrogram>=0.16.0,<0.16.10
TgCrypto>=1.1.1,<1.1.10
git+https://gitea.eponym.info/Mirrors/youtube-dl
pyrogram
TgCrypto
youtube_dl
feedparser
natsort
aiohttp

0 comments on commit 11bd13a

Please sign in to comment.