foto_ablage_telegrambot/fotoablage_bot.py

import os
import time

import toml
from peewee import SqliteDatabase, PostgresqlDatabase

import db_modell as db
from telegram_api import telegram_bot_api as api


def load_config():
    configfile = os.path.join(SKRIPTPFAD, "config_fotobot.toml")
    with open(configfile) as conffile:
        config = toml.loads(conffile.read())
    return config 
 
 
SKRIPTPFAD = os.path.abspath(os.path.dirname(__file__))       
CONFIG = load_config()
# Nur für Testzwecke!
#CONFIG["root_verzeichnis"] = SKRIPTPFAD

if CONFIG["db_adapter"] == "sqlite":
    db.DB_PROXY.initialize(SqliteDatabase(os.path.join(SKRIPTPFAD, "datenbank.db3")))
elif CONFIG["db_adapter"] == "postgres":
    db.DB_PROXY.initialize(PostgresqlDatabase(CONFIG["pg"]["database"],
                                              user=CONFIG["pg"]["user"],
                                              password=CONFIG["pg"]["password"]))
else:
    raise NameError("db_adapter: {} nicht gefunden".format(CONFIG["db_adapter"]))


class FotoBot(api.Bot):
    def __init__(self, token):
        super().__init__(token)
        
        self.menue = None
        self.umenue = None

        self.undefinierte_file_id = None
        self.undefinierte_message_id = None
        
        self.bot_kommando = {
            "undefiniertefotos": self.undefinierte_fotos_verarbeiten
        }
        
    def message_handler(self, message):
        if self.menue is not None:
            self.menue(message)
            return
    
    def bot_kommando_exec(self, kommando, message):
        self.menue = self.bot_kommando.get(kommando, self.kommando_unbekannt)
        self.menue(message)

    def undefinierte_fotos_verarbeiten(self, message):
        """
        Ist self.umenue auf None, so wird der parameter zu undefiniertes_foto
        Ist self.umenue auf 1, so wird parameter zur message
        """
        self.menue = self.undefinierte_fotos_verarbeiten
        if self.umenue is None:
            for undefiniertes_foto in db.FotoUndefiniert.select():
                self.send_message(chat_id=undefiniertes_foto.chat_id,
                                  text="Ordner und Notiz angeben",
                                  reply_to_message_id=undefiniertes_foto.message_id)
                self.undefinierte_file_id = undefiniertes_foto.file_id
                self.undefinierte_message_id = undefiniertes_foto.message_id
                self.umenue = 1
                break
            else:
                self.menue = None
                self.umenue = None
        elif self.umenue == 1:
            ordner, notiz = caption_analysieren(message["message"]["text"])
            ordner_db_id = db.ordner_struktur_erstellen(ordner)
            verzeichnis = verzeichnis_erstellen(ordner)

            file_id = db.FotoUndefiniert.select(db.FotoUndefiniert.file_id).where(
                db.FotoUndefiniert.message_id == self.undefinierte_message_id).scalar()
            file, file_path = self.get_file(file_id)
            dateiendung = get_dateiendung(file_path)
            file_name = db.file_speichern(ordner_db_id, notiz, dateiendung)
            datei_speichern(verzeichnis, file_name, file)
            db.FotoUndefiniert.delete().where(db.FotoUndefiniert.message_id == self.undefinierte_message_id).execute()
            self.umenue = None
            
            if db.FotoUndefiniert.get_or_none() is None:
                self.menue = None
            else:
                self.undefinierte_fotos_verarbeiten("_")

    @staticmethod
    def kommando_unbekannt(_):
        print("Unbekannt")
    

def verzeichnis_erstellen(ordner):
    verzeichnis = os.path.join(CONFIG["root_verzeichnis"], *ordner)
    os.makedirs(verzeichnis, mode=0o770, exist_ok=True)
    return verzeichnis


def get_biggest_size(photos):
    max_file_size = 0
    max_index = None
    for index, photo in enumerate(photos):
        if photo["file_size"] > max_file_size:
            max_index = index
    return max_index
    

def datei_speichern(verzeichnis, file_name, inhalt):
    file = os.path.join(verzeichnis, file_name)
    with open(file, "wb") as file_:
        file_.write(inhalt)


def caption_analysieren(caption):
    caption_split = caption.split(";")
    ordner = caption_split[0].split(",")
    ordner = [ordner_einzel.replace(" ", "") for ordner_einzel in ordner]
    try:
        notizen = caption_split[1]
        if notizen.startswith(" "):
            notizen = notizen.replace(" ", "", 1)
    except IndexError:
        notizen = []
    return ordner, notizen


def get_dateiendung(file_path):
    return file_path.split(".")[-1]


def nachricht_auswerten(bot, message):
    absender = message["message"]["from"].get("id", 0)
    if absender not in CONFIG["erlaubte_telegram_ids"]:
        bot.send_message(absender, "User nicht freigegeben")
        return
    if bot.menue is None:
        if "entities" in message["message"]:
            if "bot_command" in message["message"]["entities"][0]["type"]:  
                text = message["message"]["text"]
                kommando = text.replace("/", "")
                bot.bot_kommando_exec(kommando, message)

        if "photo" in message["message"]:
            max_index = get_biggest_size(message["message"]["photo"])
            if "caption" in message["message"]:
                ordner, notiz = caption_analysieren(message["message"]["caption"])
                ordner_db_id = db.ordner_struktur_erstellen(ordner)
                verzeichnis = verzeichnis_erstellen(ordner)

                file, file_path = bot.get_file(message["message"]["photo"][max_index]["file_id"])
                dateiendung = get_dateiendung(file_path)
                file_name = db.file_speichern(ordner_db_id, notiz, dateiendung)
                datei_speichern(verzeichnis, file_name, file)

            else:
                bot.send_message(chat_id=absender,
                                 text="Bild kann nicht verabeitet werden, Ordner wurde nicht angegeben",
                                 reply_to_message_id=message["message"]["message_id"])
                db.FotoUndefiniert.create(message_id=message["message"]["message_id"],
                                          chat_id=absender,
                                          file_id=message["message"]["photo"][max_index]["file_id"])

    else:
        bot.menue(message)


def main():
    bot = FotoBot(CONFIG["token"])
    db.db_create_table()
    db.update_root_verzeichnis(CONFIG["root_verzeichnis"])

    while True:
        nachricht_vorhanden = False
        for message in bot.get_updates():
            nachricht_vorhanden = True
            nachricht_auswerten(bot, message)
            
        if not nachricht_vorhanden and bot.menue is None:
            break
        time.sleep(0.1)
            
            
if __name__ == "__main__":
    main()