nim-censor-bot/src/nim_censor_bot.nim

#
# Copyright © 2022 Maestro Creativescape
#
# SPDX-License-Identifier: AGPL-3.0-or-later
#

import telebot
import std/[asyncdispatch, logging, options, strutils, random, with, os, tables, times, sequtils, json]
import norm/[model, sqlite]

# Logging Level
var L = newConsoleLogger(levelThreshold=lvlError, fmtStr="$levelname, [$time] ")
addHandler(L)

# Custom Types Defined to handle tables on norm
type
  CensoredData* = ref object of Model
    ftype*: string
    fhash*: string
    fileid*: string
    time*: float
    caption*: string

type
  BannedUsers* = ref object of Model
    userid*: int64
    bantime*: float
    bantype*: string

# Hashmaps for RateLimiting and working around telegram's album issue
var RateLimiter = initTable[int64, seq[float]]()

var GroupMedia = initTable[int, string]()

# Bot Admins, is comma-separated variable
let AdminID = getEnv("ADMIN_ID").split(",")

let dbConn* = getDb()

# Functions to assist adding/deleting entries from tables
func NewCensoredData*(ftype = ""; fhash = ""; fileid = ""; time = 0.0; caption = ""):
  CensoredData = CensoredData(ftype: ftype, fhash: fhash, fileid: fileid, time: time, caption: caption)

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())

# Ratelimits a user if there is more than 20 timestamps within 60 secs. 
# Resets their ratelimit counter if 60s has passed since first timestamp
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)

# 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)

# Generates a 6-character magic that will be used to identify the file
proc generate_hash(): string=
  result = newString(7) 
  const charset = {'a' .. 'z', 'A' .. 'Z', '0' .. '9'}
  for i in 0..6:
    result[i] = sample(charset)
  return result

# Start command handler
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 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 messages could then be forwarded to any chat for them to view")
    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)
        else:
          if TempData[0].ftype == "photo":
            discard await b.sendPhoto(c.message.chat.id, TempData[0].fileid, TempData[0].caption)
          elif TempData[0].ftype == "document":
            discard await b.sendDocument(c.message.chat.id, TempData[0].fileid, 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)

# 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!")

# 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
    discard await b.sendMessage(c.message.chat.id, "Banned!")

# Inline share handler
proc inlineHandler(b: Telebot, u: InlineQuery): Future[bool]{.async.} =
    var TempData = @[NewCensoredData()]
    var ftype = ""
    var res: InlineQueryResultArticle
    var results: seq[InlineQueryResultArticle]
    res.kind = "article"
    res.id = "1"
    if not dbConn.exists(CensoredData, "fhash = ?",u.query):
      res.title = "Media Not Found"
      res.inputMessageContent = InputTextMessageContent("Media Not Found").some
    else:
      dbConn.select(TempData, "fhash = ?", u.query)
      if len(TempData) > 1:
        ftype = "Album"
      elif len(TempData) == 1:
        ftype = TempData[0].ftype
      res.title = "Censored " & capitalizeAscii(ftype)
      res.inputMessageContent = InputMessageContent(kind: TextMessage, messageText:"*Censored " & capitalizeAscii(ftype) & "*\n\n[Tap to View](tg://resolve?domain=" & b.username & "&start=" & u.query & ")", parseMode: some("Markdown")).some
    results.add(res)
    discard await b.answerInlineQuery(u.id, results)

# Main update handler
proc updateHandler(b: Telebot, u: Update): Future[bool] {.async, gcsafe.} =
  if u.message.isSome:
    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 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)
            with dbConn:
                  insert CensoredRow
            var replybutton = InlineKeyboardButton(text: "Share", switchInlineQuery: some(filehash))
            let replymark = newInlineKeyboardMarkup(@[replybutton])
            discard await b.sendMessage(response.chat.id, "*Censored " & 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, "*Censored " & capitalizeAscii(ftype) & "*\n\n[Tap to View](tg://resolve?domain=" & b.username & "&start=" & filehash & ")", replyMarkup = replymark, parseMode="Markdown")
      OldDataCleanup()

when isMainModule:
  let bot = newTeleBot(getEnv("TELEGRAM_TOKEN"))
  echo "*********************"
  echo "CensorBot is running!"
  echo "*********************"

  var commands = @[
      BotCommand(command: "start" , description: "Start the bot!")
  ]

  discard waitFor bot.setMyCommands(commands)
  bot.onUpdate(updateHandler)
  bot.onCommand("start", startHandler)
  bot.onCommand("ban", banHandler)
  bot.onCommand("unban", unbanHandler)
  bot.onInlineQuery(inlineHandler)
  if getEnv("HOOK_DOMAIN") != "":
    bot.startWebhook(getEnv("HOOK_SECRET"), getEnv("HOOK_DOMAIN") & "/" & getEnv("HOOK_SECRET"))
  else:
    bot.poll(timeout=300)