matu
/
Heroku
mirror of https://github.com/coddrago/Heroku
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add restart lock on protection, add protection from protections' flood in bot

pull/1/head
hikariatama 2022-05-01 06:32:36 +00:00
parent 70719ac6eb
commit 855379bd75
1 changed files with 20 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

58
hikka/modules/api_protection.py
View File

@ -19,6 +19,7 @@ import time
import io import io
import json import json
import random import random
import asyncio
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -67,6 +68,7 @@ class APIRatelimiterMod(loader.Module):
_ratelimiter = [] _ratelimiter = []
_suspend_until = 0 _suspend_until = 0
_lock = False
def __init__(self): def __init__(self):
self.config = loader.ModuleConfig( self.config = loader.ModuleConfig(
@ -82,11 +84,16 @@ class APIRatelimiterMod(loader.Module):
) )
async def client_ready(self, client, db): async def client_ready(self, client, db):
if hasattr(client._call, "_old_call_rewritten"): self._client = client
asyncio.ensure_future(self._install_protection())
async def _install_protection(self):
await asyncio.sleep(30) # Restart lock
if hasattr(self._client._call, "_old_call_rewritten"):
raise loader.SelfUnload("Already installed") raise loader.SelfUnload("Already installed")
self._me = (await client.get_me()).id self._me = (await self._client.get_me()).id
old_call = client._call old_call = self._client._call
async def new_call( async def new_call(
sender: "MTProtoSender", # noqa: F821 sender: "MTProtoSender", # noqa: F821
@ -108,7 +115,11 @@ class APIRatelimiterMod(loader.Module):
) )
) )
if len(self._ratelimiter) > int(self.config["threshold"]): if (
len(self._ratelimiter) > int(self.config["threshold"])
and not self._lock
):
self._lock = True
report = io.BytesIO( report = io.BytesIO(
json.dumps( json.dumps(
self._ratelimiter, self._ratelimiter,
@ -129,12 +140,13 @@ class APIRatelimiterMod(loader.Module):
# It is intented to use time.sleep instead of asyncio.sleep # It is intented to use time.sleep instead of asyncio.sleep
time.sleep(int(self.config["local_floodwait"])) time.sleep(int(self.config["local_floodwait"]))
self._lock = False
return await old_call(sender, request, ordered, flood_sleep_threshold) return await old_call(sender, request, ordered, flood_sleep_threshold)
client._call = new_call self._client._call = new_call
client._old_call_rewritten = old_call self._client._old_call_rewritten = old_call
client._call._hikka_overwritten = True self._client._call._hikka_overwritten = True
logger.debug("Successfully installed ratelimiter") logger.debug("Successfully installed ratelimiter")
async def on_unload(self): async def on_unload(self):
@ -190,37 +202,7 @@ class APIRatelimiterMod(loader.Module):
async def _confirm_step_2(self, call: InlineCall): async def _confirm_step_2(self, call: InlineCall):
await call.edit( await call.edit(
self.strings("test"), self.strings("test"),
self._generate_silly_markup("😞", "😔", self._confirm_step_3), self._generate_silly_markup("😞", "😔", self._finish),
)
async def _confirm_step_3(self, call: InlineCall):
await call.edit(
self.strings("test"),
self._generate_silly_markup("🧑‍🏫", "👨‍🏫", self._confirm_step_4),
)
async def _confirm_step_4(self, call: InlineCall):
await call.edit(
self.strings("test"),
self._generate_silly_markup("👨‍💻", "🧑‍💻", self._confirm_step_5),
)
async def _confirm_step_5(self, call: InlineCall):
await call.edit(
self.strings("test"),
self._generate_silly_markup("🤵‍♂️", "🤵", self._confirm_step_6),
)
async def _confirm_step_6(self, call: InlineCall):
await call.edit(
self.strings("test"),
self._generate_silly_markup("👩‍🚀", "🧑‍🚀", self._confirm_step_7),
)
async def _confirm_step_7(self, call: InlineCall):
await call.edit(
self.strings("test"),
self._generate_silly_markup("👨‍🚒", "🧑‍🚒", self._finish),
) )
async def _finish(self, call: InlineCall): async def _finish(self, call: InlineCall):