mirror of
https://github.com/polyjitter/searchbot-discord.git
synced 2024-08-14 22:46:55 +00:00
Full Rejigging of Structure
This commit is contained in:
Adriene Hutchins
parent
c59de307dc
commit
e235e0cee3
4 changed files with 264 additions and 173 deletions
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -1 +1,3 @@
|
|||
config.json
|
||||
extensions/__pycache__/developer.cpython-38.pyc
|
||||
extensions/__pycache__/search.cpython-38.pyc
|
||||
|
|
|
|||
92
extensions/developer.py
Normal file
92
extensions/developer.py
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
|
||||
# Developer functions
|
||||
# Provides functions only useable for developers
|
||||
|
||||
'''Developer Cog'''
|
||||
|
||||
import discord
|
||||
from discord.ext import commands
|
||||
import aiohttp
|
||||
import random
|
||||
|
||||
class Developer(commands.Cog):
|
||||
def __init__(self, bot):
|
||||
self.bot = bot
|
||||
self.session = bot.session
|
||||
self.instances = bot.instances
|
||||
|
||||
async def _instance_check(self, instance, info):
|
||||
'''Checks the quality of an instance.'''
|
||||
|
||||
# Makes sure proper values exist
|
||||
if 'error' in info:
|
||||
return False
|
||||
if not ('engines' in info and 'initial' in info['timing']):
|
||||
return False
|
||||
if not ('google' in info['engines'] and 'enabled' in info['engines']['google']):
|
||||
return False
|
||||
|
||||
# Makes sure google is enabled
|
||||
if not info['engines']['google']['enabled']:
|
||||
return False
|
||||
|
||||
# Makes sure is not Tor
|
||||
if info['network_type'] != 'normal':
|
||||
return False
|
||||
|
||||
# Only picks instances that are fast enough
|
||||
timing = int(info['timing']['initial'])
|
||||
if timing > 0.20:
|
||||
return False
|
||||
|
||||
# Check for Google captcha
|
||||
test_search = f'{instance}/search?q=test&format=json&lang=en-US'
|
||||
try:
|
||||
async with self.session.get(test_search) as resp:
|
||||
response = await resp.json()
|
||||
response['results'][0]['content']
|
||||
except (aiohttp.ClientError, KeyError, IndexError):
|
||||
return False
|
||||
|
||||
# Reached if passes all checks
|
||||
return True
|
||||
|
||||
@commands.command()
|
||||
async def rejson(self, ctx):
|
||||
'''Refreshes the list of instances for searx.'''
|
||||
|
||||
msg = await ctx.send('<a:updating:403035325242540032> Refreshing instance list...\n\n'
|
||||
'(Due to extensive quality checks, this may take a bit.)')
|
||||
plausible = []
|
||||
|
||||
# Get, parse, and quality check all instances
|
||||
async with self.session.get('https://searx.space/data/instances.json') as r:
|
||||
# Parsing
|
||||
searx_json = await r.json()
|
||||
instances = searx_json['instances']
|
||||
|
||||
# Quality Check
|
||||
for i in instances:
|
||||
info = instances.get(i)
|
||||
is_good = await self._instance_check(i, info)
|
||||
if is_good:
|
||||
plausible.append(i)
|
||||
|
||||
# Save new list
|
||||
self.instances = plausible
|
||||
with open('searxes.txt', 'w') as f:
|
||||
f.write('\n'.join(plausible))
|
||||
|
||||
await msg.edit(content='Instances refreshed!')
|
||||
|
||||
@commands.command(aliases=['exit', 'reboot'])
|
||||
async def restart(self, ctx):
|
||||
await ctx.send(':zzz: **Restarting.**')
|
||||
exit()
|
||||
|
||||
async def cog_check(self, ctx):
|
||||
return commands.is_owner()
|
||||
|
||||
def setup(bot):
|
||||
bot.add_cog(Developer(bot))
|
||||
113
extensions/search.py
Normal file
113
extensions/search.py
Normal file
|
|
@ -0,0 +1,113 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
|
||||
# Search Functionality
|
||||
# Provides search results from SearX
|
||||
|
||||
'''Search Cog'''
|
||||
|
||||
import discord
|
||||
from discord.ext import commands
|
||||
import aiohttp
|
||||
import random
|
||||
|
||||
class Search(commands.Cog):
|
||||
def __init__(self, bot):
|
||||
self.bot = bot
|
||||
self.session = bot.session
|
||||
self.instances = bot.instances
|
||||
|
||||
async def _search_logic(self, query: str, type: str = None):
|
||||
"""Provides search logic for all search commands."""
|
||||
|
||||
# Choose an instance
|
||||
if self.instances == []:
|
||||
with open('searxes.txt') as f:
|
||||
self.instances = f.read().split('\n')
|
||||
instance = random.sample(self.instances, k=1)[0]
|
||||
print(f"Attempting to use {instance}")
|
||||
|
||||
# Error Template
|
||||
error_msg = ("**An error occured!**\n\n"
|
||||
f"There was a problem with `{instance}`. Please try again later.\n"
|
||||
f"_If problems with this instance persist, contact`{self.bot.appinfo.owner}` to have it removed._")
|
||||
|
||||
# Create the URL to make an API call to
|
||||
call = f'{instance}/search?q={query}&format=json&language=en-US'
|
||||
|
||||
# If a type is provided, add that type to the call URL
|
||||
if type:
|
||||
call += f'?type={type}'
|
||||
|
||||
# Make said API call
|
||||
try:
|
||||
async with self.session.get(call) as resp:
|
||||
response = await resp.json()
|
||||
except aiohttp.ClientError:
|
||||
return error_msg
|
||||
|
||||
# Split our response data up for parsing
|
||||
# infoboxes = response['infoboxes']
|
||||
results = response['results']
|
||||
|
||||
# Create message with results
|
||||
try:
|
||||
# Handle tiny result count
|
||||
if len(results) > 5:
|
||||
amt = 5
|
||||
else:
|
||||
amt = len(results)
|
||||
|
||||
# Header
|
||||
msg = f"Showing **{amt}** results for `{query}`. \n\n"
|
||||
# Expanded Result
|
||||
msg += (
|
||||
f"**{results[0]['title']}** <{results[0]['url']}>\n"
|
||||
f"{results[0]['content']}\n\n")
|
||||
# Other Results
|
||||
msg += "\n".join(
|
||||
[f"**{entry['title']}** <{entry['url']}>" for entry in results[1:5]])
|
||||
# Instance Info
|
||||
msg += f"\n\n_Results retrieved from instance `{instance}`._"
|
||||
|
||||
# Reached if error with returned results
|
||||
except (KeyError, IndexError) as e:
|
||||
# Logging
|
||||
print(f"{e} with instance {instance}, trying again.")
|
||||
|
||||
self.instances.remove(instance) # Weed the instance out
|
||||
return await self._search_logic(query) # Recurse until good response
|
||||
|
||||
return msg
|
||||
|
||||
@commands.command()
|
||||
async def search(self, ctx, *, query: str):
|
||||
"""Search online for results."""
|
||||
|
||||
# Logging
|
||||
print(f"\n\nNEW CALL: {ctx.author} from {ctx.guild}.\n")
|
||||
|
||||
# Handling
|
||||
async with ctx.typing():
|
||||
msg = await self._search_logic(query)
|
||||
await ctx.send(msg)
|
||||
|
||||
@commands.Cog.listener()
|
||||
async def on_command_error(self, ctx, error):
|
||||
"""Listener makes no command fallback to searching."""
|
||||
|
||||
if isinstance(error, commands.CommandNotFound):
|
||||
# Logging
|
||||
print(f"\n\nNEW CALL: {ctx.author} from {ctx.guild}.\n")
|
||||
|
||||
# Handling
|
||||
async with ctx.typing():
|
||||
# Prepares term
|
||||
term = ctx.message.content.replace(ctx.prefix, '', 1)
|
||||
term = term.lstrip(' ')
|
||||
# Does search
|
||||
msg = await self._search_logic(term)
|
||||
# Sends result
|
||||
await ctx.send(msg)
|
||||
|
||||
def setup(bot):
|
||||
bot.add_cog(Search(bot))
|
||||
226
main.py
226
main.py
|
|
@ -9,7 +9,10 @@
|
|||
|
||||
import discord
|
||||
from discord.ext import commands
|
||||
import traceback
|
||||
import json
|
||||
import os
|
||||
import asyncio
|
||||
import aiohttp
|
||||
import random
|
||||
|
||||
|
|
@ -20,26 +23,44 @@ class Bot(commands.Bot):
|
|||
def __init__(self, **options):
|
||||
super().__init__(self.get_prefix_new, **options)
|
||||
print('Performing initialization...\n')
|
||||
|
||||
# Create Session
|
||||
# NOTE Coro created due to a deprecation in aiohttp.
|
||||
asyncio.run(self.create_session())
|
||||
|
||||
# Get Config Values
|
||||
with open('config.json') as f:
|
||||
self.config = json.load(f)
|
||||
self.prefix = self.config.get('PREFIX')
|
||||
self.version = self.config.get('VERSION')
|
||||
self.maintenance = self.config.get('MAINTENANCE')
|
||||
|
||||
# Get Instances
|
||||
with open('searxes.txt') as f:
|
||||
self.instances = f.read().split('\n')
|
||||
|
||||
self.init_extensions()
|
||||
|
||||
print('Initialization complete.\n\n')
|
||||
|
||||
async def get_prefix_new(self, bot, msg):
|
||||
return commands.when_mentioned_or(*self.prefix)(bot, msg)
|
||||
|
||||
async def on_ready(self):
|
||||
self.appinfo = await bot.application_info()
|
||||
async def create_session(self):
|
||||
self.session = aiohttp.ClientSession()
|
||||
|
||||
def init_extensions(self):
|
||||
for ext in os.listdir('extensions'):
|
||||
if ext.endswith('.py'):
|
||||
self.load_extension(f'extensions.{ext[:-3]}')
|
||||
|
||||
async def on_ready(self):
|
||||
appinfo = await self.application_info()
|
||||
|
||||
msg = "CONNECTED!\n"
|
||||
msg += "-----------------------------\n"
|
||||
msg += f"ACCOUNT: {bot.user}\n"
|
||||
msg += f"OWNER: {self.appinfo.owner}\n"
|
||||
msg += f"OWNER: {appinfo.owner}\n"
|
||||
msg += "-----------------------------\n"
|
||||
|
||||
print(msg)
|
||||
|
|
@ -57,192 +78,55 @@ class Bot(commands.Bot):
|
|||
elif message.content in mentions:
|
||||
assist_msg = (
|
||||
"**Hi there! How can I help?**\n\n"
|
||||
f"You may use **{self.user.mention} `term here`** to search, or **{self.user.mention} `help`** for assistance.")
|
||||
# Two New Lines Here
|
||||
f"You may use **{self.user.mention} `term here`** to search,"
|
||||
"or **{self.user.mention} `help`** for assistance.")
|
||||
await ctx.send(assist_msg)
|
||||
else:
|
||||
await self.process_commands(message)
|
||||
|
||||
|
||||
bot = Bot(
|
||||
description='search - a tiny little search utility bot for discord.',
|
||||
case_insensitive=True)
|
||||
|
||||
@bot.command()
|
||||
async def search(ctx, *, query: str):
|
||||
"""Search online for results."""
|
||||
|
||||
print(f"\n\nNEW CALL: {ctx.author} from {ctx.guild}.\n")
|
||||
|
||||
async with ctx.typing():
|
||||
msg = await search_logic(query)
|
||||
await ctx.send(msg)
|
||||
|
||||
@bot.command(aliases=['exit', 'reboot'])
|
||||
@commands.is_owner()
|
||||
async def restart(ctx):
|
||||
await ctx.send(':zzz: **Restarting.**')
|
||||
exit()
|
||||
|
||||
@bot.command()
|
||||
@commands.is_owner()
|
||||
async def rejson(ctx):
|
||||
'''Refreshes the list of instances for searx.'''
|
||||
|
||||
msg = await ctx.send('<a:updating:403035325242540032> Refreshing instance list...\n\n'
|
||||
'(Due to extensive quality checks, this may take a bit.)')
|
||||
plausible = []
|
||||
|
||||
# Get, parse, and quality check all instances
|
||||
async with bot.session.get('https://searx.space/data/instances.json') as r:
|
||||
|
||||
# Parsing
|
||||
searx_json = await r.json()
|
||||
instances = searx_json['instances']
|
||||
|
||||
# Quality Check
|
||||
for i in instances:
|
||||
info = instances.get(i)
|
||||
is_good = await instance_check(i, info)
|
||||
if is_good:
|
||||
plausible.append(i)
|
||||
|
||||
# Save new list
|
||||
with open('searxes.txt', 'w') as f:
|
||||
f.write('\n'.join(plausible))
|
||||
|
||||
await msg.edit(content='Instances refreshed!')
|
||||
|
||||
async def search_logic(query: str, type: str = None):
|
||||
'''Provides search logic for all search commands.'''
|
||||
|
||||
# Choose an instance & distribute load
|
||||
if bot.instances == []:
|
||||
with open('searxes.txt') as f:
|
||||
bot.instances = f.read().split('\n')
|
||||
instance = random.sample(bot.instances, k=1)[0]
|
||||
print(f"Attempting to use {instance}")
|
||||
|
||||
# Error Template
|
||||
error_msg = ("**An error occured!**\n\n"
|
||||
f"There was a problem with `{instance}`. Please try again later.\n"
|
||||
f"_If problems with this instance persist, contact`{bot.appinfo.owner}` to have it removed._")
|
||||
|
||||
# Create the URL to make an API call to
|
||||
call = f'{instance}/search?q={query}&format=json&language=en-US'
|
||||
|
||||
if type:
|
||||
call += f'?type={type}'
|
||||
|
||||
# Make said API call
|
||||
try:
|
||||
async with bot.session.get(call) as resp:
|
||||
response = await resp.json()
|
||||
except aiohttp.ClientError:
|
||||
return error_msg
|
||||
|
||||
# Split our response data up for parsing
|
||||
# infoboxes = response['infoboxes']
|
||||
results = response['results']
|
||||
|
||||
# Create message with results
|
||||
try:
|
||||
msg = f"Showing **5** results for `{query}`. \n\n"
|
||||
msg += (f"**{results[0]['title']}** <{results[0]['url']}>\n"
|
||||
f"{results[0]['content']}\n\n")
|
||||
msg += "\n".join(
|
||||
[f"**{entry['title']}** <{entry['url']}>" for entry in results[1:5]])
|
||||
msg += f"\n\n_Results retrieved from instance `{instance}`._"
|
||||
except (KeyError, IndexError) as e:
|
||||
# Reached if error with returned results
|
||||
print(f"{e} with instance {instance}, trying again.")
|
||||
bot.instances.remove(instance)
|
||||
return await search_logic(query) # Recurse until good response
|
||||
|
||||
# Send message
|
||||
return msg
|
||||
|
||||
async def instance_check(instance, info):
|
||||
'''Checks the quality of an instance.'''
|
||||
|
||||
# Makes sure proper values exist
|
||||
if 'error' in info:
|
||||
return False
|
||||
if not ('engines' in info and 'initial' in info['timing']):
|
||||
return False
|
||||
if not ('google' in info['engines'] and 'enabled' in info['engines']['google']):
|
||||
return False
|
||||
|
||||
# Makes sure google is enabled
|
||||
if not info['engines']['google']['enabled']:
|
||||
return False
|
||||
|
||||
# Makes sure is not Tor
|
||||
if info['network_type'] != 'normal':
|
||||
return False
|
||||
|
||||
# Only picks instances that are fast enough
|
||||
timing = int(info['timing']['initial'])
|
||||
if timing > 0.20:
|
||||
return False
|
||||
|
||||
# Check for Google captcha
|
||||
test_search = f'{instance}/search?q=test&format=json&lang=en-US'
|
||||
try:
|
||||
async with bot.session.get(test_search) as resp:
|
||||
response = await resp.json()
|
||||
response['results'][0]['content']
|
||||
except (aiohttp.ClientError, KeyError, IndexError):
|
||||
return False
|
||||
|
||||
# Reached if passes all checks
|
||||
return True
|
||||
|
||||
@bot.listen()
|
||||
async def on_command_error(ctx, error):
|
||||
|
||||
if isinstance(error, commands.CommandNotFound):
|
||||
return
|
||||
elif isinstance(error, commands.CommandInvokeError):
|
||||
error = error.original
|
||||
_traceback = traceback.format_tb(error.__traceback__)
|
||||
_traceback = ''.join(_traceback)
|
||||
appinfo = await bot.application_info()
|
||||
|
||||
print(f"\n\nNEW CALL: {ctx.author} from {ctx.guild}.\n")
|
||||
embed_fallback = f"**An error occured: {type(error).__name__}. Please contact {appinfo.owner}.**"
|
||||
|
||||
async with ctx.typing():
|
||||
# Prepares term
|
||||
term = ctx.message.content.replace(ctx.prefix, '', 1)
|
||||
term = term.lstrip(' ')
|
||||
# Does search
|
||||
msg = await search_logic(term)
|
||||
# Sends result
|
||||
await ctx.send(msg)
|
||||
error_embed = discord.Embed(
|
||||
title=f"{type(error).__name__}",
|
||||
color=0xFF0000,
|
||||
description=(
|
||||
"This is (probably) a bug. This has been not been automatically "
|
||||
f"reported, so please give **{appinfo.owner}** a heads-up in DMs.")
|
||||
)
|
||||
|
||||
# elif isinstance(error, commands.CommandInvokeError):
|
||||
# error = error.original
|
||||
# _traceback = traceback.format_tb(error.__traceback__)
|
||||
# _traceback = ''.join(_traceback)
|
||||
# embed_fallback = "**ERROR: <@97788939196182528>**"
|
||||
|
||||
# error_embed = discord.Embed(
|
||||
# title="An error has occurred.",
|
||||
# color=0xFF0000,
|
||||
# description=(
|
||||
# "This is (probably) a bug. This has been automatically "
|
||||
# f"reported, but give **{bot.appinfo.owner}** a heads-up in DMs.")
|
||||
# )
|
||||
|
||||
# trace_content = (
|
||||
# "```py\n\nTraceback (most recent call last):"
|
||||
# "\n{}{}: {}```").format(
|
||||
# _traceback,
|
||||
# type(error).__name__,
|
||||
# error)
|
||||
|
||||
# error_embed.add_field(
|
||||
# name="`{}` in command `{}`".format(
|
||||
# type(error).__name__, ctx.command.qualified_name),
|
||||
# value=(trace_content[:1018] + '...```')
|
||||
# if len(trace_content) > 1024
|
||||
# else trace_content)
|
||||
# await ctx.send(embed_fallback, embed=error_embed)
|
||||
trace_content = (
|
||||
"```py\n\nTraceback (most recent call last):"
|
||||
"\n{}{}: {}```").format(
|
||||
_traceback,
|
||||
type(error).__name__,
|
||||
error)
|
||||
|
||||
error_embed.add_field(
|
||||
name="`{}` in command `{}`".format(
|
||||
type(error).__name__, ctx.command.qualified_name),
|
||||
value=(trace_content[:1018] + '...```')
|
||||
if len(trace_content) > 1024
|
||||
else trace_content)
|
||||
await ctx.send(embed_fallback, embed=error_embed)
|
||||
else:
|
||||
await ctx.send(error)
|
||||
|
||||
bot.run(bot.config['TOKEN'])
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue