AO3-DL/python/admin.py
2026-01-07 21:28:10 -08:00

678 lines
24 KiB
Python
Executable file

#!/usr/bin/env python3
# This file is part of AOMO, "Archive Of My Own", a collection of Python and PHP
# scripts designed act as an improved local backup system
# for works published on https://archiveofourown.org
#
# Copyright (c) 2025 Cyberpro123, except where otherwise noted.
#
# This project is available under the GNU General Public License (GPL) v2.0,
# available at https://www.gnu.org/licenses/old-licenses/gpl-2.0.en.html
# or in the 'LICENSE.md' file that should be distributed alongside this one.
#
# Report issues to https://codeberg.org/Cyberpro123/AOMO
# Contact author at cyberpro123@posteo.com
from collections.abc import Callable
import curses
import datetime
import os.path
import textwrap
import sys
import bcrypt
import common
import sql
class MenuItem:
def __init__(
self,
prompt: str,
actionType: str,
action: Callable | str = "",
key: str | int = "",
value: str | bool = "",
):
self.prompt = str(prompt)
self.actionType = str(actionType)
self.action = action
self.key = str(key)[:1]
self.value = value
self.defaultValue = value
def getHeight(self, width) -> int:
table = textwrap.wrap(self.value, width=width, initial_indent=4, subsequent_indent=4)
return max(1, len(table))
def render(self, width) -> list:
if self.key:
outStr = f"[{self.key}] "
else:
outStr = "[ ] "
outStr += self.prompt
if self.actionType in ("checkbox", "getTxt") and (not outStr.endswith(" ")):
outStr += " "
if self.actionType == "checkbox":
if self.value:
outStr += "[X]"
else:
outStr += "[ ]"
elif self.actionType == "getTxt":
if "password" in self.prompt.lower():
for i in self.value:
outStr += "*"
else:
outStr += self.value
return textwrap.wrap(outStr, width=width, subsequent_indent=" ")
class Menu:
def __init__(self, *args, defaultHeader: str = ""):
self.children = []
self.header = defaultHeader
for i in args:
if isinstance(i, MenuItem):
self.children.append(i)
else:
raise TypeError(f"Class 'Menu' only accepts 'MenuItem' args, not [{type(i)}]")
def getHeight(self, width) -> int:
sum = 0
for i in self.children:
sum += i.getHeight(width)
return sum
def render(self, width) -> list:
out = []
for i in self.children:
out.append(i.render(width))
return out
def getKeys(self) -> list:
out = []
for i in self.children:
if i.key:
out.append(i.key)
return out
def getMenuItemFromKey(self, key) -> MenuItem | None:
for i in self.children:
if i.key and i.key == key:
return i
def resetValues(self) -> None:
for i in self.children:
i.value = i.defaultValue
menus = {
"menu_main": Menu(
MenuItem(prompt="Manage Users", actionType="chMn", action="menu_users", key=1),
MenuItem(
prompt="Manage Site Skins",
actionType="chMn",
action="menu_site_skins",
key=2,
),
MenuItem(prompt="Manage Works", actionType="chMn", action="menu_works", key=3),
MenuItem(prompt="Close Menu", actionType="close", key=9),
),
"menu_users": Menu(
MenuItem(prompt="Add User", actionType="chMn", action="userAdd", key=1),
MenuItem(prompt="Change Password", actionType="chMn", action="chPass", key=2),
MenuItem(
prompt="Change Stylesheet",
actionType="chMn",
action="userStylesheet",
key=3,
),
MenuItem(
prompt="Change Work Rating Blacklist",
actionType="chMn",
action="ratingBlacklist",
key=4,
),
MenuItem(prompt="Delete User", actionType="chMn", action="userDel", key=5),
MenuItem(prompt="Back", actionType="chMn", action="menu_main", key=9),
),
"menu_site_skins": Menu(
MenuItem(prompt="Add Site Skin", actionType="chMn", action="skinAdd", key=1),
MenuItem(prompt="Edit Site Skin", actionType="chMn", action="skinEdit", key=2),
MenuItem(prompt="Delete Site Skin", actionType="chMn", action="skinDel", key=3),
MenuItem(prompt="Back", actionType="chMn", action="menu_main", key=9),
),
"menu_works": Menu(
MenuItem(prompt="Blacklist Work", actionType="chMn", action="blacklistAdd", key=1),
MenuItem(prompt="Un-Blacklist Work", actionType="chMn", action="blacklistDel", key=2),
MenuItem(prompt="Delete Work", actionType="chMn", action="workDel", key=3),
MenuItem(prompt="Back", actionType="chMn", action="menu_main", key=9),
),
}
loaded = {
"highestExistingUserID": None,
"highestExistingSkinID": None,
}
def doMenu(screen):
exitMenu = False
char = 0
optionSelected = 0
optionsStart = 0
menu = menus["menu_main"]
header = menu.header
while not exitMenu:
screen.erase()
skipInput = False
height, width = screen.getmaxyx()
headerList = textwrap.wrap(header, width=width)
optionsStart = len(headerList)
common.logger.debug(f"Char: [{char}]")
for num, line in enumerate(headerList):
screen.addstr(num, 0, line)
optionsStartLines = [optionsStart]
for menuItem in menu.render(width):
for num, line in enumerate(menuItem):
screen.addstr((num + optionsStartLines[-1]), 0, line)
optionsStartLines.append(len(menuItem) + optionsStartLines[-1])
optionsStartLines.pop(-1)
if (chr(char) in menu.getKeys()) or (char in (curses.KEY_RIGHT, curses.KEY_ENTER, 10, 13)):
# 10 = newline 13 = carriage return
if chr(char) in menu.getKeys():
item = menu.getMenuItemFromKey(chr(char))
else:
item = menu.children[optionSelected]
match item.actionType:
case "close":
sys.exit(0)
case "chMn":
menu.resetValues()
menu = menus[item.action]
header = menu.header
case "checkbox":
item.value = not item.value
case "function":
header = item.action(menu)
menu.resetValues()
if item.actionType not in ("checkbox", "none"):
optionSelected = 0
skipInput = True
del item
elif char in (curses.KEY_DOWN, 9): # 9 is tab
optionSelected += 1
elif char in (curses.KEY_UP, curses.KEY_BTAB):
optionSelected -= 1
elif chr(char).isprintable() and char != curses.KEY_LEFT:
item = menu.children[optionSelected]
if char == curses.KEY_BACKSPACE:
item.value = item.value[:-1]
else:
item.value = item.value + chr(char)
del item
skipInput = True
if char in (curses.KEY_HOME, curses.KEY_PPAGE) or optionSelected == len(menu.children):
optionSelected = 0
elif char in (curses.KEY_END, curses.KEY_NPAGE) or optionSelected == -1:
optionSelected = len(menu.children) - 1
screen.move(optionsStartLines[optionSelected], 1)
if skipInput:
char = 0
else:
char = screen.getch()
def getHighestExistingUserID() -> int:
if not loaded["highestExistingUserID"]:
loaded["highestExistingUserID"] = sql.execute(
commands=[("SELECT MAX(id) FROM users", ())],
config=common.config,
logger=common.logger,
)[0][0][0]
if not loaded["highestExistingUserID"]:
loaded["highestExistingUserID"] = 0
return loaded["highestExistingUserID"]
def getHighestExistingSkinID() -> int:
if not loaded["highestExistingSkinID"]:
loaded["highestExistingSkinID"] = sql.execute(
commands=[("SELECT MAX(id) FROM site_skins", ())],
config=common.config,
logger=common.logger,
)[0][0][0]
if not loaded["highestExistingSkinID"]:
loaded["highestExistingSkinID"] = 0
return loaded["highestExistingSkinID"]
def getIdsFromIdOrFile(input: str) -> set:
input = str(input)
output = set(())
if input.isnumeric():
output.add(int(input))
elif input:
with open(input) as file:
output = common.parseInfile(file)
return output
def getSkinIdAndNameFromEither(input: str) -> str | tuple[str, str]:
if str(input).isnumeric():
id = input
sheetName = sql.execute(
commands=[("SELECT name FROM site_skins WHERE id = ?", (id,))],
config=common.config,
logger=common.logger,
)[0]
if sheetName:
sheetName = sheetName[0][0]
else:
return f"ERROR: Sheet with ID [{id}] doesn't exist"
else:
sheetName = input
id = sql.execute(
commands=[("SELECT id FROM site_skins WHERE name = ?", (sheetName,))],
config=common.config,
logger=common.logger,
)[0]
if id:
id = id[0][0]
else:
return f"ERROR: Sheet with name [{sheetName}] doesn't exist"
return (id, sheetName)
################################################################
# Single-Menu Functions
################################################################
################################
# Manage Users
################################
def addUser(menu: MenuItem) -> str:
id = 1 + getHighestExistingUserID()
username = menu.children[0].value
pass1 = menu.children[1].value
pass2 = menu.children[2].value
preHashed = menu.children[3].value
if pass1 != pass2:
return "ERROR: Passwords don't match. User not added."
if len(username) > 100:
return "ERROR: Username too long (max 100 characters)"
if len(pass1) > 72:
return "ERROR: Password too long (max 72 characters)"
if preHashed:
hashed = pass1
else:
hashed = bcrypt.hashpw(pass1.encode(), bcrypt.gensalt()).decode()
if not bcrypt.checkpw(pass1.encode(), hashed.encode()):
return "ERROR: Password hashing failed"
try:
sql.execute(
commands=[
(
"INSERT INTO users (id, username, password_hash) VALUES (?, ?, ?)",
(id, username, hashed),
)
],
config=common.config,
logger=common.logger,
)
except sql.UniqueConstraintFailed:
return "ERROR: User with that name already exists."
else:
loaded["highestExistingUserID"] += 1
return f"Successfully added user [{username}] as ID [{id}]"
def changePassword(menu: MenuItem) -> str:
username = menu.children[0].value
pass1 = menu.children[1].value
pass2 = menu.children[2].value
preHashed = menu.children[3].value
if pass1 != pass2:
return "ERROR: Passwords don't match. User not added."
if len(username) > 100:
return "ERROR: Username too long (max 100 characters)"
if len(pass1) > 72:
return "ERROR: Password too long (max 72 characters)"
if preHashed:
hashed = pass1
else:
hashed = bcrypt.hashpw(pass1.encode(), bcrypt.gensalt()).decode()
if not bcrypt.checkpw(pass1.encode(), hashed.encode()):
return "ERROR: Password hashing failed"
sql.execute(
commands=[
(
"UPDATE users SET password_hash = ? WHERE username = ?",
(hashed, username),
)
],
config=common.config,
logger=common.logger,
)
if menu.children[4].value:
id = sql.execute(
commands=[("SELECT id FROM users WHERE username = ?", (username,))],
config=common.config,
logger=common.logger,
)[0]
if not id:
return f"ERROR: User [{username}] doesn't exist?"
sql.execute(
commands=[("DELETE FROM login_tokens WHERE id = ?", (id[0][0],))],
config=common.config,
logger=common.logger,
)
return f"Successfully changed password for user [{username}]"
def changeUserStylesheet(menu: MenuItem) -> str:
username = menu.children[0].value
tmp = getSkinIdAndNameFromEither(menu.children[1].value)
if isinstance(tmp, str):
return tmp
else:
id, sheetName = tmp
del tmp
sql.execute(
commands=[("UPDATE users SET site_skin = ? WHERE username = ?", (id, username))],
config=common.config,
logger=common.logger,
)
return f"Successfully set user [{username}]'s Site Skin to id [{id}] name [{sheetName}]"
def changeUserBlacklist(menu: MenuItem) -> str:
username = menu.children[0].value
value = 0
for i in range(5):
if not menu.children[1 + i].value:
value += pow(2, i)
sql.execute(
commands=[
(
"UPDATE users SET can_see_explicit_works = ? WHERE username = ?",
(value, username),
)
],
config=common.config,
logger=common.logger,
)
return f"Set user [{username}]'s Blacklist value to [{value}]"
def deleteUser(menu: MenuItem) -> str:
tmp = menu.children[0].value
if str(tmp).isnumeric():
id = tmp
del tmp
username = sql.execute(
commands=[("SELECT username FROM users WHERE id = ?", (id,))],
config=common.config,
logger=common.logger,
)[0]
if username:
username = username[0][0]
else:
return f"ERROR: User with ID [{id}] doesn't exist"
else:
username = tmp
del tmp
id = sql.execute(
commands=[("SELECT id FROM users WHERE username = ?", (username,))],
config=common.config,
logger=common.logger,
)[0]
if id:
id = id[0][0]
else:
return f"ERROR: User with name [{username}] doesn't exist"
sql.execute(
commands=[
("DELETE FROM users WHERE id = ?", (id,)),
("DELETE FROM login_tokens WHERE id = ?", (id,)),
("DELETE FROM progress_tracker WHERE user_id = ?", (id,)),
],
config=common.config,
logger=common.logger,
)
return f"Deleted id [{id}] name [{username}]"
################################
# Manage Site Skins
################################
def addSkin(menu: MenuItem) -> str:
id = 1 + getHighestExistingSkinID()
name = menu.children[0].value
filePath = os.path.abspath(os.path.expanduser(menu.children[1].value))
url = menu.children[2].value
if len(name) > 100:
return "ERROR: Skin Name too long (max 100 characters)"
if len(url) > 300:
return "ERROR: Preview Image URL too long (max 300 characters)"
if os.path.exists(filePath) and os.path.getsize(filePath):
with open(filePath, "r") as file:
fileStr = file.read()
try:
sql.execute(
commands=[
(
"INSERT INTO site_skins VALUES (?, ?, ?, ?)",
(id, fileStr, name, url),
),
],
config=common.config,
logger=common.logger,
)
except sql.UniqueConstraintFailed:
return "ERROR: Skin with that name already exists"
else:
loaded["highestExistingSkinID"] += 1
return f"Successfully added Site Skin [{name}] as ID [{id}]"
else:
return f"ERROR: File [{filePath}] doesn't exist or is 0 bytes long"
def editSkin(menu: MenuItem) -> str:
tmp = getSkinIdAndNameFromEither(menu.children[0].value)
if isinstance(tmp, str):
return tmp
else:
id, sheetName = tmp
del tmp
filePath = os.path.abspath(os.path.expanduser(menu.children[1].value))
url = menu.children[2].value
fileExists = os.path.exists(filePath) or os.path.getsize(filePath)
if (not url) and (not fileExists):
return "ERROR: One of 'File' or 'Preview Image URL' must be set"
sqlCommands = []
if fileExists:
with open(filePath, "r") as file:
fileStr = file.read()
sqlCommands.append(("UPDATE site_skins SET body = ? WHERE id = ?", (fileStr, id)))
if url:
sqlCommands.append(("UPDATE site_skins SET url = ? WHERE id = ?", (url, id)))
sql.execute(
commands=sqlCommands,
config=common.config,
logger=common.logger,
)
return f"Successfully edited Site Skin id [{id}] name [{sheetName}]"
def deleteSkin(menu: MenuItem) -> str:
tmp = getSkinIdAndNameFromEither(menu.children[0].value)
if isinstance(tmp, str):
return tmp
else:
id, sheetName = tmp
del tmp
sql.execute(
commands=[
("DELETE FROM site_skins WHERE id = ?", (id,)),
("UPDATE users SET site_skin = 0 WHERE site_skin = ?", (id,)),
],
config=common.config,
logger=common.logger,
)
return f"Deleted id [{id}] name [{sheetName}]"
################################
# Manage Works
################################
def addToBlacklist(menu: MenuItem) -> str:
ids = getIdsFromIdOrFile(menu.children[0].value)
reason = str(menu.children[1].value)[:1000]
now = int(datetime.datetime.now().timestamp())
if not isinstance(now, int):
raise Exception(f"{type(now)}")
sqlStr = "INSERT INTO blacklist VALUES (?, ?, ?) ON CONFLICT DO NOTHING"
commands = []
for i in ids:
commands.append((sqlStr, (i, now, reason)))
sql.execute(commands=commands, config=common.config, logger=common.logger)
if len(ids) == 1:
for i in ids:
return f"Added workID [{i}] to blacklist"
else:
return f"Added [{len(ids)}] works to blacklist"
def removeFromBlacklist(menu: MenuItem) -> str:
commands = []
for i in getIdsFromIdOrFile(menu.children[0].value):
commands.append(("DELETE FROM blacklist WHERE id = ?", (i,)))
reason = menu.children[1].value
if reason != "":
commands.append(("DELETE FROM blacklist WHERE reason = ?", (reason,)))
timestamp = menu.children[2].value
if timestamp != "":
commands.append(("DELETE FROM blacklist WHERE timestamp_added = ?", (timestamp,)))
sql.execute(commands=commands, config=common.config, logger=common.logger)
return "Removed specified works from blacklist"
def deleteWork(menu: MenuItem) -> str:
ids = getIdsFromIdOrFile(menu.children[0].value)
commands = []
commandStrs = [
"DELETE FROM works WHERE id = ?",
"DELETE FROM tags WHERE id = ?",
"DELETE FROM text WHERE id = ?",
"DELETE FROM progress_tracker WHERE id = ?",
"DELETE FROM inspirations WHERE id = ? AND NOT EXISTS(SELECT id FROM works WHERE works.id = inspirations.parent_id)",
"DELETE FROM inspirations WHERE parent_id = ? AND NOT EXISTS(SELECT id FROM works WHERE works.id = inspirations.id)",
]
for i in ids:
for cmd in commandStrs:
commands.append((cmd, (i,)))
sql.execute(commands=commands, config=common.config, logger=common.logger)
return "Deleted specified works"
################################################################
# Menu() Definitions
################################################################
################################
# Manage Users
################################
menus["userAdd"] = Menu(
MenuItem(prompt="Username: ", actionType="getTxt"),
MenuItem(prompt="Enter Password: ", actionType="getTxt"),
MenuItem(prompt="Confirm Password: ", actionType="getTxt"),
MenuItem(prompt="Password Pre-Hashed?", actionType="checkbox", value=False),
MenuItem(prompt="[Submit]", actionType="function", action=addUser),
MenuItem(prompt="Back", actionType="chMn", action="menu_users"),
)
menus["chPass"] = Menu(
MenuItem(prompt="Username: ", actionType="getTxt"),
MenuItem(prompt="Enter Password: ", actionType="getTxt"),
MenuItem(prompt="Confirm Password: ", actionType="getTxt"),
MenuItem(prompt="Password Pre-Hashed?", actionType="checkbox", value=False),
MenuItem(prompt="Delete login tokens?", actionType="checkbox", value=True),
MenuItem(prompt="[Submit]", actionType="function", action=changePassword),
MenuItem(prompt="Back", actionType="chMn", action="menu_users"),
)
menus["userStylesheet"] = Menu(
MenuItem(prompt="Username: ", actionType="getTxt"),
MenuItem(prompt="Stylesheet: ", actionType="getTxt"),
MenuItem(prompt="[Submit]", actionType="function", action=changeUserStylesheet),
MenuItem(prompt="Back", actionType="chMn", action="menu_users"),
)
menus["ratingBlacklist"] = Menu(
MenuItem(prompt="Username: ", actionType="getTxt"),
MenuItem(prompt="Can See Gen Works?", actionType="checkbox", value=True),
MenuItem(prompt="Can See Teen+ Works?", actionType="checkbox", value=True),
MenuItem(prompt="Can See Mature Works?", actionType="checkbox", value=True),
MenuItem(prompt="Can See Explicit Works?", actionType="checkbox", value=True),
MenuItem(prompt="Can See Unrated Works?", actionType="checkbox", value=True),
MenuItem(prompt="[Submit]", actionType="function", action=changeUserBlacklist),
MenuItem(prompt="Back", actionType="chMn", action="menu_users"),
defaultHeader="WARNING: These settings aren't actually checked by the current WebUI version at all.",
)
menus["userDel"] = Menu(
MenuItem(prompt="Username: ", actionType="getTxt"),
MenuItem(prompt="[Submit]", actionType="function", action=deleteUser),
MenuItem(prompt="Back", actionType="chMn", action="menu_users"),
defaultHeader=(
"WARNING: This will irreversibly delete all information associated with the user, including the 'progress tracker'"
+ " information showing what chapters of what works they have read. For a less permenant option, consider changing their password."
),
)
################################
# Manage Site Skins
################################
menus["skinAdd"] = Menu(
MenuItem(prompt="Name: ", actionType="getTxt"),
MenuItem(prompt="File: ", actionType="getTxt"),
MenuItem(prompt="Preview Image URL (Optional): ", actionType="getTxt"),
MenuItem(prompt="[Submit]", actionType="function", action=addSkin),
MenuItem(prompt="Back", actionType="chMn", action="menu_site_skins"),
)
menus["skinEdit"] = Menu(
MenuItem(prompt="Name or ID: ", actionType="getTxt"),
MenuItem(prompt="New File: ", actionType="getTxt"),
MenuItem(prompt="Preview Image URL: ", actionType="getTxt"),
MenuItem(prompt="[Submit]", actionType="function", action=editSkin),
MenuItem(prompt="Back", actionType="chMn", action="menu_site_skins"),
)
menus["skinDel"] = Menu(
MenuItem(prompt="Name or ID: ", actionType="getTxt"),
MenuItem(prompt="[Submit]", actionType="function", action=deleteSkin),
MenuItem(prompt="Back", actionType="chMn", action="menu_site_skins"),
defaultHeader="WARNING: This can't be undone. Any users with the given skin selected will have their selected skin set to 0",
)
################################
# Manage Works
################################
menus["blacklistAdd"] = Menu(
MenuItem(prompt="ID or filename: ", actionType="getTxt"),
MenuItem(prompt="Reason (Opt.): ", actionType="getTxt"),
MenuItem(prompt="[Submit]", actionType="function", action=addToBlacklist),
MenuItem(prompt="Back", actionType="chMn", action="menu_works"),
)
menus["blacklistDel"] = Menu(
MenuItem(prompt="ID or filename: ", actionType="getTxt"),
MenuItem(prompt="Reason: ", actionType="getTxt"),
MenuItem(prompt="Timestamp: ", actionType="getTxt"),
MenuItem(prompt="[Submit]", actionType="function", action=removeFromBlacklist),
MenuItem(prompt="Back", actionType="chMn", action="menu_works"),
)
menus["workDel"] = Menu(
MenuItem(prompt="ID or filename: ", actionType="getTxt"),
MenuItem(prompt="[Submit]", actionType="function", action=deleteWork),
MenuItem(prompt="Back", actionType="chMn", action="menu_works"),
)
if __name__ == "__main__":
common.init(defaultOutfile="/dev/null")
curses.wrapper(doMenu)