sin0: add telegram danbooru_img_bot

This commit is contained in:
Guanran Wang 2024-10-07 16:07:24 +08:00
parent 0a60d56e8d
commit 9f3510a13a
Signed by: nyancat
GPG key ID: 91F97D9ED12639CF
5 changed files with 269 additions and 1 deletions

View file

@ -23,6 +23,11 @@ creation_rules:
- age: - age:
- *guanranwang - *guanranwang
- *tyo0 - *tyo0
- path_regex: ^hosts/vultr/sin0/secrets.yaml$
key_groups:
- age:
- *guanranwang
- *sin0
# shared # shared
- path_regex: ^nixos/profiles/restic/secrets.yaml$ - path_regex: ^nixos/profiles/restic/secrets.yaml$

View file

@ -1,9 +1,10 @@
{ ... }: { lib, ... }:
{ {
imports = [ imports = [
./anti-feature.nix ./anti-feature.nix
./ports.nix ./ports.nix
./services/telegram-bot/danbooru_img_bot.nix
./services/redlib.nix ./services/redlib.nix
../../../nixos/profiles/sing-box-server ../../../nixos/profiles/sing-box-server
@ -17,6 +18,10 @@
443 443
]; ];
sops.secrets = lib.mapAttrs (_n: v: v // { sopsFile = ./secrets.yaml; }) {
"tg/danbooru_img_bot" = { };
};
services.caddy.enable = true; services.caddy.enable = true;
services.caddy.settings.apps.http.servers.srv0 = { services.caddy.settings.apps.http.servers.srv0 = {
listen = [ ":443" ]; listen = [ ":443" ];

View file

@ -0,0 +1,31 @@
tg:
danbooru_img_bot: ENC[AES256_GCM,data:ZGwbGt8M4WGX2Ex0yj7NM0hv+Tre+CN48dCyuuPykafJJOWBdhGNLRzayz09ESqwmABexLeaBSYykfLMCVrETdkR7FIzd6E=,iv:p8f/aU0tPLdDLtjRIlwzq6DL+nL8n+MQutMij42Cx0M=,tag:yEefgb/c2wdRDjMyXpxZGw==,type:str]
sops:
kms: []
gcp_kms: []
azure_kv: []
hc_vault: []
age:
- recipient: age129yyxyz686qj88ce5v77ahelqqwt6zz94mzzls0ny4hq76psrd9qhc79kq
enc: |
-----BEGIN AGE ENCRYPTED FILE-----
YWdlLWVuY3J5cHRpb24ub3JnL3YxCi0+IFgyNTUxOSByRGdXamQrbEZRMEpZaVA0
aGtGNytENkozTmZmQVR2cXdQb2RRMXZLWVRVCmdPY3VBVW14M0MrTkNVb2RlNFhP
blpFclhuR1VnQ0ovZXV6QXpId3VhckEKLS0tIHpaWUxlVy8wYll6enVTNW1vWWxu
ZGVmRnhZYmNVQndzYUo5UGVWdUJqTXcKZSrvh1OzJigx0JKxKem7DnLwBbWauQNR
+KxZVoD3ZP1gJUTdpSci5Zbmp0szl0bNCJyXPdaF/IOwCckoaO4Ucw==
-----END AGE ENCRYPTED FILE-----
- recipient: age1u7srtfpgf83hesmsvtqdqftl8xrjmmp33mlg0aze6ken866ad55qxmzdqd
enc: |
-----BEGIN AGE ENCRYPTED FILE-----
YWdlLWVuY3J5cHRpb24ub3JnL3YxCi0+IFgyNTUxOSBQYThEaEgxemdaK3lRV1Vn
MHIxZUNUS3E2MVBDMVpJdGF1SXJuN1FVZTI4CmRmVEpna1g1d3ZxQTBoVklDNWQ4
VUZhanEzcDcvWGsvQTB1SThhcVd0QTQKLS0tIDJzOXYxSHJEOW14THVydmRydExW
RTIrUHBYVFU1N2NraVdSc0ZKRy9JSEEKjE9uVvnMxYzQGKKht+lMIMJBQ3P+MFBb
AQmxkPy5PH/jDpYgIsKbkBl/oyLs6qUNp8tkZa7Mn1uIl9oK8MO8lw==
-----END AGE ENCRYPTED FILE-----
lastmodified: "2024-10-07T08:05:28Z"
mac: ENC[AES256_GCM,data:C+j61xqKGkAWtHfPgdp2U4x6f66qXFnTgUjdkX4PIDQ69eoW4985V3qz/ycLhVTX2vSJVHAD0CZP5S0r3mJ2+ZsW2o9p/2ed8eIb/Nk0yT0vYdtfcx31B7ueZUWz/ezha+ysVIUO4ohdtu11hVpzFbYWG7KOTuO0g2hrdaioWiI=,iv:1UinD7iQs8FearWLWIqNDgTXDW4IBJJnBE3IX+Xze+c=,tag:TwKC9Ku19NtBk3ZSGfl3Vw==,type:str]
pgp: []
unencrypted_suffix: _unencrypted
version: 3.9.1

View file

@ -0,0 +1,42 @@
{ pkgs, config, ... }:
{
systemd.services."tg-danbooru_img_bot" = {
wantedBy = [ "multi-user.target" ];
serviceConfig = {
# TODO: un-vendor this file
ExecStart = pkgs.writers.writePython3 "pytest" {
doCheck = false;
libraries = with pkgs.python3Packages; [
python-telegram-bot
aiohttp
];
} (builtins.readFile ./danbooru_img_bot.py);
EnvironmentFile = config.sops.secrets."tg/danbooru_img_bot".path;
CapabilityBoundingSet = "";
DynamicUser = true;
LockPersonality = true;
MemoryDenyWriteExecute = true;
NoNewPrivileges = true;
PrivateDevices = true;
PrivateUsers = true;
ProcSubset = "pid";
ProtectClock = true;
ProtectControlGroups = true;
ProtectHome = true;
ProtectHostname = true;
ProtectKernelLogs = true;
ProtectKernelModules = true;
ProtectKernelTunables = true;
ProtectProc = "invisible";
ProtectSystem = "strict";
RestrictNamespaces = true;
RestrictRealtime = true;
RestrictSUIDSGID = true;
SystemCallArchitectures = "native";
SystemCallErrorNumber = "EPERM";
SystemCallFilter = "@system-service";
UMask = "0077";
};
};
}

View file

@ -0,0 +1,185 @@
import logging
import aiohttp
import random
import time
import os
from urllib.parse import urlparse
from telegram import Update, constants, InlineKeyboardMarkup, InlineKeyboardButton
from telegram.ext import ApplicationBuilder, ContextTypes, CommandHandler
logging.basicConfig(format="[%(levelname)s] %(message)s", level=logging.INFO)
class DanbooruFetcher:
def __init__(self) -> None:
self.session = None
self.proxy = os.getenv("http_proxy")
self.cache = {}
self.cache_expiry = 3600 # 1 hour
def _is_cache_valid(self, key: str) -> bool:
return key in self.cache and (
time.time() - self.cache[key]["timestamp"] < self.cache_expiry
)
async def create_session(self) -> None:
if not self.session:
self.session = aiohttp.ClientSession()
async def fetch_danbooru(self, type: str) -> dict:
await self.create_session()
cache_key = f"danbooru_{type}"
# Check if the data is already cached and valid
if self._is_cache_valid(cache_key):
logging.info(f"Returning cached data for type: {type}")
return random.choice(self.cache[cache_key]["data"])
# If not cached or expired, make a request
logging.info(f"Fetching new data for type: {type}")
async with self.session.get(
url="https://danbooru.donmai.us/posts.json",
params={
"limit": 100,
"tags": " ".join(
[
type,
"is:nsfw",
"order:rank",
]
),
},
proxy=self.proxy,
) as response:
response.raise_for_status()
data = await response.json()
# Cache the data with a timestamp
self.cache[cache_key] = {
"data": data,
"timestamp": time.time(),
}
return random.choice(data)
fetcher = DanbooruFetcher()
def format_source(record: dict) -> str:
if record["pixiv_id"] is not None:
return f"https://www.pixiv.net/en/artworks/{record['pixiv_id']}"
else:
return record["source"]
def format_tags(tags: str) -> str:
return tags.replace(" ", ", ").replace("_", " ")
def generate_caption(record: dict) -> str:
return "\n".join(
f"<pre><code class=\"language-{tag}\">{format_tags(record.get(f'tag_string_{tag}', '(no tag)'))}</code></pre>"
for tag in ["character", "copyright", "artist"]
)
def get_source_name(record: dict) -> str:
source = record["source"]
domain = urlparse(source).netloc
if record["pixiv_id"] is not None:
return "Pixiv"
if any(
domain.startswith(prefix)
for prefix in [
"twitter.com",
"x.com",
"fxtwitter.com",
"fixupx.com",
"vxtwitter.com",
]
):
return "X (Twitter)"
return domain
def get_reply_markup(r: dict) -> InlineKeyboardMarkup:
return InlineKeyboardMarkup(
[
[
InlineKeyboardButton(
text="Danbooru",
url=f"https://danbooru.donmai.us/posts/{r['id']}",
),
InlineKeyboardButton(
text=get_source_name(r),
url=format_source(r),
),
],
]
)
async def help(update: Update, _: ContextTypes.DEFAULT_TYPE) -> None:
if update.message:
await update.message.reply_text(
do_quote=True,
text="/help - get help\n/image - fetch image\n/video - fetch video",
)
async def image(update: Update, _: ContextTypes.DEFAULT_TYPE) -> None:
r = await fetcher.fetch_danbooru(type="-video")
if update.message:
try:
await update.message.reply_photo(
do_quote=True,
photo=r["large_file_url"],
parse_mode=constants.ParseMode.HTML,
caption=generate_caption(r),
reply_markup=get_reply_markup(r),
)
except Exception as e:
logging.error(f"Error in /image: {e}")
logging.error(f"{r['id']}, {r['file_url']}")
await update.message.reply_text(
do_quote=True,
text=str(e),
)
async def video(update: Update, _: ContextTypes.DEFAULT_TYPE) -> None:
r = await fetcher.fetch_danbooru(type="video")
if update.message:
try:
await update.message.reply_video(
do_quote=True,
video=r["large_file_url"],
parse_mode=constants.ParseMode.HTML,
caption=generate_caption(r),
reply_markup=get_reply_markup(r),
)
except Exception as e:
logging.error(f"Error in /video: {e}")
logging.error(f"{r['id']}, {r['file_url']}")
await update.message.reply_text(
do_quote=True,
text=str(e),
)
if __name__ == "__main__":
application = (
ApplicationBuilder()
.token(os.environ["DANBOORU_TELEGRAM_TOKEN"])
.build()
)
application.add_handler(CommandHandler("help", help))
application.add_handler(CommandHandler("image", image))
application.add_handler(CommandHandler("video", video))
application.run_polling()