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",