You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

128 lines
4.2 KiB

#!/usr/bin/env python3
import magic
import sqlite3
import sys
import random
import os
import os.path as op
from contextlib import AbstractContextManager
from common import load_config
from datetime import datetime, timedelta
from typing import Any, Dict, Optional
from mastodon import Mastodon
from saucenao_api import SauceNao
from saucenao_api.containers import BasicSauce
class ArtDB(AbstractContextManager):
conn: sqlite3.Connection
def __init__(self, db_file: str) -> None:
self.conn = sqlite3.connect(db_file, detect_types=sqlite3.PARSE_DECLTYPES)
self.conn.row_factory = sqlite3.Row
self.conn.execute('''CREATE TABLE IF NOT EXISTS art(
filename TEXT PRIMARY KEY,
last_used timestamp
)''')
self.conn.commit()
def __enter__(self) -> 'ArtDB':
return self
def __exit__(self, exc_type, exc_value, traceback) -> None:
self.conn.close()
def get_art(self, art: str) -> Optional[sqlite3.Row]:
info = self.conn.execute('SELECT * FROM art WHERE filename=?', (art,)).fetchone()
return info
def update_art(self, art: str) -> None:
self.conn.execute('''INSERT INTO art(filename, last_used) VALUES (:art, :last_used)
ON CONFLICT(filename) DO UPDATE SET last_used=:last_used''',
{'art': art, 'last_used': datetime.now()})
self.conn.commit()
# --------------------------------------------------
def art_recently_used(art_info: Optional[sqlite3.Row], within: timedelta) -> bool:
if art_info is None:
return False
last_used = art_info['last_used']
return datetime.now() - last_used <= within
def get_source(art: str, saucenao_config: Dict[str, Any]) -> Optional[BasicSauce]:
'''
Return a source for the given image file.
A source is returned iff there's a source with a similarity score >= min_similarity
and it contains at least one URL.
'''
sauce = SauceNao(saucenao_config['api_key'])
with open(art, 'rb') as f:
results = sauce.from_file(f)
results = [result for result in results
if result.similarity >= saucenao_config['min_similarity']
and len(result.urls) > 0
and all(domain not in result.urls[0] for domain in saucenao_config['forbidden_domains'])]
if len(results) == 0:
return None
return results[0]
def main():
config = load_config()
DIR_SFW = config['local']['dir_sfw']
DIR_NSFW = config['local']['dir_nsfw']
SAUCENAO_CONFIG = config['saucenao']
mastodon = Mastodon(
access_token = 'token.dat',
api_base_url = config["api_base_url"]
)
sfwcount = len([name for name in os.listdir(DIR_SFW) if op.isfile(op.join(DIR_SFW, name))])
nsfwcount = len([name for name in os.listdir(DIR_NSFW) if op.isfile(op.join(DIR_NSFW, name))])
random_choice = random.randint(1, sfwcount + nsfwcount)
print(f'\ns:{sfwcount} n:{nsfwcount} r:{random_choice}')
is_safe = False if random_choice < nsfwcount else True
art = ""
dir_ = DIR_SFW if is_safe else DIR_NSFW
files = [f for f in os.listdir(dir_) if op.isfile(op.join(dir_, f))]
with ArtDB('art.db') as db:
while True:
art = op.join(dir_, random.choice(files))
art_info = db.get_art(art)
if not art_recently_used(art_info, timedelta(hours=5)):
db.update_art(art)
break
else:
print(f"Art {art} was used in the last 5 hours, trying again")
fformat = magic.from_file(art, mime=True)
with open(art, 'rb') as picture:
data = picture.read()
media = mastodon.media_post(data, fformat)
toot = config["local"]["toot_text"]
if SAUCENAO_CONFIG['enabled']:
source = get_source(art, SAUCENAO_CONFIG)
if source:
toot = add_source(source, toot)
mastodon.status_post(toot, media_ids=[media], visibility=config["visibility"], sensitive=not is_safe)
print(str(datetime.now()) + ': ' + art)
def add_source(source, toot):
toot += f'\n{source.urls[0]}'
if source.author:
toot += f'\nAuthor: {source.author}'
return toot
if __name__ == '__main__':
sys.exit(main())