diff --git a/src/nim_censor_bot.nim b/src/nim_censor_bot.nim index eeccc68..587e175 100644 --- a/src/nim_censor_bot.nim +++ b/src/nim_censor_bot.nim @@ -12,14 +12,13 @@ import std/[asyncdispatch, options, strutils, random, - with, os, tables, times, sequtils, json] -import norm/[model, sqlite] +import norm/[model, sqlite, pool] # Logging Level var L = newConsoleLogger(levelThreshold=lvlError, fmtStr="$levelname, [$time] ") @@ -49,6 +48,7 @@ var GroupMedia = initTable[int, string]() let AdminID = getEnv("ADMIN_ID").split(",") let dbConn* = getDb() +var connPool = newPool[DbConn](15) # Functions to assist adding/deleting entries from tables func NewCensoredData*(ftype = ""; fhash = ""; fileid = ""; time = 0.0; caption = ""): @@ -58,8 +58,9 @@ func NewBannedUsers*(userid = int64(0), bantime = 0.0, bantype = ""): BannedUsers = BannedUsers(userid: userid, bantime: bantime, bantype: bantype) # Create tables if they dont exist -dbConn.createTables(NewCensoredData()) -dbConn.createTables(NewBannedUsers()) +withDb(connPool): + db.createTables(NewCensoredData()) + db.createTables(NewBannedUsers()) # Ratelimits a user if there is more than 30 timestamps within 60 secs. # Resets their ratelimit counter if 60s has passed since first timestamp @@ -70,24 +71,26 @@ proc ManageRateLimit(): void= elif len(RateLimiter[i]) >= 30: 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) + withDb(connPool): + db.insert(BannedUser) + withDb(connPool): + if db.exists(BannedUsers, "bantype = ? and ? - bantime >= 1800", "auto", epochTime()): + var TempData = @[NewBannedUsers()] + db.select(TempData, "bantype = ? and ? - bantime >= 1800", "auto", epochTime()) + for i in TempData: + var e = i + db.delete(e) # Cleanup old data to save space # Deletes old data after 6 months of entry proc OldDataCleanup(): void= - if dbConn.exists(CensoredData, "? - time >= 15780000", epochTime()): - var TempData = @[NewCensoredData()] - dbConn.select(TempData, "? - time >= 15780000", epochTime()) - for i in TempData: - var e = i - dbConn.delete(e) + withDb(connPool): + if db.exists(CensoredData, "? - time >= 15780000", epochTime()): + var TempData = @[NewCensoredData()] + db.select(TempData, "? - time >= 15780000", epochTime()) + for i in TempData: + var e = i + db.delete(e) # Generates a 6-character magic that will be used to identify the file proc generate_hash(): string= @@ -101,86 +104,89 @@ proc generate_hash(): string= proc startHandler(b: Telebot, c: Command): Future[bool] {.gcsafe, async.} = let param = c.params ManageRateLimit() - if not dbConn.exists(BannedUsers, "userid = ?", c.message.chat.id): - # Update rate-limit counter - if not AdminID.contains($c.message.chat.id): - if RateLimiter.contains(c.message.chat.id): - RateLimiter[c.message.chat.id].insert(epochTime()) - else: - RateLimiter[c.message.chat.id] = @[epochTime()] - if param == "": - discard await b.sendMessage(c.message.chat.id, - "Hey, To create a censored post, you can share any album, video, photo, gif, sticker, etc. " & - "The response of the bot can be used to share the media with other chats.") - else: - if not dbConn.exists(CensoredData, "fhash = ?", param): - discard await b.sendMessage(c.message.chat.id, - "Media does not exist on database, ask the sender to censor this again!") - else: - var TempData = @[NewCensoredData()] - dbConn.select(TempData, "fhash = ?", param) - if len(TempData) > 1: - var inputmedia = newSeq[InputMediaPhoto]() - for i in TempData: - if i.caption != "": - inputmedia.insert(InputMediaPhoto(kind: i.ftype, media: i.fileid, caption: some(i.caption))) - else: - inputmedia.insert(InputMediaPhoto(kind: i.ftype, media: i.fileid)) - discard await b.sendMediaGroup($c.message.chat.id, media=inputmedia) + withDb(connPool): + if not db.exists(BannedUsers, "userid = ?", c.message.chat.id): + # Update rate-limit counter + if not AdminID.contains($c.message.chat.id): + if RateLimiter.contains(c.message.chat.id): + RateLimiter[c.message.chat.id].insert(epochTime()) else: - if TempData[0].ftype == "photo": - discard await b.sendPhoto(chatId=c.message.chat.id, photo=TempData[0].fileid, caption=TempData[0].caption) - elif TempData[0].ftype == "document": - discard await b.sendDocument(c.message.chat.id, TempData[0].fileid, caption=TempData[0].caption) - elif TempData[0].ftype == "video": - discard await b.sendVideo(c.message.chat.id, TempData[0].fileid, caption=TempData[0].caption) - elif TempData[0].ftype == "videoNote": - discard await b.sendVideoNote(c.message.chat.id, TempData[0].fileid) - elif TempData[0].ftype == "animation": - discard await b.sendAnimation(c.message.chat.id, TempData[0].fileid, caption=TempData[0].caption) - elif TempData[0].ftype == "sticker": - discard await b.sendSticker(c.message.chat.id, TempData[0].fileid) + RateLimiter[c.message.chat.id] = @[epochTime()] + if param == "": + discard await b.sendMessage(c.message.chat.id, + "Hey, To create a censored post, you can share any album, video, photo, gif, sticker, etc. " & + "The response of the bot can be used to share the media with other chats.") + else: + if not db.exists(CensoredData, "fhash = ?", param): + discard await b.sendMessage(c.message.chat.id, + "Media does not exist on database, ask the sender to censor this again!") + else: + var TempData = @[NewCensoredData()] + db.select(TempData, "fhash = ?", param) + if len(TempData) > 1: + var inputmedia = newSeq[InputMediaPhoto]() + for i in TempData: + if i.caption != "": + inputmedia.insert(InputMediaPhoto(kind: i.ftype, media: i.fileid, caption: some(i.caption))) + else: + inputmedia.insert(InputMediaPhoto(kind: i.ftype, media: i.fileid)) + discard await b.sendMediaGroup($c.message.chat.id, media=inputmedia) + else: + if TempData[0].ftype == "photo": + discard await b.sendPhoto(chatId=c.message.chat.id, photo=TempData[0].fileid, caption=TempData[0].caption) + elif TempData[0].ftype == "document": + discard await b.sendDocument(c.message.chat.id, TempData[0].fileid, caption=TempData[0].caption) + elif TempData[0].ftype == "video": + discard await b.sendVideo(c.message.chat.id, TempData[0].fileid, caption=TempData[0].caption) + elif TempData[0].ftype == "videoNote": + discard await b.sendVideoNote(c.message.chat.id, TempData[0].fileid) + elif TempData[0].ftype == "animation": + discard await b.sendAnimation(c.message.chat.id, TempData[0].fileid, caption=TempData[0].caption) + elif TempData[0].ftype == "sticker": + discard await b.sendSticker(c.message.chat.id, TempData[0].fileid) # Give them source url proc sourceHandler(b: Telebot, c: Command): Future[bool] {.gcsafe, async.} = ManageRateLimit() - if not dbConn.exists(BannedUsers, "userid = ?", c.message.chat.id): - # Update rate-limit counter - if not AdminID.contains($c.message.chat.id): - if RateLimiter.contains(c.message.chat.id): - RateLimiter[c.message.chat.id].insert(epochTime()) - else: - RateLimiter[c.message.chat.id] = @[epochTime()] - discard await b.sendMessage(c.message.chat.id, - "Hey, this bot is open-source and licensed under AGPL-v3! " & - "\n\nYou are welcome to selfhost your own instance of this bot" & - "\n\nThe source code is [here](https://git.baalajimaestro.me/baalajimaestro/nim-censor-bot)", - parseMode="Markdown", - disableWebPagePreview = true) + withDb(connPool): + if not db.exists(BannedUsers, "userid = ?", c.message.chat.id): + # Update rate-limit counter + if not AdminID.contains($c.message.chat.id): + if RateLimiter.contains(c.message.chat.id): + RateLimiter[c.message.chat.id].insert(epochTime()) + else: + RateLimiter[c.message.chat.id] = @[epochTime()] + discard await b.sendMessage(c.message.chat.id, + "Hey, this bot is open-source and licensed under AGPL-v3! " & + "\n\nYou are welcome to selfhost your own instance of this bot" & + "\n\nThe source code is [here](https://git.baalajimaestro.me/baalajimaestro/nim-censor-bot)", + parseMode="Markdown", + disableWebPagePreview = true) # UnBan Handler proc unbanHandler(b: Telebot, c: Command): Future[bool] {.gcsafe, async.} = if $c.message.chat.id in AdminID: let user = c.params - if dbConn.exists(BannedUsers, "userid = ?", int64(parseInt(user))): - var TempData = @[NewBannedUsers()] - dbConn.select(TempData, "userid = ?", int64(parseInt(user))) - for i in TempData: - var e = i - dbConn.delete(e) - discard await b.sendMessage(c.message.chat.id, "Unbanned!") + withDb(connPool): + if db.exists(BannedUsers, "userid = ?", int64(parseInt(user))): + var TempData = @[NewBannedUsers()] + db.select(TempData, "userid = ?", int64(parseInt(user))) + for i in TempData: + var e = i + db.delete(e) + discard await b.sendMessage(c.message.chat.id, "Unbanned!") # ban Handler proc banHandler(b: Telebot, c: Command): Future[bool] {.gcsafe, async.} = if $c.message.chat.id in AdminID: let user = c.params var BannedUser = NewBannedUsers(int64(parseInt(user)), epochTime(), "permanent") - with dbConn: - insert BannedUser + withDb(connPool): + db.insert(BannedUser) discard await b.sendMessage(c.message.chat.id, "Banned!") # Inline share handler -proc inlineHandler(b: Telebot, u: InlineQuery): Future[bool]{.async.} = +proc inlineHandler(b: Telebot, u: InlineQuery): Future[bool]{.gcsafe, async.} = var TempData = @[NewCensoredData()] var ftype = "" var res: InlineQueryResultArticle @@ -188,25 +194,26 @@ proc inlineHandler(b: Telebot, u: InlineQuery): Future[bool]{.async.} = res.kind = "article" res.id = "1" if u.query != "": - if not dbConn.exists(CensoredData, "fhash = ?",u.query): - res.title = "Media Not Found" - res.inputMessageContent = InputTextMessageContent( - "Media does not exist on database, ask the sender to censor this again!").some - else: - dbConn.select(TempData, "fhash = ?", u.query) - if len(TempData) > 1: - ftype = "Album" - elif len(TempData) == 1: - ftype = TempData[0].ftype - res.title = "NSFW " & capitalizeAscii(ftype) - res.inputMessageContent = InputMessageContent(kind: TextMessage, - messageText:"*NSFW " & - capitalizeAscii(ftype) & - "*\n\n[Tap to View](tg://resolve?domain=" & - b.username & "&start=" & - u.query & - ")", - parseMode: some("Markdown")).some + withDb(connPool): + if not db.exists(CensoredData, "fhash = ?",u.query): + res.title = "Media Not Found" + res.inputMessageContent = InputTextMessageContent( + "Media does not exist on database, ask the sender to censor this again!").some + else: + db.select(TempData, "fhash = ?", u.query) + if len(TempData) > 1: + ftype = "Album" + elif len(TempData) == 1: + ftype = TempData[0].ftype + res.title = "NSFW " & capitalizeAscii(ftype) + res.inputMessageContent = InputMessageContent(kind: TextMessage, + messageText:"*NSFW " & + capitalizeAscii(ftype) & + "*\n\n[Tap to View](tg://resolve?domain=" & + b.username & "&start=" & + u.query & + ")", + parseMode: some("Markdown")).some else: res.title = "Waiting for File Hash" res.inputMessageContent = InputTextMessageContent( @@ -220,82 +227,83 @@ proc updateHandler(b: Telebot, u: Update): Future[bool] {.async, gcsafe.} = let response = u.message.get # Refresh rate-limits ManageRateLimit() - # Dont bother about rate-limited/banned users - if not dbConn.exists(BannedUsers, "userid = ?", response.chat.id): - # Update rate-limit counter - if not AdminID.contains($response.chat.id): - if RateLimiter.contains(response.chat.id): - RateLimiter[response.chat.id].insert(epochTime()) - else: - RateLimiter[response.chat.id] = @[epochTime()] - if not response.text.isSome: - var fileid = "" - var ftype = "" - var fcaption = "" - if response.caption.isSome: - fcaption = response.caption.get - if response.document.isSome: - fileid = response.document.get.fileId - ftype = "document" - elif response.video.isSome: - fileid = response.video.get.fileId - ftype = "video" - elif response.videoNote.isSome: - fileid = response.videoNote.get.fileId - ftype = "videoNote" - elif response.animation.isSome: - fileid = response.animation.get.fileId - ftype = "animation" - elif response.photo.isSome: - fileid = response.photo.get[0].fileId - ftype = "photo" - elif response.sticker.isSome: - fileid = response.sticker.get.fileId - ftype = "sticker" - # Workaround for groupmedia using hashtables - # Telegram is sending multiple updates, instead of 1, so just workaround it - if response.mediaGroupId.isSome: - if parseInt(response.mediaGroupId.get) notin GroupMedia.keys().toSeq(): + withDb(connPool): + # Dont bother about rate-limited/banned users + if not db.exists(BannedUsers, "userid = ?", response.chat.id): + # Update rate-limit counter + if not AdminID.contains($response.chat.id): + if RateLimiter.contains(response.chat.id): + RateLimiter[response.chat.id].insert(epochTime()) + else: + RateLimiter[response.chat.id] = @[epochTime()] + if not response.text.isSome: + var fileid = "" + var ftype = "" + var fcaption = "" + if response.caption.isSome: + fcaption = response.caption.get + if response.document.isSome: + fileid = response.document.get.fileId + ftype = "document" + elif response.video.isSome: + fileid = response.video.get.fileId + ftype = "video" + elif response.videoNote.isSome: + fileid = response.videoNote.get.fileId + ftype = "videoNote" + elif response.animation.isSome: + fileid = response.animation.get.fileId + ftype = "animation" + elif response.photo.isSome: + fileid = response.photo.get[0].fileId + ftype = "photo" + elif response.sticker.isSome: + fileid = response.sticker.get.fileId + ftype = "sticker" + # Workaround for groupmedia using hashtables + # Telegram is sending multiple updates, instead of 1, so just workaround it + if response.mediaGroupId.isSome: + if parseInt(response.mediaGroupId.get) notin GroupMedia.keys().toSeq(): + let filehash = generate_hash() + GroupMedia[parseInt(response.mediaGroupId.get)] = filehash + var CensoredRow = NewCensoredData(ftype, filehash, fileid, epochTime(), fcaption) + withDb(connPool): + db.insert(CensoredRow) + var replybutton = InlineKeyboardButton(text: "Share", switchInlineQuery: some(filehash)) + let replymark = newInlineKeyboardMarkup(@[replybutton]) + discard await b.sendMessage(response.chat.id, + "*NSFW " & + "Album" & + "*\n\n[Tap to View](tg://resolve?domain=" & + b.username & + "&start=" & + filehash & + ")", + replyMarkup = replymark, + parseMode="Markdown") + else: + let filehash = GroupMedia[parseInt(response.mediaGroupId.get)] + var CensoredRow = NewCensoredData(ftype, filehash, fileid, epochTime(), fcaption) + withDb(connPool): + db.insert(CensoredRow) + + else: let filehash = generate_hash() - GroupMedia[parseInt(response.mediaGroupId.get)] = filehash var CensoredRow = NewCensoredData(ftype, filehash, fileid, epochTime(), fcaption) - with dbConn: - insert CensoredRow + withDb(connPool): + db.insert(CensoredRow) var replybutton = InlineKeyboardButton(text: "Share", switchInlineQuery: some(filehash)) let replymark = newInlineKeyboardMarkup(@[replybutton]) discard await b.sendMessage(response.chat.id, "*NSFW " & - "Album" & + capitalizeAscii(ftype) & "*\n\n[Tap to View](tg://resolve?domain=" & b.username & "&start=" & filehash & ")", replyMarkup = replymark, - parseMode="Markdown") - else: - let filehash = GroupMedia[parseInt(response.mediaGroupId.get)] - var CensoredRow = NewCensoredData(ftype, filehash, fileid, epochTime(), fcaption) - with dbConn: - insert CensoredRow - - else: - let filehash = generate_hash() - var CensoredRow = NewCensoredData(ftype, filehash, fileid, epochTime(), fcaption) - with dbConn: - insert CensoredRow - var replybutton = InlineKeyboardButton(text: "Share", switchInlineQuery: some(filehash)) - let replymark = newInlineKeyboardMarkup(@[replybutton]) - discard await b.sendMessage(response.chat.id, - "*NSFW " & - capitalizeAscii(ftype) & - "*\n\n[Tap to View](tg://resolve?domain=" & - b.username & - "&start=" & - filehash & - ")", - replyMarkup = replymark, - parseMode="Markdown") + parseMode="Markdown") OldDataCleanup() when isMainModule: