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

# Main update handler
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 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) & "*", 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) & "*", replyMarkup = replymark, parseMode="Markdown")
    OldDataCleanup()

# Start the bot
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)
if getEnv("HOOK_DOMAIN") != "":
  bot.startWebhook(getEnv("HOOK_SECRET"), getEnv("HOOK_DOMAIN") & "/" & getEnv("HOOK_SECRET"))
else:
  bot.poll(timeout=300)