2024-07-06 21:10:18 +00:00
2024-07-05 19:37:11 +00:00
import argparse
from mastodon import Mastodon
import os
import sqlite3
2024-07-06 21:10:18 +00:00
import logging
import time
2024-07-05 19:37:11 +00:00
import tomllib
ACCOUNT_FETCH_LIMIT = 100000
def check_db_exists ( cursor : sqlite3 . Cursor ) - > bool :
res = cursor . execute ( " SELECT COUNT(*) FROM sqlite_master WHERE type= ' table ' AND name= ' welcome_log ' " )
if res . fetchone ( ) [ 0 ] == 0 :
cursor . execute ( " CREATE TABLE welcome_log (id INTEGER PRIMARY KEY, username TEXT, userdb_id INTEGER, welcomed INTEGER DEFAULT 0) " )
return False
return True
def user_exists ( cursor : sqlite3 . Cursor , userid : int ) - > bool :
res = cursor . execute ( " SELECT COUNT(*) FROM welcome_log WHERE userdb_id = ? " , ( userid , ) )
return res . fetchone ( ) [ 0 ] > 0
def user_welcomed ( cursor : sqlite3 . Cursor , userid : int ) - > bool :
res = cursor . execute ( " SELECT COUNT(*) FROM welcome_log WHERE userdb_id = ? AND welcomed = 1 " , ( userid , ) )
return res . fetchone ( ) [ 0 ] > 0
def set_user_welcomed ( cursor : sqlite3 . Cursor , userid : int ) :
cursor . execute ( " UPDATE welcome_log SET welcomed = 1 WHERE userdb_id = ? " , ( userid , ) )
def create_user ( cursor : sqlite3 . Cursor , userid : int , username : str ) :
cursor . execute ( " INSERT INTO welcome_log (userdb_id, username) VALUES(?, ?) " , ( userid , username ) )
2024-07-06 21:10:18 +00:00
def check_and_welcome_new_users ( ) :
logging . debug ( " Checking if database exists " )
fresh_database = not check_db_exists ( cursor )
if fresh_database :
logging . info ( " Database was freshly created - we ' ll set all pre-existing users as of now to ' welcomed ' to avoid spamming everyone on the server. " )
logging . info ( " Starting the welcome process " )
all_accounts = mastodon . admin_accounts ( remote = False , status = ' active ' , limit = ACCOUNT_FETCH_LIMIT )
for account in all_accounts :
logging . debug ( f " Checking if user { account [ ' id ' ] } exists " )
if not user_exists ( cursor , account [ ' id ' ] ) :
logging . debug ( f " Creating user { account [ ' username ' ] } with ID { account [ ' id ' ] } " )
create_user ( cursor , account [ ' id ' ] , account [ ' username ' ] )
connection . commit ( )
logging . debug ( f " Checking if user { account [ ' id ' ] } was welcomed " )
if fresh_database :
set_user_welcomed ( cursor , account [ ' id ' ] )
connection . commit ( )
elif not user_welcomed ( cursor , account [ ' id ' ] ) :
result_id = None
for message in config [ ' messages ' ] :
content_warning = message [ ' content_warning ' ] if ' content_warning ' in message else None
result = mastodon . status_post ( status = f " @ { account [ ' username ' ] } , { message [ ' content ' ] } " , in_reply_to_id = result_id , visibility = ' public ' , spoiler_text = content_warning )
result_id = result . id
set_user_welcomed ( cursor , account [ ' id ' ] )
connection . commit ( )
2024-07-05 19:37:11 +00:00
if __name__ == ' __main__ ' :
2024-07-06 21:10:18 +00:00
try :
import tomllib
except ModuleNotFoundError :
import tomli as tomllib
logging . basicConfig ( filename = ' /home/archos/welcome_bot/welcome_bot.log ' , level = logging . DEBUG , format = ' %(asctime)s %(levelname)s : %(message)s ' )
2024-07-05 19:37:11 +00:00
arg_parser = argparse . ArgumentParser (
2024-07-06 21:10:18 +00:00
description = " Welcomes users to the Mastodon instance " ,
epilog = " Note: the user this bot logs in as must have admin:read:accounts access "
2024-07-05 19:37:11 +00:00
)
arg_parser . add_argument ( ' --email ' , required = False , help = " Only required for first execution " )
arg_parser . add_argument ( ' --password ' , required = False , help = " Only required for first execution " )
arg_parser . add_argument ( ' --config ' , default = " config.toml " )
args = arg_parser . parse_args ( )
with open ( args . config , " rb " ) as toml_file :
config = tomllib . load ( toml_file )
mastodon = None
if not os . path . exists ( config [ ' mastodon ' ] [ ' credential_storage ' ] ) :
if not ( args . email and args . password ) :
2024-07-06 21:10:18 +00:00
logging . error ( " Initial login has not yet occurred - this is required to generate the credential file. Please supply login credentials (--email, --password) " )
2024-07-05 19:37:11 +00:00
exit ( - 1 )
else :
2024-07-06 21:10:18 +00:00
logging . info ( " Registering app " )
2024-07-05 19:37:11 +00:00
Mastodon . create_app (
2024-07-06 21:10:18 +00:00
client_name = " Welcome Bot " ,
scopes = [ " write:statuses " , " admin:read:accounts " ] ,
to_file = config [ ' mastodon ' ] [ ' secret_storage ' ] ,
api_base_url = config [ ' mastodon ' ] [ ' base_url ' ]
2024-07-05 19:37:11 +00:00
)
2024-07-06 21:10:18 +00:00
logging . info ( " Initializing client " )
2024-07-05 19:37:11 +00:00
mastodon = Mastodon (
2024-07-06 21:10:18 +00:00
client_id = config [ ' mastodon ' ] [ ' secret_storage ' ] ,
api_base_url = config [ ' mastodon ' ] [ ' base_url ' ]
2024-07-05 19:37:11 +00:00
)
2024-07-06 21:10:18 +00:00
logging . info ( " Performing login " )
2024-07-05 19:37:11 +00:00
mastodon . log_in (
2024-07-06 21:10:18 +00:00
username = args . email ,
password = args . password ,
to_file = config [ ' mastodon ' ] [ ' credential_storage ' ] ,
scopes = [ " write:statuses " , " admin:read:accounts " ]
2024-07-05 19:37:11 +00:00
)
else :
mastodon = Mastodon (
2024-07-06 21:10:18 +00:00
access_token = config [ ' mastodon ' ] [ ' credential_storage ' ] ,
api_base_url = config [ ' mastodon ' ] [ ' base_url ' ]
2024-07-05 19:37:11 +00:00
)
assert mastodon is not None , " Mastodon client not initialized "
connection = sqlite3 . connect ( config [ ' database ' ] [ ' sqlite_path ' ] )
cursor = connection . cursor ( )
2024-07-06 21:10:18 +00:00
while True :
check_and_welcome_new_users ( )
logging . info ( " Sleeping for 10 minutes " )
time . sleep ( 600 )