sp-mastodon-module/mastodon-kanidm-sync.py

from io import DEFAULT_BUFFER_SIZE
import os
import json
import requests
import psycopg2 as ps

def read_file(path):
    with open(path, "r", encoding="utf-8") as f:
        return f.read()

def getenv(name):
    try:
        return os.environ[name]
    except KeyError:
        print(f"[ERROR] Missing environment variable {name}. You should NOT run this script by hand, please use systemd mastodon-kanidm-sync.service.")
        exit(1)


# Import configuration
KANIDM_URL = getenv("KANIDM_URL")
KANIDM_TOKEN = read_file(getenv("KANIDM_TOKEN_PATH")).strip()
OWNER_USERNAME = getenv("OWNER_USERNAME")

# Fetch kanidm users list from userdata file
# Userdata file is json list with information about what users are configured by kanidm
try:
    USERDATA = read_file(getenv("USERDATA_FILE_PATH")).strip()
    userdata = json.loads(USERDATA)
    print("[INFO] ")
except FileNotFoundError:
  userdata = []

# Load database
conn = ps.connect(
    dbname=getenv("POSTGRES_DBNAME"),
    user=getenv("POSTGRES_USER"),
    host=getenv("POSTGRES_HOST")
)

# Fetch current userdata from database
cur = conn.cursor()
cur.execute('''
    SELECT identities.uid, users.id, user_roles.name
    FROM users
    JOIN identities
      ON users.id = identities.id
    LEFT JOIN user_roles
      ON users.role_id = user_roles.id;
    '''
)

state = cur.fetchall()

users = {}
for i in state:
    users[i[0]] = {
        "id": i[1],
        "role": i[2],
        "isKanidmUser": False
    }

# Fetch Kanidm userdata
kanidm_users_raw = requests.get(
    f"{KANIDM_URL}/v1/person",
    headers={
        "Authorization": f"Bearer {KANIDM_TOKEN}",
        "Content-Type": "application/json",
    },
    timeout=5,
).json()

def give_role(uid, role, putUserdata = True):
    if (uid not in userdata) and (putUserdata):
        userdata.append(uid)
    users[uid]["isKanidmUser"] = True
    users[uid]["role"] = role
    print(f"[INFO] {uid} is marked as {role}")
    

for i in kanidm_users_raw:
    i = i["attrs"]
    for uid in i["name"]: # [user].attrs.name is a list
      if uid in users: # Don't apply anything for users who have no mastodon access (sp.mastodon.users) or didn't register
        if uid == OWNER_USERNAME:
            give_role(uid, "Owner", False)
            continue

        for group in i["memberof"]:
          if group.startswith("sp.mastodon.admins@") or group.startswith("sp.admins@"):
            give_role(uid, "Admin")
            break

          elif group.startswith("sp.mastodon.moderators@"):
            give_role(uid, "Moderator")
            break

          elif uid in userdata:
            # If user, who previously had a role, has no roles set by Kanidm, delete them from userdata list so allow setting roles directly by mastodon
            give_role(uid, None, False)
            userdata.remove(uid)

print("[DEBUG] ", users) # DEBUG

# Fetch RoleIDs
cur = conn.cursor()
cur.execute("SELECT id, name FROM user_roles;")

roles_raw = cur.fetchall()
roles = {}
for i in roles_raw:
    roles[i[1]] = i[0]

print("[debug]", roles)

# Give roles
for uid in users:
    if not users[uid]["isKanidmUser"]:
        continue
    
    if users[uid]["role"]:
        rolename = users[uid]["role"]
        roleid = roles[rolename]
    else:
        roleid = "NULL"

    sqlcommand = f"UPDATE users SET role_id = {roleid} WHERE id = {users[uid]["id"]};"
    print("[debug] SQL:", sqlcommand)
    cur.execute(sqlcommand)

conn.commit()
cur.close()
conn.close()

print("[INFO] Final userdata.json file content: ", userdata)

def write_userdata(mode):
    with open(getenv("USERDATA_FILE_PATH"), mode) as f:
        f.write(json.dumps(userdata))
        f.close()

try:
    write_userdata("w")
except FileNotFoundError:
    print("[INFO] userdata.json file doesn't exist. Creating it")
    write_userdata("x")