fedi-board-bot/bot.py

from dataclasses import dataclass
from mastodon import Mastodon
from os import environ
from sys import argv

import json
import pathlib
import random
import requests

class Config:
    block_tags: list[str]
    nsfw_tags: list[str]
    search_tags: list[str]
    text_cw: str
    text_filler: list[str]
    type: str
    url: str

    def __init__(self, config_name: str):
        with open(f"config/{config_name}.json") as file:
            config = json.load(file)

            self.block_tags = config["blockTags"]
            self.nsfw_tags = config["nsfwTags"]
            self.search_tags = config["searchTags"]
            self.text_cw = config["textCW"]
            self.text_filler = config["textFiller"]
            self.type = config["type"]
            self.url = config["url"]

def random_image_danbooru(block_tags: list[str], search_tags: list[str], url: str) -> dict | None:
    tags = []
    for tag in block_tags:
        tags.append("-" + tag.strip())
    for tag in search_tags:
        tags.append(tag.strip())
    tags = " ".join(tags)

    params = {
        "random": "1",
        "tags": tags
    }
    if environ.get("BOARD_API_KEY") and environ.get("BOARD_USERNAME"):
        params["api_key"] = environ.get("BOARD_API_KEY")
        params["login"] = environ.get("BOARD_USERNAME")
    resp = requests.get(url + "/posts.json", params=params)

    if resp.ok == False:
        print(resp.json())
        return None
    elif resp.json() == []:
        return None

    return random.choice(resp.json())

def random_image_gelbooru(block_tags: list[str], search_tags: list[str]) -> dict | None:
    tags = ["sort:random"]
    for tag in block_tags:
        tags.append("-" + tag.strip())
    for tag in search_tags:
        tags.append(tag.strip())
    tags = " ".join(tags)

    params = {
        "page": "dapi",
        "s": "post",
        "q": "index",
        "pid": 0,
        "json": "1",
        "limit": "1",
        "tags": tags
    }
    resp = requests.get("https://gelbooru.com/index.php", params=params)

    if resp.ok == False:
        print(resp.json())
        return None

    post = resp.json()

    if len(post["post"]) > 0:
        return random.choice(resp.json()["post"])
    
    return None


# Exit if config isn't specified
if len(argv) < 2:
    print(f"Usage: python3 {argv[0]} <config-name>")
    exit()

config = Config(argv[1])

match config.type:
    case "danbooru":
        post = random_image_danbooru(config.block_tags, config.search_tags, config.url)
    case "gelbooru":
        post = random_image_gelbooru(config.block_tags, config.search_tags)
    case _:
        print("'type' in config must be 'danbooru' or 'gelbooru'")
        exit()

if post == None:
    print("Couldn't get an image")
    exit()

match config.type:
    case "danbooru":
        characters = post.get("tag_string_character") or ""
        filename = f'image.{post["file_ext"]}'
        post_url = config.url + "/posts/" + str(post["id"])
    case "gelbooru":
        characters = ""
        filename = post["image"]
        post_url = "https://gelbooru.com/index.php?page=post&s=view&id=" + str(post["id"])
    case _:
        characters = ""
        filename = "img.png"
        post_url = ""

print(post)

image = requests.get(post["file_url"])

if image.ok == False:
    print("Couldn't download image")
    exit()

with open(filename, "wb") as f:
    f.write(image.content)

# Image is NSFW if rated Explicit, Sensitive or Unknown
is_nsfw = post["rating"].startswith("e") or post["rating"].startswith("s") or post["rating"] == None

if is_nsfw == False:
    for tag in config.nsfw_tags:
        # Check if any config defined NSFW tags are present
        match config.type:
            case "danbooru":
                if tag in post["tag_string"] or tag in post["tag_string_general"] or tag in post["tag_string_artist"] or tag in post["tag_string_copyright"] or tag in post["tag_string_character"] or tag in post["tag_string_meta"]:
                    is_nsfw = True
                    break
            case "gelbooru":
                if tag in post["tags"]:
                    is_nsfw = True
                    break
            case _:
                is_nsfw = True
                break

# This sucks but it works :3
post_content = [
    random.choice(config.text_filler),
    "\n\n",
    "Artist: " + post["tag_string_artist"] + "\n" if config.type == "danbooru" and len(post["tag_string_artist"]) > 0 else "",
    "Characters: " + characters + "\n" if characters != "" else "",
    "Source: " + post["source"] if len(post["source"]) > 0 else "",
    "\n\n" + post_url if post_url != "" else ""
]

mastodon = Mastodon(access_token=f"auth/{argv[1]}.auth.secret")

match config.type:
    case "danbooru":
        tags = post["tag_string"]
    case "gelbooru":
        tags = post["tags"]
    case _:
        tags = "unknown"

media = mastodon.media_post(
    description="Automatically posted image. Tagged: " + tags,
    media_file=filename
)

status = mastodon.status_post(
    "".join(post_content),
    media_ids=[media["id"]],
    sensitive=is_nsfw,
    spoiler_text="(NSFW) " + config.text_cw if is_nsfw else config.text_cw,
    visibility="unlisted"
)
print(status["url"])

pathlib.Path(filename).unlink()