190 lines
4.4 KiB
Python
190 lines
4.4 KiB
Python
import logging
|
|
import configparser
|
|
import sys
|
|
|
|
# from tgbarebot.bot import Bot
|
|
from tgbarebot.bot import Bot
|
|
from src.markov import MarkovLite, MarkovPost, Markov
|
|
|
|
|
|
# logging.basicConfig(filename='debug.log', level=logging.DEBUG)
|
|
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
|
|
|
|
|
|
def hi(bot, text, chat):
|
|
return "hi."
|
|
|
|
|
|
def debug(bot, text, chat):
|
|
return str(chat)
|
|
|
|
|
|
def makestore(config, markov: Markov):
|
|
"""Generate the store command according to markov object and config"""
|
|
try:
|
|
groups = config["AccessControl"]["groups"].split(",")
|
|
|
|
def groupcheckstore(bot, text, chat):
|
|
|
|
grp = chat["chat"]["id"]
|
|
|
|
if str(grp) in groups:
|
|
markov.store(text, grp)
|
|
|
|
return ""
|
|
|
|
return groupcheckstore
|
|
|
|
except KeyError:
|
|
|
|
def groupnocheck(bot, text, chat):
|
|
|
|
markov.store(text, 0) # set group as 0 as a global value
|
|
|
|
return ""
|
|
|
|
return groupnocheck
|
|
|
|
|
|
def makegen(config, markov: Markov):
|
|
"""Make gen command according to markov object and config"""
|
|
try:
|
|
groups = config["AccessControl"]["groups"].split(",")
|
|
|
|
def groupcheckgen(bot, text, chat):
|
|
|
|
grp = chat["chat"]["id"]
|
|
|
|
if str(grp) in groups:
|
|
|
|
try: # FUCK IT, UNDOCUMENTED FEATURE. /gen <n> will generate n messages
|
|
num = int(text.split(" ")[1])
|
|
|
|
res = ""
|
|
if num < 20: # Soft limit of 20 messages in a row.
|
|
for _ in range(num):
|
|
res += " " + markov.generate(grp)
|
|
return str(res)
|
|
|
|
except:
|
|
return markov.generate(grp)
|
|
else:
|
|
return ""
|
|
|
|
return groupcheckgen
|
|
|
|
except KeyError:
|
|
|
|
def groupnocheck(bot, text, chat):
|
|
|
|
return markov.generate(0) # set group as 0 as a global value
|
|
|
|
return groupnocheck
|
|
|
|
|
|
def makestatus(config, markov: Markov):
|
|
"""Make status command according to markov object and options"""
|
|
try:
|
|
groups = config["AccessControl"]["groups"].split(",")
|
|
|
|
def groupcheckgen(bot, text, chat):
|
|
|
|
grp = chat["chat"]["id"]
|
|
|
|
if str(grp) in groups:
|
|
|
|
return markov.status(grp)
|
|
else:
|
|
return ""
|
|
|
|
return groupcheckgen
|
|
|
|
except KeyError:
|
|
|
|
def groupnocheck(bot, text, chat):
|
|
|
|
return markov.status(0) # set group as 0 as a global value
|
|
|
|
return groupnocheck
|
|
|
|
|
|
def makedelete(config, markov: Markov):
|
|
"""Make delete command"""
|
|
try:
|
|
groups = config["AccessControl"]["groups"].split(",")
|
|
|
|
def deletecheck(bot, text, chat):
|
|
|
|
grp = chat["chat"]["id"]
|
|
args = text.split(" ")
|
|
|
|
if str(grp) in groups:
|
|
|
|
if len(args) > 1 and str(grp) == args[1]:
|
|
|
|
markov.delete(grp)
|
|
return "Database wiped."
|
|
else:
|
|
|
|
return f"To wipe the databse, write /delete {grp}"
|
|
|
|
else:
|
|
return ""
|
|
|
|
return deletecheck
|
|
|
|
except KeyError:
|
|
|
|
def deletenocheck(bot, text, chat):
|
|
# TODO Let the bot admin wipe the database
|
|
return "The database is shared across various groups. Please ask the bot admin to wipe it."
|
|
|
|
return deletenocheck
|
|
|
|
|
|
def main():
|
|
print("starting bot lmao")
|
|
|
|
# read config
|
|
config = configparser.ConfigParser()
|
|
try:
|
|
config.read("conf.ini")
|
|
except:
|
|
print("Couldn't read config file. Exiting")
|
|
exit()
|
|
|
|
# Get token
|
|
try:
|
|
token = config["Telegram"]["token"]
|
|
except KeyError:
|
|
print("No TOKEN in conf. Exiting.")
|
|
exit()
|
|
|
|
# Get working mode. Defaults to SQLITE3
|
|
try:
|
|
if config["Database"]["backend"] == "PostgreSQL":
|
|
pass # TODO
|
|
markov = MarkovLite()
|
|
except KeyError:
|
|
markov = MarkovLite(inmemory=False)
|
|
|
|
# generate the commands to feed to the telegram bot library
|
|
store = makestore(config, markov)
|
|
gen = makegen(config, markov)
|
|
status = makestatus(config, markov)
|
|
delete = makedelete(config, markov)
|
|
|
|
commands = {
|
|
"/start": hi,
|
|
"/debug": debug,
|
|
"/gen": gen,
|
|
"/status": status,
|
|
"/delete": delete,
|
|
}
|
|
|
|
b = Bot(token, commands, fallback=store)
|
|
b.poll()
|
|
|
|
|
|
main()
|