From 5bf33a49ea91044e897f52cd13ce1986d32dd0d7 Mon Sep 17 00:00:00 2001
From: baalajimaestro <me@baalajimaestro.me>
Date: Sun, 10 Apr 2022 23:24:48 +0530
Subject: [PATCH] Introduce Rate-Limiting

Ratelimits a user for 30mins if they send more than 20 requests in a minute, also a schema change

Signed-off-by: baalajimaestro <me@baalajimaestro.me>
---
 nim.cfg                |  2 +-
 nim_censor_bot.nimble  |  2 +-
 src/nim_censor_bot.nim | 42 +++++++++++++++++++++++++++++++++---------
 3 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/nim.cfg b/nim.cfg
index 302d700..f9a5a1f 100644
--- a/nim.cfg
+++ b/nim.cfg
@@ -1,2 +1,2 @@
 -d:ssl
---warning:GCUnsafe2:off
\ No newline at end of file
+--warning:GCUnsafe2:off
diff --git a/nim_censor_bot.nimble b/nim_censor_bot.nimble
index 364e823..e6f91fc 100644
--- a/nim_censor_bot.nimble
+++ b/nim_censor_bot.nimble
@@ -12,4 +12,4 @@ bin           = @["nim_censor_bot"]
 
 requires "nim >= 1.6.0"
 requires "telebot"
-requires "norm"
\ No newline at end of file
+requires "norm"
diff --git a/src/nim_censor_bot.nim b/src/nim_censor_bot.nim
index eff123d..431cc7e 100644
--- a/src/nim_censor_bot.nim
+++ b/src/nim_censor_bot.nim
@@ -1,5 +1,5 @@
 import telebot
-import std/[asyncdispatch, logging, options, strutils, random, with, os]
+import std/[asyncdispatch, logging, options, strutils, random, with, os, tables, times, sequtils]
 import norm/[model, sqlite]
 
 type
@@ -11,17 +11,38 @@ type
 
 type
   BannedUsers* = ref object of Model
-    userid*: string
+    userid*: int64
+    bantime*: float
+    bantype*: string
+
+var RateLimiter = initTable[int64, seq[float]]()
+
+let AdminID = getEnv("ADMIN_ID").split(",")
 
 func NewCensoredData*(ftype = ""; fhash = ""; fileid = ""; caption = ""):
   CensoredData = CensoredData(ftype: ftype, fhash: fhash, fileid: fileid, caption: caption)
 
-func NewBannedUsers*(userid = ""):
-  BannedUsers = BannedUsers(userid: userid)
-
+func NewBannedUsers*(userid = int64(0), bantime = 0.0, bantype = ""):
+  BannedUsers = BannedUsers(userid: userid, bantime: bantime, bantype: bantype)
 
 let dbConn* = sqlite.open("censordata.db", "", "", "")
 
+proc ManageRateLimit(): void=
+  for i in RateLimiter.keys().toSeq():
+    if RateLimiter[i][^1] - RateLimiter[i][0] >= 60:
+      RateLimiter.del(i)
+    elif len(RateLimiter[i]) >= 20:
+      var BannedUser = NewBannedUsers(i, epochTime(), "auto")
+      RateLimiter.del(i)
+      with dbConn:
+        insert BannedUser
+  if dbConn.exists(BannedUsers, "bantype = ? and ? - bantime >= 1800", "auto", epochTime()):
+    var TempData = @[NewBannedUsers()]
+    dbConn.select(TempData,  "bantype = ? and ? - bantime >= 1800", "auto", epochTime())
+    for i in TempData:
+      var e = i
+      dbConn.delete(e)
+
 dbConn.createTables(NewCensoredData())
 dbConn.createTables(NewBannedUsers())
 
@@ -35,18 +56,21 @@ proc generate_hash(): string=
     result[i] = sample(charset)
   return result
 
-const API_KEY = getEnv("TELEGRAM_TOKEN")
-
 proc updateHandler(b: Telebot, u: Update): Future[bool] {.async, gcsafe.} =
   let response = u.message.get
+  ManageRateLimit()
   if not dbConn.exists(BannedUsers, "userid = ?", response.chat.id):
+    if RateLimiter.contains(response.chat.id):
+      RateLimiter[response.chat.id].insert(epochTime())
+    else:
+      RateLimiter[response.chat.id] = @[epochTime()]
     if response.text.isSome:
       let message = response.text.get
       if message == "/start":
         discard await b.sendMessage(response.chat.id, "Hey, To create a censored post, you can share any album, video, photo, gif, sticker, etc. The messages could then be forwarded to any chat for them to view")
       elif message.contains("/ban"):
         let user = message.split(" ")
-        var BannedUser = NewBannedUsers(user[1])
+        var BannedUser = NewBannedUsers(int64(parseInt(user[1])), epochTime(), "permanent")
         with dbConn:
           insert BannedUser
         discard await b.sendMessage(response.chat.id, "Banned!")
@@ -104,6 +128,6 @@ proc updateHandler(b: Telebot, u: Update): Future[bool] {.async, gcsafe.} =
 
       discard await b.sendMessage(response.chat.id, "Censored --> " & "tg://resolve?domain=botbotbotnotabot&start=" & filehash)
 
-let bot = newTeleBot(API_KEY)
+let bot = newTeleBot(getEnv("TELEGRAM_TOKEN"))
 bot.onUpdate(updateHandler)
 bot.poll(timeout=300)