Local-Booru/run.py

import json
import os
import subprocess
import traceback
from pathlib import Path

import requests
from flask import (Flask, jsonify, render_template, render_template_string,
                    request, send_from_directory, make_response, send_file)

from system_module import db, gelbooru, image_processor, logger, tasks

if Path("config_gb.json").exists():
    CONFIG_FILE = "config_gb.json"
else:
    CONFIG_FILE = "config.json"

CONFIG_ROOT = json.load(open(CONFIG_FILE, "r", encoding="utf-8"))
CONFIG_HOST = CONFIG_ROOT["server"]["host"]
CONFIG_PORT = CONFIG_ROOT["server"]["port"]
PAGINATION_COUNT = CONFIG_ROOT["server"]["count_page"]
TARGET_MIN_SIZE = CONFIG_ROOT["server"]["thumb_size"]
GB_ID = CONFIG_ROOT["gb"]["id"]
GB_HASH = CONFIG_ROOT["gb"]["hash"]
GB_HEADERS = CONFIG_ROOT["gb"]["headers"]
FILE_FOLDERS = CONFIG_ROOT["file"]["repos"]
FILE_DW_FOLDERS = CONFIG_ROOT["file"]["dw"]
THUMB_FOLDER = CONFIG_ROOT["file"]["thumb"]
DATA_FOLDER = CONFIG_ROOT["file"]["data"]


gb = gelbooru.GB(GB_ID, GB_HASH, GB_HEADERS, FILE_DW_FOLDERS)
log = logger.Logger()
log.init_global_exception_hook()
database = db.DB(DATA_FOLDER)
app = Flask(__name__)

if not image_processor.check_ffmpeg(): print("WARNING!\n\nУ вас не установлен ffmpeg. Генерация превью для видео не доступна!")

@app.route("/")
def index():
    return render_template("index.html")

@app.route("/analytics")
def analytics():
    return render_template("analytics.html")

@app.route("/task")
def task():
    return render_template("task.html")

@app.route("/post/<path:post_id>")
def post(post_id):
    return render_template("post.html")

@app.route('/file/<path:filename>')
def serve_image(filename):
    folder_path = tasks.find_file(filename, FILE_FOLDERS)
    print(filename, FILE_FOLDERS)
    if folder_path == None:
        return make_response("Not found", 404)
    else:
        return send_from_directory(folder_path, filename)

@app.route('/thumb/<path:filename>')
def serve_thumb(filename):
    return send_from_directory(THUMB_FOLDER, filename)

@app.route("/api", methods=["GET", "POST"])
def api():
    data = request.get_json()
    if data["type"] == "search":
        search_query = data["query"]
        search_rating = data["rating"]
        if search_query == "" and search_rating == "e+q+s+g":
            results = database.get_all()
        else:
            search_query = search_query.split(" ")
            search_rating = search_rating.split("+")
            results = database.get_search(search_rating, search_query)
        if len(results) <= PAGINATION_COUNT:
            results_send = {"results_info":{"total":len(results), "page":int(data["page"]), "count_page":PAGINATION_COUNT}, "results":results}
            return jsonify(results_send)
        else:
            results_slile = results[int(data["page"])*PAGINATION_COUNT:(int(data["page"])+1)*PAGINATION_COUNT]
            results_send = {"results_info":{"total":len(results), "page":int(data["page"]), "count_page":PAGINATION_COUNT}, "results":results_slile}
            return jsonify(results_send)
    elif data["type"] == "autocomplete":
        tag_slise = data["q"]
        result_autocomplete = database.get_top_tag(tag_slise, count_tag_get=10)
        return jsonify(result_autocomplete)
    elif data["type"] == "post":
        post_id = data["post_id"]
        result_post = database.get_from_id(post_id)
        return jsonify(result_post)
    elif data["type"] == "analytic":
        return send_file(os.path.join(DATA_FOLDER, "analytic.json"), mimetype="application/json")
    elif data["type"] == "task":
        if data["task"] == "analytic":
            if tasks.calculate_analytic(database, CONFIG_ROOT):
                return "", 200
            else:
                return "", 500
        if data["task"] == "not_found_file":
            list_md5_not_found = tasks.get_files_not_found(database, CONFIG_ROOT)
            return jsonify({'list': list_md5_not_found}), 200
    elif data["type"] == "download":
        if not data or 'id' not in data:
            return jsonify({'status': 'error', 'message': 'URL not get'}), 400
        image_id = data["id"]
        try:
            data_get, url_file, file_name = gb.get_from_id(image_id)
            if database.check_uni(data_get["md5"]):
                message = f"File '{file_name}' already exists. Skipping."
                log.send("warning", "all", message)
                return jsonify({'status': 'skipped', 'message': message})
            with requests.get(url_file, stream=True, headers=GB_HEADERS, timeout=30) as r:
                r.raise_for_status()
                content_type = r.headers.get('Content-Type', '')
                if 'text/html' in content_type:
                    log.send("error", "all", "The server returned HTML instead of an image. Bot protection.")
                    return False
                with open(FILE_DW_FOLDERS+"/"+file_name, 'wb') as f:
                    for chunk in r.iter_content(chunk_size=8192):
                        f.write(chunk)
            database.add_raw(data_get)
            res_data = image_processor.generate_thumb(FILE_DW_FOLDERS+"/"+file_name, THUMB_FOLDER, TARGET_MIN_SIZE)
            if not res_data["status"]:
                print(res_data["message"])
            return jsonify({'status': 'success'})
        except requests.exceptions.RequestException as e:
            error_message = f"Network error {image_id}: {e}"
            log.send("error", "all", error_message)
            return jsonify({'status': 'error', 'message': error_message}), 500
        except Exception as e:
            traceback.print_exc()
            error_message = f"Unknown error: {e}"
            log.send("error", "all", error_message)
            return jsonify({'status': 'error', 'message': error_message}), 500
    else:
        pass
    return "", 204


if __name__ == "__main__":
    app.run(host=CONFIG_HOST, port=CONFIG_PORT)