async def lavalink_mitm():
await websocket.accept()
headers = websocket.headers
authorization = headers.get("authorization", None)
num_shards = headers.get("num-shards", None)
user_id = headers.get("user-id", None)
print(headers)
if authorization and \
num_shards and \
user_id:
async with aiohttp.ClientSession() as _session:
async with _session.ws_connect(lavalink_address,
headers={
"Authorization": authorization,
"Num-Shards": num_shards,
"User-Id": user_id
}) as lavalink_ws:
loop = asyncio.get_event_loop()
try:
await asyncio.gather(
copy_current_websocket_context(mitm_sender)(
lavalink_ws, authorization, num_shards, user_id),
copy_current_websocket_context(mitm_receiver)(
lavalink_ws, authorization, num_shards, user_id))
except Exception as e:
rockutils.prefix_print(str(e),
prefix="IPC Slave",
prefix_colour="light red",
text_colour="red")
else:
return abort(403)
def save_session(self, app, session, response):
domain = self.get_cookie_domain(app)
if not session:
rockutils.prefix_print(
f"Deleting session {self.prefix + session.sid}. URL: {request.path}",
prefix_colour="light red")
r.table("sessions").get(self.prefix + session.sid).delete().run(
self.rethink)
if session.modified:
response.delete_cookie(app.session_cookie_name, domain=domain)
return
cookie_exp = self.get_expiration_time(app, session)
val = self.serialize.dumps(self.serializer.dumps(dict(session)))
rockutils.prefix_print(
f"Updating session {self.prefix + session.sid}. URL: {request.path}",
prefix_colour="light yellow")
if r.table("sessions").get(self.prefix + session.sid).run(
self.rethink):
r.table("sessions").get(self.prefix + session.sid).replace({
"id":
self.prefix + session.sid,
"data":
val
}).run(self.rethink)
else:
r.table("sessions").insert({
"id": self.prefix + session.sid,
"data": val
}).run(self.rethink)
response.set_cookie(app.session_cookie_name,
session.sid,
expires=cookie_exp,
httponly=True,
domain=domain)
async def ipc_slave(cluster, auth):
if auth != config['ipc']['auth_key']:
return "Invalid authentication", 403
else:
await websocket.accept()
cluster = str(cluster).lower()
ipc_jobs[cluster] = []
clusters_initialized.add(cluster)
if not cluster in cluster_jobs:
cluster_jobs[cluster] = []
rockutils.prefix_print(f"Connected to cluster {cluster}")
last_ping = time.time()
loop = asyncio.get_event_loop()
try:
await asyncio.gather(
copy_current_websocket_context(sending)(cluster),
copy_current_websocket_context(receiving)(cluster))
except Exception as e:
rockutils.prefix_print(str(e),
prefix="IPC Slave",
prefix_colour="light red",
text_colour="red")
def reloadall(self):
ts = datetime.now()
wm = []
bm = []
try:
for ext in os.listdir("modules"):
if os.path.isfile("modules/" + ext):
if ext[-3:] == ".py":
module = ext[:-3]
work, _ = self.reload(module)
if work:
wm.append(module)
else:
bm.append(module)
te = datetime.now()
tl = te - ts
rockutils.prefix_print(
f"Loaded {len(wm+bm)} modules in {(tl.seconds * 1000000 + tl.microseconds)/1000}ms",
prefix="Modules",
prefix_colour="yellow")
return True, wm, bm, (tl.seconds * 1000000 +
tl.microseconds) / 1000
except Exception as e:
print(str(e))
return False, str(e), 0, 0
def load(self, module):
ts = datetime.now()
try:
self.bot.load_extension(f"modules.{module}")
te = datetime.now()
tl = te - ts
rockutils.prefix_print(
f"Loaded {module} in {str((tl.seconds * 1000000 + tl.microseconds) / 1000)}ms",
prefix="Modules",
prefix_colour="yellow")
return True, (tl.seconds * 1000000 + tl.microseconds) / 1000
except Exception as e:
return False, str(e)
def open_session(self, app, request):
sid = request.cookies.get(app.session_cookie_name)
if not sid:
sid = self.generate_sid()
return self.session_class(sid=sid, new=True)
rockutils.prefix_print(
f"Retreiving session {self.prefix + sid}. URL: {request.path}",
prefix_colour="light blue")
val = r.table("sessions").get(self.prefix + sid).run(self.rethink)
if val is not None:
data = self.serializer.loads(self.serialize.loads(val['data']))
return self.session_class(data, sid=sid)
return self.session_class(sid=sid, new=True)
async def get_guild_info(id, refer=""):
try:
rockutils.prefix_print(
f"{f'[Refer: {refer}] ' if refer != '' else ''}Getting information for G:{id}",
prefix="Guild Info:Get",
prefix_colour="light green")
guild_info = r.table("guilds").get(str(id)).run(self.bot.connection)
return guild_info or False
except:
rockutils.prefix_print(
f"Error occured whilst retrieving info for G:{id}. {e}",
prefix="Guild Info:Update",
prefix_colour="red",
text_colour="light red")
return False
async def receiving(cluster):
rockutils.prefix_print(f"Started receiving for {cluster}")
try:
while True:
_data = json.loads(await websocket.receive())
o = _data['o']
if o == "SUBMIT" and "ping" in _data['k']:
last_ping[cluster] = time.time()
cluster_data[cluster] = _data['d']
rockutils.prefix_print(
f"Retrieved PING from cluster {cluster}")
elif o == "STATUS_UPDATE":
d = _data['d']
cluster_status[cluster] = d
rockutils.prefix_print(
f"Cluster {cluster} is now {_status_name.get(d)}")
elif o == "SUBMIT" and _data['k'] != "push":
k = _data['k']
r = _data['r']
d = _data['d']
if k in ipc_jobs:
ipc_jobs[k][r] = d
elif o == "PUSH_OPCODE":
d = _data['d']
_opcode = d[0]
_args = d[1]
r = d[2]
recepients = []
if r == "*":
for i in clusters_initialized:
recepients.append(i)
else:
try:
for i in json.loads(r):
_c = str(i).lower()
if _c in clusters_initialized:
recepients.append(_c)
except:
_c = str(r).lower()
if _c in clusters_initialized:
recepients.append(_c)
for _r in recepients:
ipc_jobs[_r].append({
"o": _opcode,
"a": _args,
"k": "push"
})
await asyncio.sleep(0.05)
except Exception as e:
rockutils.prefix_print(str(e),
prefix="IPC Receiver",
prefix_colour="light red",
text_colour="red")
async def connect(_debug, _config):
global rethink
host = config['db']['host']
port = config['db']['port']
db = config['db']['table']
if _debug:
db += "debug"
rockutils.prefix_print(f"Connecting to DB {host}:{port}@{db}")
try:
rethink = await r.connect(host=host, port=port, db=db)
except Exception as e:
rockutils.prefix_print(f"Failed to connect to DB: {e}.",
prefix_colour="light red",
text_colour="red")
exit()
async def cache_discord(url,
bot_type,
key=None,
custom_token=None,
default={},
cachetime=120):
if not bot_type in config['tokens']:
return False, default
if not custom_token:
token = config['tokens'].get(bot_type)
else:
token = custom_token
key = token['access_token']
if not key:
key = url
url = f"{_oauth2['api_base']}/v6/{url}"
_t = time.time()
if not key in discord_cache or _t - discord_cache.get(key)['s'] > 0:
try:
rockutils.prefix_print(f"Retrieving {url}", prefix="Cacher")
async with aiohttp.Session() as _session:
async with _session.get(
url, headers={"Authorization": f"Bot {token}"}) as r:
data = await r.json()
if type(data) == dict and data.get("code"):
rockutils.prefix_print(
f"Encountered bad response: {data}",
prefix="Cacher")
discord_cache[key] = {
"d": default,
"s": _t + cachetime
}
return False, data.get("code", -1)
discord_cache[key] = {"d": data, "s": _t + cachetime}
except Exception as e:
print(f"[cache_discord] {e}")
return False, []
return True, discord_cache[key]['d']
async def sending(cluster):
rockutils.prefix_print(f"Started sending for {cluster}")
_last_ping = 0
try:
while True:
_t = time.time()
_jobs = cluster_jobs[cluster]
if _t - _last_ping > 60:
_jobs.append({"o": "PING", "a": "", "k": f"ping.{cluster}"})
_last_ping = _t
if len(_jobs) > 0:
await websocket.send(json.dumps(_jobs))
cluster_jobs[cluster] = []
await asyncio.sleep(0.05)
except Exception as e:
rockutils.prefix_print(str(e),
prefix="IPC Sender",
prefix_colour="light red",
text_colour="red")
def load(self, module):
ts = datetime.now()
try:
if f"modules.{module}" not in self.bot.extensions:
self.bot.load_extension(f"modules.{module}")
else:
self.bot.reload_extension(f"modules.{module}")
te = datetime.now()
tl = te - ts
rockutils.prefix_print(
f"Loaded {module} in {str((tl.seconds * 1000000 + tl.microseconds) / 1000)}ms",
prefix="Modules",
prefix_colour="yellow")
return True, (tl.seconds * 1000000 + tl.microseconds) / 1000
except discord.ext.commands.errors.ExtensionNotLoaded:
return False, ""
except ModuleNotFoundError:
return False, ""
except Exception as e:
return False, str(e)
async def connect(_debug, _config):
global connection
host = config['db']['host']
db = config['db']['db']
password = config['db']['password']
user = config['db']['user']
rockutils.prefix_print(f"Connecting to DB {user}@{host}")
try:
connection = await asyncpg.create_pool(user=user,
password=password,
database=db,
host=host,
max_size=20,
command_timeout=30)
except Exception as e:
rockutils.prefix_print(f"Failed to connect to DB: {e}.",
prefix_colour="light red",
text_colour="red")
exit()
async def execute_worker(self, workdata):
_tstart = time.time()
# stats
# timeroles
# analytics saver
# logs saver
# stats format
{
"c": 0,
"t": "minecraft",
"d": none,
"f": ">> {ip}"
}
workdata['processing'] = time.time() - _tstart
if workdata['processing'] > self.delay/2:
rockutils.prefix_print(f"Work job took {math.ceil(workdata['processing'] * 1000)}ms", prefix="Worker Task", prefix_colour="light red")
return workdata
def __init__(self, *args, **kwargs):
kwargs['command_prefix'] = "+"
super().__init__(*args, **kwargs)
rockutils.prefix_print("Init", text_colour="light blue")
self.connection = kwargs.get("connection")
self.affinity = kwargs.get("affinity")
self.cluster_id = kwargs.get("cluster_id")
self.cluster_count = kwargs.get("cluster_count")
print("shard_ids", self.shard_ids)
print("shard_count", self.shard_count)
self.donator = kwargs.get("donator", False)
self.debug = kwargs.get("testing", False)
self.ipc_ws = None
self.ipc_queue = []
self.cache = {}
self.jobs = []
self.config = config
self.version = config['bot']['version']
self.init_time = time.time()
# self.integrated_sync = True
# if true will process jobs in sync handler
self.load_extension("modules.core")
setattr(self, "command_prefix", self.get_prefix)
try:
self.load_extension("modules.worker")
except Exception:
exc_info = sys.exc_info()
traceback.print_exception(*exc_info)
self.sync_handle = self.loop.create_task(self.sync_task())
self.worker_handle = self.loop.create_task(self.worker_task())
async def update_guild_info(id, data, forceupdate=False, refer=""):
try:
rockutils.prefix_print(
f"{f'[Refer: {refer}] ' if refer != '' else ''}Updating information for G:{id}",
prefix="Guild Info:Update",
prefix_colour="light green")
t = time.time()
if forceupdate:
r.table("guilds").get(str(id)).update(data).run(
self.bot.connection)
else:
r.table("guilds").get(str(id)).replace(data).run(
self.bot.connection)
te = time.time()
if te - t > 1:
rockutils.prefix_print(
f"Updating guild info took {math.floor((te-t)*1000)}ms",
prefix="Guild Info:Update",
prefix_colour="red",
text_colour="light red")
return True
except Exception as e:
rockutils.prefix_print(
f"Error occured whilst updating info for G:{id}. {e}",
prefix="Guild Info:Update",
prefix_colour="red",
text_colour="light red")
return False
def __init__(self, *args, **kwargs):
rockutils.prefix_print("Init", text_colour="light blue")
self.connection = kwargs.get("connection")
self.affinity = kwargs.get("affinity")
self.shard_range = kwargs.get("shard_range")
self.shard_count = kwargs.get("shard_count")
self.cluster_id = kwargs.get("cluster_id")
self.cluster_count = kwargs.get("cluster_count")
self.donator = kwargs.get("donator", False)
self.debug = kwargs.get("testing", False)
self.ipc_ws = None
self.ipc_queue = []
self.cache = {}
self.jobs = []
self.config = config
self.init_time = time.time()
self.integrated_sync = True
# if true will process jobs in sync handler
kwargs['command_prefix'] = "+"
super().__init__(*args, **kwargs)
self.load_extension("modules.core")
setattr(self, "command_prefix", self.get_prefix)
self.load_extension("modules.worker")
self.sync_handle = self.loop.create_task(self.sync_task())
self.worker_handle = self.loop.create_task(self.worker_task())
# self.job_handle = self.loop.create_task(self.job_task())
@self.event
async def on_message(message):
return
async def exit(self, opcode, args):
if args == 0:
await self.push_ipc({"o": "STATUS_UPDATE", "d": 5})
rockutils.prefix_print("Received job to terminate")
os._exit(100)
elif args == 1:
await self.push_ipc({"o": "STATUS_UPDATE", "d": 3})
rockutils.prefix_print("Received job to hang")
os._exit(102)
elif args == 2:
await self.push_ipc({"o": "STATUS_UPDATE", "d": 4})
rockutils.prefix_print("Received job to restart")
os._exit(101)
async def worker_task(self):
rockutils.prefix_print(f"Awaiting ready", prefix="Worker Task")
while not self.bot.is_ready():
await asyncio.sleep(5)
rockutils.prefix_print(f"Starting worker", prefix="Worker Task")
_time = time.time()
workdata = {}
while True:
try:
workdata = await self.bot.execute_worker(workdata)
except Exception as e:
rockutils.prefix_print(str(e), prefix="Worker Task", prefix_colour="light red", text_colour="red")
await asyncio.sleep(self.delay - workdata['processing'])
async def before_identify_hook(self, shard_id, initial):
identify_key = f"{hashlib.md5(self.http.token.encode()).hexdigest()}:{shard_id % self.config['bot']['max_concurrency']}:{self.shard_count}"
domain = f"http://{self.config['ipc']['host']}:{self.config['ipc']['port']}/api/ipc_identify/{identify_key}/{self.cluster_id}/{self.config['ipc']['auth_key']}"
start = time.time()
rockutils.prefix_print(
f"Retrieving identify lock for shard {shard_id}",
prefix="IdentifyHook")
async with aiohttp.ClientSession() as session:
while True:
async with session.post(domain) as resp:
data = await resp.json()
if data['available']:
rockutils.prefix_print(
f"Retireved identify lock for shard {shard_id} in {math.ceil((time.time()-start)*1000)} ms",
prefix="IdentifyHook")
break
else:
rockutils.prefix_print(
f"Waiting {data['sleep']} seconds for identify lock for shard {shard_id}",
prefix="IdentifyHook")
await asyncio.sleep(data['sleep'])
from rockutils import rockutils
import sys
import time
import os
#argument = "LD_PRELOAD=/usr/local/lib/libjemalloc.so " + sys.argv[1]
argument = sys.argv[1]
while True:
rockutils.prefix_print(argument, prefix="Wrapper")
try:
exit_code = os.system(argument)
except Exception as e:
rockutils.prefix_print(
"Cluster has encountered fatal error and has to stop",
prefix="Wrapper",
prefix_colour="red")
rockutils.prefix_print(
str(e),
prefix="Wrapper",
prefix_colour="red",
text_colour="light red")
time.sleep(5)
if exit_code == 100:
rockutils.prefix_print("Cluster has stopped",
prefix="Wrapper", prefix_colour="red")
elif exit_code == 101:
rockutils.prefix_print("Cluster is restarting",
prefix="Wrapper", prefix_colour="green")
elif exit_code == 102:
import asyncio
import aiohttp
import os
import ujson as json
from rockutils import rockutils
from quart import Response, Quart, send_file, request, jsonify, abort, websocket
from quart.ctx import copy_current_websocket_context
app = Quart(__name__)
config = rockutils.load_json("cfg/config.json")
os.environ['OAUTHLIB_INSECURE_TRANSPORT'] = 'true'
lavalink_address = f"ws://{':'.join(map(str,config['lavalink']['server']))}"
rockutils.prefix_print(f"Lavalink address set to: {lavalink_address}")
async def mitm_sender(lavalink_ws, authorization, num_shards, user_id):
# get information from bot and relay to lavalink
while True:
data = await websocket.receive()
print(f"[>>] {data}")
await lavalink_ws.send_str(data)
async def mitm_receiver(lavalink_ws, authorization, num_shards, user_id):
# get information from lavalink and relay to bot
while True:
async for msg in lavalink_ws:
print(f"[<<] {msg.data}")
await websocket.send(msg.data)
import rethinkdb as r
import ujson as json
from datetime import datetime
from discord.ext import commands
from rockutils import gameretriever, rockutils
_arg = sys.argv
# rockutils.prefix_print("test", prefix="test 123")
# rockutils.add_lang("en")
# print(rockutils._("Test file %s","en"))
# print(rockutils._("Test file %s"))
if len(_arg) == 1:
rockutils.prefix_print(f"{_arg[0]} <cluster id>")
exit()
r.set_loop_type("asyncio")
config = rockutils.load_json("cfg/config.json")
rockutils.prefix_print(
f"IPC address: {config['ipc']['host']}:{config['ipc']['port']}")
rockutils.prefix_print(
f"CDN address: {config['cdn']['host']}:{config['cdn']['port']}")
cluster_count = config['bot']['clusters']
cluster_id = _arg[1]
if cluster_id.lower() == "debug":
from rockutils import rockutils
logger = logging.getLogger('discord')
logger.setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO,
filename='discord.log',
filemode='a',
format='%(asctime)s:%(levelname)s:%(name)s: %(message)s')
logging.getLogger().addHandler(logging.StreamHandler())
# logger.addHandler(handler)
_arg = sys.argv
if len(_arg) == 1:
rockutils.prefix_print(f"{_arg[0]} <cluster id>")
exit()
logging.basicConfig(level=logging.DEBUG)
rockutils.prefix_print(" _ __ __ ",
text_colour="cyan")
rockutils.prefix_print("| | / /__ / /________ ____ ___ ___ _____",
text_colour="cyan")
rockutils.prefix_print(r"| | /| / / _ \/ / ___/ __ \/ __ `__ \/ _ \/ ___/",
text_colour="cyan")
rockutils.prefix_print("| |/ |/ / __/ / /__/ /_/ / / / / / / __/ / ",
text_colour="cyan")
rockutils.prefix_print(r"|__/|__/\___/_/\___/\____/_/ /_/ /_/\___/_/ ",
text_colour="cyan")
rockutils.prefix_print("Made by ImRock#0001 - Ver. 09/05/2021",
